Over the last decade, there has been significant progress in the field of interactive virtual rehabilitation. Physical therapy (PT) stands as a highly effective approach for enhancing physical impairments. However, patient motivation and progress tracking in rehabilitation outcomes remain a challenge. This work addresses the gap through a machine learning-based approach to objectively measure outcomes of the upper limb virtual therapy system in a user study with non-clinical participants. In this study, we use virtual reality to perform several tracing tasks while collecting motion and movement data using a KinArm robot and a custom-made wearable sleeve sensor. We introduce a two-step machine learning architecture to predict the motion intention of participants. The first step predicts reaching task segments to which the participant-marked points belonged using gaze, while the second step employs a Long Short-Term Memory (LSTM) model to predict directional movements based on resistance change values from the wearable sensor and the KinArm. We specifically propose to transpose our raw resistance data to the time-domain which significantly improves the accuracy of the models by 34.6%.
To evaluate the effectiveness of our model, we compared different classification techniques with various data configurations. The results show that our proposed computational method is exceptional at predicting participant's actions with accuracy values of 96.72% for diamond reaching task, and 97.44% for circle reaching task, which demonstrates the great promise of using multimodal data, including eye-tracking and resistance change, to objectively measure the performance and intention in virtual rehabilitation settings.

*Data and Code Availability
This paper uses two types of data: gaze data collected from the HTC Vive Pro Eye tracker and resistance data collected from a wearable sleeve sensor made of carbon nanotubes. Detailed explanations of the data collection process are provided in later sections.
The data and code generated during the current study are publicly accessible via the following

*Institutional Review Board (IRB)
The research has been sanctioned by the Institutional Review Board (IRB) of the University of Delaware, ensuring compliance with all ethical norms concerning research with human participants. This encompasses the protection of participant confidentiality and the mitigation of any potential harm. Ethical approval was obtained on Feb 3, 2022 under protocol Number 1982585-1, and applies for the duration of the project.

§ INTRODUCTION

Over the past years, the use of virtual reality (VR) has increased significantly. VR involves various technologies that create an engaging, simulated digital world. Users can interact within this environment, which responds to their movements, fostering a sense of presence in the virtual environment.
Moreover, there have been substantial developments in these interactive virtual settings specifically designed for the motor skills rehabilitation. These improvements are especially beneficial for patients with cognitive issues and those undergoing orthopedic rehabilitation.

In addition, technology advancements in robotics and artificial intelligence progress the rehabilitation process for upper extremity patients [Araújo et al., 2020]. Patients impacted by from stroke need to perform reaching and stretching movements to regain mobility in the upper limb. The VR setup for the upper extremity offers benefits like an increase in patient engagement in therapy, easy setup for home, and no distraction and disturbance from the environment [Baron et al., 2021, Baron et al., 2023]. Robotics rehabilitation also shows a huge improvement in terms of the recovery process. Despite its potential, the research on combining VR and the endpoint robot is limited [Mubin et al., 2019, Tarnita et al., 2022].

In this work, we utilized a framework for upper extremity rehabilitation that integrates both VR and endpoint robotics. The study involved 16 participants, each engaging in two distinct reaching tasks within the framework: drawing shapes such as a circle and a diamond. The participants were required to hit each of the 40 points arranged in the shapes of a diamond or circle.

This setup was crucial for evaluating their ability to execute planar movements across multiple joints, offering insights into the smoothness and precision of their reaching trajectories [Kwakkel et al., 2019]. We have extracted two different types of data using this framework: gaze data and resistance values. The gaze data was collected from the VR headset which provides information about the visual focus throughout the task. The resistance data was captured from nano-composite sensors embedded in the wearable sleeve. These values were recorded during muscle flexion and extension while performing reaching tasks with the sensor effectively converting the physical energy exerted during muscle movement into measurable electrical signals [Saggio et al., 2015]. The above-described framework used for data collection exhibits significant advancements in the field of rehabilitation and improvement of motor function.

In this study, we used these two types of data, gaze and resistance values, to predict the motion intentions of participants. To assess the effectiveness of our proposed approach, we applied different machine learning models with various data setups and compared the results with our proposed architecture. This comparative analysis provides insights into the performance and advantages of our proposed methodology over alternative approaches. The study on predicting the motion intentions of participants allows therapists to identify specific motor deviations. Therapists can use this information to tailor rehabilitation exercises, potentially leading to a more efficient and effective rehabilitation process.

The contributions of this paper include:

* Introducing a machine learning framework for prediction of the segmentation of trajectory shapes and directional movement.

* Integrating and leveraging multiple data modalities, such as gaze data and resistance measurements from sleeve sensors.

* Converting resistance measurements obtained from wearable sleeve sensors into time-domain features to attain higher accuracy.

* Incorporating KinArm robot data for analyzing motion intention, thereby enhancing the efficacy of predictive models for potential upper extremity rehabilitation tasks.

This study is organized as follows: Following the introduction, the related works are discussed in Section <ref>. The methodology of the work is explained in Section <ref>. The results are presented in Section <ref>. The discussion, including implications, limitations, and future research directions, is in Section <ref>. Lastly, the conclusion is presented in Section <ref>

§ RELATED WORK

In this section, we review some previous work related to virtual reality and robotics for upper extremity rehabilitation and the use of machine learning for upper limb motion prediction.

§.§ VR Therapy and Robots

Recent technological advancements have significantly increased the demand for human-machine interaction, particularly in the realm of robotics. This surge in interest has notably impacted the field of rehabilitation, leading to a deeper understanding of motor control. This work focuses on developing a better machine learning algorithm that can precisely analyze the motion intentions. Robots can analyze the movement of limbs help in controlling the motions of the joint and provide precise motion direction for effective therapy [Durandau et al., 2019]. When compared with normal physical therapy, assistive robots result in a more positive outcome in rehabilitation because patients can train for longer periods [Daunoraviciene et al., 2018]. Robots can also assist in improving motor control and restoring neurological functions [Dixit and Tedla, 2019]. [Kim et al., 2023] developed an electromagnetic robot for the rehabilitation of the upper extremity in sub-acute stroke patients. Its real-time position tracking feature allows for active flexion and extension, irrespective of the hand’s position. The researchers found that this robotic rehabilitation system aided stroke patients in improving their hand-motor function. This system also helps in tracking the motion in real-time and get feedback from it.

Virtual reality-based rehabilitation in stroke patients demonstrates improvements in patient motivation, functional recovery, motor function, and dual-task performance [Aderinto et al., 2023]. [Anwar et al., 2021] conducted a study on 68 stroke patients to compare VR training and physical therapy, revealing that VR was more effective in improving balance and lower limb function. [Vibhuti et al., 2023] presented a systematic review of VR therapy for various neuromotor impairments such as stroke, cerebral palsy, spinal cord injury, and parkinson’s disease reviewing forty-five studies. They reported that home-based VR therapy systems significantly improved the functional mobility of the patients. VR and gaming for upper extremity rehabilitation have also been shown to be effective in improving the motor recovery of the patients [Karamians et al., 2020].

[Walker et al., 2023] specified the potential of robots and VR in improving the treatment procedure by using assistive robots. [Wonsick and Padir, 2020] proposed the increase in growth of technology, the integration of both robots and the VR interface has more potential to improve rehabilitation techniques for people suffering from motor impairments.

Wearable sensors provide continuous tracking and valuable clinical insights on physical activities such as the practice of diminished skills, thus playing an important role in rehabilitation systems [Dobkin and Dorsch, 2011]. [Caviedes et al., 2020] proposed a wearable sensor array design for monitoring spine posture during physical therapy exercises. The sensor array consists of stretch sensors integrated into a custom-made sports garment, forming a triangular pattern resembling spinal exercise. Their pilot study showcased the potential of the wearable sensor array in monitoring both the adherence and accuracy of therapeutic spinal exercises during unsupervised home-based programs. [Luo et al., 2010] developed an interactive VR system for arm and hand rehabilitation that includes a wearable sensor system consisting of an arm suit with an Optical Linear Encoder, two Inertial Measurement Units for tracking arm motion, and a SmartGlove for tracking finger motion. They developed two VR games that increase patient motivation and enable objective evaluation of patient progress.

The VR environment integrated with eye tracking shows improvement to the rehabilitation process [Gavas et al., 2018]. This data provides a clear understating of the participant's interaction with the VR environment; the eye tracking system provides proper data that can track the vision. This data is crucial in grasping and locomotion tasks [Desanghere and Marotta, 2015]. Hence, gaze data can play a crucial role in rehabilitation and assistance of both upper and lower limb impairments [Cognolato et al., 2018]. Gaze data helps to provide support to participants while doing rehabilitation tasks [Novak and Riener, 2013]. Park et al. ([Park et al., 2019])  worked on integrating eye tracking with the VR environment for vestibular rehabilitation. The integration of eye tracking in rehabilitation parallels our approach of combining the VR environment integrated with eye tracking and endpoint robotics. Park et al. observed a significant improvement in rehabilitation outcomes with the addition of eye tracking, motivating us to utilize eye tracking data in our work.

§.§ Machine Learning for Upper Limb Motion Prediction

The integration of machine learning techniques with wearable sensor technology can enhance the efficacy of rehabilitation procedures for patients suffering from conditions requiring prolonged recovery periods, such as stroke [Wei and Wu, 2023]. [García-Vellisca et al., 2021] introduced a method for predicting hand kinematics for prosthesis control using high-density Electromyography (EMG) data. The actual hand kinematics were collected using the 18 sensors of CyberGlove II by CyberGlove Systems LLC. They employed artificial neural networks, which effectively predicted joint trajectories over 13 key hand movements. [Trigili et al., 2019] introduced a novel concept of using the signals by analyzing the behaviors of participants doing the reaching task. They then converted the sEMG signals into 14 time-domain features and used these features to train the machine learning algorithm to predict reaching tasks, go-forward, and go-backward movements. There research focused on basic task prediction without considering the motion trajectory or directional components. However, in our study, we focused on developing multimodal analysis to predict the motion intention.

Little et al. worked on the application of neural network in predicting the elbow motion trajectory focuses on dealing with different key elements for the model training process. Their work emphasised more on algorithm, which is rigid in individual variances [Little et al., 2021]. Sibo Yand et al, did similar research on the topic of motion intention prediction; they proposed a machine learning model to predict the hand positions. This work suggests that modern machine learning algorithms can effectively predict the motion intention by taking the data in a sequence type [Yang et al., 2023].

§ MATERIALS AND METHODS

The research project was approved by the Institutional Review Board (IRB). In this upcoming session, we describe briefly the participants, equipment used, data collection procedures, data preprocessing methods, and the study procedure.

§.§ Participants

The dataset consists of data from 16 healthy participants from the university of Delaware, with 16 participants (eight males), aged between 22 and 37 years (M=27.38, SD=4.13). The participants volunteered for the study and were not provided any financial compensation. Nine participants reported familiarity with video games, varying from several times a year to daily engagement, while seven had previous virtual reality experience. All participants self-reported as right-handed.

The upper extremity rehabilitation framework setup for data collection. The participant completes the circle task by holding the KinArm robot handle.

§.§ Apparatus

In this section, we explain the upper extremity rehabilitation system utilized for our experiments [Chheang et al., 2023]. The details of the setup can be found in Figure <ref>.
The virtual environment was developed using Unity game engine (version 2021.3.10f1).

Gaze data was collected at each hit point using the advanced eye-tracking technology of the HTC Vive Pro Eye tracker, providing detailed information on eye and head movements [Baron et al., 2023]

The KinArm robot, featuring six degrees of freedom with a cylindrical handle for horizontal movement, was integrated into the VR setup. The KinArm model was constructed in the virtual environment with the HTC Vive tracker capturing the handle's position and inverse kinematics calculated through simulation. The KinArm robot does not provide assistive support for the participants, and in this study we didn't focus on analysis of the data from the KinArm robot. Instead, it is used for consistent hand movement of participants in 2-Dimensional space, supporting and understand the motor functions in individuals [Thostenson et al., 2019]. KinArm also helps to provide a controlled space for motor tasks [Mochizuki et al., 2019].

Elbow movement data was obtained using a wearable sleeve sensor made of a carbon nanotube.
The sensor is incorporated into a knit fabric sleeve. Due to the piezo-resistive nature of the sensor, the resistance of the electrode changes due to the strain generated in the sensor during elbow flexion/extension. Participants used this sensor during reaching tasks. Data on resistance values were recorded using an Arduino-based voltage divider circuit, and data transmission occurred through Microsoft Excel. Each participant engaged in two study conditions: reaching tasks in the geometry of a circle and a diamond.

The list of 11 time-domain features and their explanations.

1cName of the Feature     1cEquation     1cExplanation

Integrated Absolute Value     \(IAV = \sum_{n=1}^{N} |x_n|\)     the summation of the absolute values of the resistance values between two hit positions

Mean Absolute Value     \(MAV = \frac{1}{N} \sum_{n=1}^{N} |x_n|\)     the average resistance values

Modified Mean Absolute Value 1     \(MMAV1 = \frac{1}{N} \sum_{n=1}^{N} w_n |x_n|\)     a modified version of MAV with a weight component (\(w_n\)) in the equation

\(w_n = \begin{cases}
1, & \text{if } 0.25N \leq n \leq 0.75N\\
0.5, & \text{otherwise }
\end{cases}\)

Modified Mean Absolute Value 2     \(MMAV2 = \frac{1}{N} \sum_{n=1}^{N} |x_n|\)     a modified version of MMAV1 with a different weight function

\(w_n = \begin{cases}
1, & \text{if } 0.25N \leq n \leq 0.75N\\
4n/N, & \text{if } 0.25N > n\\
4(n - N)/N, & \text{if } 0.75N < n
\end{cases}\)

Simple Square Integral     \(SSI = \sum_{n=1}^{N} x_n^2\)     the energy of the signal

Variance     \(VAR = \frac{1}{N-1} \sum_{n=1}^{N} (x_n - \mu)^2\)     the variance of the resistance values, where \(\mu\) is the mean value of the resistance

Root Mean Square     \(RMS = \sqrt{\frac{1}{N} \sum_{n=1}^{N} x_n^2}\)     the root mean square of the resistance values

Waveform Length     \(WL = \sum_{n=1}^{N} |x_{n+1} - x_n|\)     the summation of the length of the waveform of the resistance signal

Logarithm     \(LOG = \frac{1}{N} \sum_{n=1}^{N} \log_{10}(|x_n|)\)     the average of the total summation logarithm of resistance values between consecutive hit positions

Skewness     \(SKEW = \frac{\frac{1}{N} \sum_{n=1}^{N} (x_n - \mu)^3}{(\frac{1}{N} \sum_{n=1}^{N} (x_n - \mu)^2)^{3/2}}\)     a measure of imbalance in the resistance data

Kurtosis     \(KURT = \frac{\frac{1}{N} \sum_{n=1}^{N} (x_n - \mu)^4}{(\frac{1}{N-1} \sum_{n=1}^{N} (x_n - \mu)^2)^{2}}\)
a statistical equation which is used to compare the general profile of data, it reveals the outlier frequency

3p0.8\(N\): the number of sensor values generated between two hitting points.

§.§ Data Collection

Example plots of the measurements. The variation in resistance values captured by the sleeve sensor over time (a) for the circle task and (b) for the diamond tasks. The lowest values are observed in elbow extension and higher values for the flexion.

Participants were equipped with a sleeve sensor and a VR headset, and seated in front of the KinArm robot. Subsequently, participants engaged in two distinct reaching tasks: following the outline of a diamond shape and a circle. Each shape was composed of 40 points along its circumference, and participants were tasked with hitting the points sequentially. Backtracking was disallowed if a point was missed, and participants were expected to trace the shape in both clockwise and counterclockwise directions. Participants were instructed based on their group, with eight participants performing counterclockwise motions and the others executing clockwise motions.

Throughout the reaching tasks, data on resistance values were collected for each time instance, capturing the flexion/extension of muscles. Additionally, gaze data, encompassing the position and orientation of participants' heads and eyes at the moment of dot contact, was systematically recorded.

Synchronization among all the system components was ensured using packages from the Lab Streaming Layer (LSL) [Kothe et al., 2024] for the collection of synchronized data among the sensors. LSL is an open-source system designed for the collection of time series in research experiments. Besides using LSL, manual checking of the data was necessary to ensure identical start and end times for accurate data synchronization.

Figure <ref> provides visual examples of the data collected during the experiments. Figure <ref> and Figure <ref> illustrate the variation in resistance values recorded by the sleeve sensor throughout the reaching task for one participant, showing data for both the circle and diamond tasks, respectively. Additionally, technical setups displaying the angle of flexion and conceptual views of the task progress are included for better illustration.

§.§ Data Preprocessing

We utilized two distinct types of data collected from the framework: gaze data for the first model and resistance data for the second model. In this section, we detail the preprocessing steps applied to the resistance data collected through carbon nanotube sensors. The objective is to transform raw resistance values into 11 time-domain features, each serving a specific purpose in characterizing upper extremity movements [Trigili et al., 2019]. For instance, the absolute value reflects overall muscle activation levels, and the root mean square relates to the force of muscle contractions. Combining all these features provides a refined overview of muscle behavior, offering a more comprehensive analysis than a single sensor value. These features are crucial for obtaining a detailed pattern of muscle behaviors and understanding the subtle aspects of human movement. Encompassing a wide range of characteristics, from the mean to the rate of signal variations, these features enhance the accuracy and robustness of analysis using machine learning.
The comprehensive list of these features and their explanations are provided in Table <ref>.

We divided the dataset into two subsets: 80% for training and 20% for testing.  This split was applied to ensure that the model was trained on a sufficiently large portion of the data while also having an independent set for evaluating its performance.
Furthermore, to ensure consistency and comparability across different features, we applied regular min-max normalization to both the training and testing data. This standardized data was then used as input for subsequent machine learning models.

Overview of the two-step machine learning architecture. In the first step, we proposed a network for predicting segments using a neural network and gaze data. In the second step, we trained a model for predicting movement direction using an LSTM model with the output of the first model and 11 time-domain features.

§.§ Study Procedure

In this study, we propose a two-step machine learning framework designed to predict both the specific segment to which participant-marked points belong and the participant's intention to move either in a clockwise or counterclockwise direction. Figure <ref> provides an overview of our two-step machine learning architecture.  In this section, we provide the details about our two models that use the data explained in Section <ref>. Our approach integrates traditional machine learning and deep learning approaches for a comprehensive understanding of motion intention.

§.§.§ Neural Network for Segment Prediction

The neural network model is designed to predict segments, which takes gaze data as input. The architecture of the neural network consists of four layers: one input layer, two hidden layers with 64 and 32 neurons, and one output layer. We used the ReLU activation function, chosen for its efficiency compared to other activation functions [Sharma et al., 2017]. The neural network was trained using the Adam optimizer with a learning rate of 0.001, 50 epochs, and a batch size of 32. We trained the input layer of the neural network by gaze data for predicting the four segments of reaching task shapes.

§.§.§ Long Short-Term Memory (LSTM) Model for Direction Prediction

The output of our neural network model, representing the predicted segments, and 11 time-domain features of resistance data are utilized to create a specialized LSTM model. The LSTM model predicts the intended direction of movement of the participant, whether it is clockwise or counterclockwise. We chose the LSTM model for its ability to capture the temporal dependencies and patterns in the data, which are crucial in determining the direction of motion. The architecture of the LSTM model is designed as a sequential model. It consists of 40 data points arranged in sequence. The input layer takes the outputs of the neural network model and 11 time-domain features. The architecture further includes two hidden layers, each containing 50 neurons, and an output layer. We applied L2 regularization with a factor of 0.01 and used the ReLU activation function. We trained the LSTM model using the Adam optimizer with a learning rate of 0.001, 50 epochs, and a batch size of 32.

§ EXPERIMENTS AND RESULTS

In this section, we present the results of our experiments aimed at understanding the effect of different dataset setups on the performance of our two-step machine learning architecture. Our study focused on predicting four segments of reaching tasks in the first model and predicting the direction of movement in the second model. We explored various combinations of raw resistance data—11 time-domain features extracted from the raw resistance value, and gaze data—and investigated different machine learning models to identify the most effective approaches. We used all the data except raw resistance data as consecutive pairs of values sequentially. This sequential process involved deriving features from consecutive pairs of data points.  We present the details of all these different setups in Table <ref>. The size listed in the table represents the number of samples and features in each dataset, with size (samples x features). For the raw resistance data (D1), we included all 40 data points per participant, resulting in a dataset size of $640 \times 1$ ($640 = 40 \text{ data points} \times 16 \text{ participants}$).
For the remaining, due to the sequential feature extraction process, the 40 data points per participant were reduced to 39. This resulted in a total of 624 ($39 \text{ data points} \times 16 \text{ participants}$) combined data points (samples) across all participants.

Specifications of data setups used in the experiments.

1cUsed data     Name     Size

Raw resistance data     D1     640x1

11 time-domain features     D2     624x11

Gaze data     D3     624x24

Output of the first model     D4     624x4

The combination of D2&D3     D5     624x35

The combination of D2&D4     D6     624x15

The combination of D3&D4     D7     624x28

The combination of D2&D3&D4     D8     624x39

§.§ Segment Prediction Experiments

The accuracy and F1 score results of the machine learning models for predicting segments in Diamond reaching task

to 0.95X[0.4,l] *4X[c]
D1     D2     D3     D5

NN     39.67 [0.325]     49.37 [0.402]     92.31 [0.987]     91.03 [0.857]

KNN     32.50 [0.308]     41.00 [0.365]     88.60 [0.897]     85.90 [0.859]

SVM     38.40 [0.350]     42.70 [0.346]     87.20 [0.870]     83.34 [0.830]

LR     31.25 [0.254]     41.00 [0.344]     87.12 [0.911]     84.20 [0.935]

to 0.95@ X @
All entities are in the format: Accuracy (%) [F1 Score]. A higher Accuracy/F1 score indicates a better performance.

\(D1\): Raw resistance data, \(D2\): 11 time-domain features, \(D3\): Gaze data, \(D5\): D2&D3

The accuracy and F1 score results of the machine learning models for predicting segments in Circle reaching task

to 0.95X[0.4,l] *4X[c]
D1     D2     D3     D5

NN     38.05 [0.258]     47.25 [0.451]     94.87 [0.884]     85.06 [0.854]

KNN     36.25 [0.349]     37.80 [0.577]     87.10 [0.946]     76.92 [0.894]

SVM     25.76 [0.317]     33.57 [0.570]     87.55 [0.797]     84.28 [0.922]

LR     36.25 [0.230]     23.08 [0.422]     86.76 [0.823]     84.61 [0.962]

to 0.95@ X @
All entities are in the format: Accuracy (%) [F1 Score]. A higher Accuracy/F1 score indicates a better performance.

\(D1\): Raw resistance data, \(D2\): 11 time-domain features, \(D3\): Gaze data, \(D5\): D2&D3

Random Guess Prediction Results - Baseline for Task Segmentation and Direction Prediction

1cPrediction problem     Accuracy (%)

Segment prediction for diamond     26.44

Segment prediction for circle     25.77

Direction prediction for diamond     50.71

Direction prediction for circle     51.90

We examined various data setups (D1, D2, D3, and D5) and machine learning techniques in our first set of experiments focusing on predicting four segments of the reaching tasks. We trained models using four different machine learning methods, namely Neural Network (NN), k-Nearest Neighbors (KNN), Support Vector Machine (SVM), and Logistic Regression (LR). We used four different data setups, to explore their impact on the models' performance.

For two reaching task shapes, diamond and circle, we listed the accuracy results and F1 scores of the trained four different machine learning models obtained with the different data setups in Table <ref> and Table <ref>, respectively. All entities in these tables are presented in the format "Accuracy [F1 Score]", where the F1 Score is a metric that offers a balanced measure of a model's performance. We also compared our results with a random guess baseline. The random guess accuracy for segment prediction tasks, as shown in Table <ref>, is around 25% for both the Diamond and Circle. In addition to our selected NN model, the other machine learning models we experimented with also achieved better results than the baseline in most data setups. This further indicates the suitability of this problem for machine learning approaches.

After analyzing the accuracy and F1 score results presented in Tables <ref> and <ref>, we chose the NN model trained with gaze data. The NN model trained with gaze data demonstrated higher accuracy and F1 scores for both the diamond and circle reaching tasks, achieving 92.31% [0.987] and 94.87% [0.884], respectively. These results show the critical role of gaze data, as it consistently contributed to the models' best performance.  Incorporating gaze data in D3 and D5 provides supplementary cues that enhance the models' predictive abilities, thus resulting in improved performance compared to setups lacking gaze data (D1 and D2). The differences in performance between D3 and D5 highlight the possibility that integrating features from different modalities could introduce complexities or redundancies, potentially affecting the overall performance of the models.

§.§ Direction Prediction Experiments

In this section, we present the results of our experiments focused on predicting the direction of movements, as clockwise or counterclockwise direction. We explored all data combinations shown in Table <ref>, and applied multiple machine learning techniques to determine the one with the highest accuracy.
We trained models with four different machine learning methods: LSTM, KNN, SVM, and LR.

Tables <ref> and <ref> present both the accuracy and F1 score results of the machine learning models for predicting the direction of movements in the diamond and circle reaching tasks, respectively. These tables provide the accuracy and F1 scores achieved by each model for different data setups. The highest accuracy and F1 score values, highlighted in bold, identify the most successful models for predicting tasks.

The accuracy and F1 score results of the machine learning models for predicting the direction in Diamond reaching task

to 0.95X[0.2,l] *4X[c]
LSTM     KNN     SVM     LR

D1     52.50 [0.665]     62.50  [0.576]     50.00 [0.666]     50.00 [0.667]

D2     87.18 [0.796]     94.44 [0.939]     71.79  [0.731]     74.36 [0.750]

D3     76.92 [0.698]      47.44 [0.782]     58.97 [0.620]     62.82 [0.701]

D4     56.25 [0.674]     57.50 [0.512]     50.00 [0.666]     50.00 [0.667]

D5     73.08 [0.887]     78.21 [0.798]     79.49 [0.789]     71.79 [0.755]

D6     96.72 [0.946]     76.92 [0.833]     75.64 [0.769]     74.36 [0.725]

D7     94.44 [0.939]     43.59 [0.435]     60.26 [0.629]     67.95 [0.721]

D8     96.15 [0.896]     70.51 [0.794]     65.38 [0.761]     71.79 [0.769]

to 0.95@ X @
All entities are in the format: Accuracy (%) [F1 Score]. A higher Accuracy/F1 score indicates a better performance.

\(D1\):  Raw resistance data, \(D2\): 11 time-domain features, \(D3\):  Gaze data, \(D4\): Output of the first model,

\(D5\): D2&D3, \(D6\): D2&D4, \(D7\): D3&D4, \(D8\): D2&D3&D4

The accuracy and F1 score results of the machine learning models for predicting the direction in Circle reaching task

to 0.95X[0.2,l] *4X[c]
LSTM     KNN     SVM     LR

D1     64.7 [0.529]     44.75 [0.473]     52.80 [0.570]     53.75 [0.539]

D2     89.74 [0.720]     51.28 [0.538]     76.19 [0.712]     56.41 [0.727]

D3     61.54 [0.602]     51.28 [0.525]     61.53 [0.646]     46.15 [0.502]

D4     57.67 [0.678]     48.75 [0.518]     51.34 [0.472]     51.21 [0.460]

D5     89.74 [0.917]     67.94 [0.679]     42.37 [0.377]     35.90 [0.307]

D6     97.44 [0.923]     53.84 [0.577]     74.35 [0.739]     58.97 [0.725]

D7     91.03 [0.894]     61.53 [0.602]     71.79 [0.741]     56.41 [0.733]

D8     96.15 [0.894]     66.67 [0.699]     51.28 [0.458]     46.15 [0.337]

to 0.95@ X @
All entities are in the format: Accuracy (%) [F1 Score]. A higher Accuracy/F1 score indicates a better performance.

\(D1\):  Raw resistance data, \(D2\): 11 time-domain features, \(D3\):  Gaze data, \(D4\): Output of the first model,

\(D5\): D2&D3, \(D6\): D2&D4, \(D7\): D3&D4, \(D8\): D2&D3&D4

Analysis of the results shows that the LSTM model trained with the combination of D2&D4 data is the optimal choice. This model demonstrated superior accuracy and F1 score, reaching 96.72% [0.946] in the diamond reaching task and 97.44% [0.923] in the Circle reaching task. The incorporation of both 11 time-domain-specific resistance features and the output of the previous segment prediction model (D2&D4) significantly improved the accuracy and F1 scores of direction predictions for both task shapes.

Similar to the segment prediction problem, in this context, machine learning results were compared with a random guess baseline. The random guess accuracy for direction prediction tasks, presented in Table <ref>, is approximately 50% for both diamond and circle tasks. Similar to the segment prediction experiments, the machine learning models for direction prediction outperformed the random guess baseline in most data setups.

§ DISCUSSION

In this study, we propose a novel two-step machine learning architecture for predicting the motion intention of participants engaged in upper extremity rehabilitation tasks. The framework we used for extracting data integrated VR with the KinArm robot, and we utilized gaze data and resistance values to enhance the precision of motion intention predictions. The KinArm robot is primarily utilized for support in the rehabilitation, playing a role in improving the process without directly contributing to prediction analysis.

Our first model, based on a neural network, successfully predicted the four reaching task segments. Our experiments demonstrate the importance of incorporating gaze data for accurate predictions. Eye movement information significantly improved the model's performance, highlighting the relevance of multi-modal data integration in rehabilitation studies.

Our second model, an LSTM model, accurately determined the participant's intention to move in a clockwise or counterclockwise direction. The integration of 11 time-domain features extracted from resistance data, combined with the output of the first model, achieved high accuracy in predicting bidirectional rotational behavior. The 11 time-domain features carry a wide range of statistical and temporal properties that capture the dynamic nature of resistance over time.

While some may question whether more complex models could offer better alternatives for this prediction problem, our approach highlights that, despite its simplicity, basic machine learning techniques can reveal valuable insights. By leveraging features, our models effectively capture essential information for prediction motion intentions in upper extremity rehabilitation tasks.

Limitations and Future Directions
While our study provides useful insights into upper extremity movements using the proposed machine learning architecture, we must acknowledge some limitations. Firstly, the performance of our models depends on the quality and quantity of the data used, which may be impacted by sensor noise or data collection errors. Secondly, our study was focused on healthy, non-clinical participants. So, our future work will include validation with clinical populations to assess the applicability and generalizability of our approach in rehabilitation settings.

§ CONCLUSION

This study presents a comprehensive approach for personalized and effective rehabilitation practices, contributing to improvements in integrating robotics, VR, and biomechanical data analysis. Our study and findings demonstrate the potential of incorporating machine learning techniques to improve the effectiveness of virtual rehabilitation assessment. For instance, our NN model achieved an impressive accuracy of 92.31% in predicting segment tasks for the Diamond reaching task and 94.87% for the Circle reaching task when trained with gaze data. Furthermore, our LSTM model, when trained with a combination of 11 time-domain features and the output of the segment prediction model, achieved accuracy rates of 96.72% and 97.44% for the Diamond and Circle reaching tasks, respectively. There are exciting possibilities for our research to explore. This research introduces an attractive application for real-time predictions, enabling robotic arms to obtain feedback from our proposed two-step model to guide participants accurately based on their motion intentions without diverting from the actual task.

§ ACKNOWLEDGMENTS

We wish to acknowledge the support from our sponsors, the National Science Foundation (2222663), and the National Institute of General Medical Sciences of the National Institutes of Health (P20 GM103446 and P20 GM103446E). We also express our gratitude to Drs. Vuthea Chheang, Joshua Cashback, and Erik Thostenson for their invaluable efforts in ideation, data collection, and resource sharing. We also extend our sincere appreciation to the entire research team and the study participants for their contributions. Any opinions, findings, conclusions, or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the sponsors.

[Aderinto et al., 2023]
Nicholas Aderinto, Gbolahan Olatunji, Muili Opeyemi Abdulbasit, Mariam Edun, Gbolahan Aboderin, and Emmanuel Egbunu.
Exploring the efficacy of virtual reality-based rehabilitation in stroke: a narrative review of current evidence.
Annals of Medicine, 550 (2):0 2285907, 2023.

[Anwar et al., 2021]
Naveed Anwar, Hossein Karimi, Ashfaq Ahmad, Nazia Mumtaz, Ghulam Saqulain, and Syed Amir Gilani.
A novel virtual reality training strategy for poststroke patients: a randomized clinical trial.
Journal of Healthcare Engineering, 2021:0 1–6, 2021.

[Araújo et al., 2020]
Nuno Miguel Faria Araújo et al.
Impact of the fourth industrial revolution on the health sector: a qualitative study.
Healthcare informatics research, 260 (4):0 328–334, 2020.

[Baron et al., 2021]
Lauren Baron, Qile Wang, Sydney Segear, Brian A Cohn, Kangsoo Kim, and Roghayeh Barmaki.
Enjoyable physical therapy experience with interactive drawing games in immersive virtual reality.
In Proceedings of the 2021 ACM Symposium on Spatial User Interaction, pages 1–8, 2021.

[Baron et al., 2023]
Lauren Baron, Vuthea Chheang, Amit Chaudhari, Arooj Liaqat, Aishwarya Chandrasekaran, Yufan Wang, Joshua Cashaback, Erik Thostenson, and Roghayeh Leila Barmaki.
Virtual therapy exergame for upper extremity rehabilitation using smart wearable sensors.
In Proceedings of the 8th ACM/IEEE International Conference on Connected Health: Applications, Systems and Engineering Technologies, pages 92–101, 2023.

[Caviedes et al., 2020]
Jorge E Caviedes, Baoxin Li, and Varun C Jammula.
Wearable sensor array design for spine posture monitoring during exercise incorporating biofeedback.
IEEE Transactions on Biomedical Engineering, 670 (10):0 2828–2838, 2020.

[Chheang et al., 2023]
Vuthea Chheang, Rakshith Lokesh, Amit Chaudhari, Qile Wang, Lauren Baron, Behdokht Kiafar, Sagar Doshi, Erik Thostenson, Joshua Cashaback, and Roghayeh Leila Barmaki.
Immersive virtual reality and robotics for upper extremity rehabilitation.
arXiv preprint arXiv:2304.11110, 2023.

[Cognolato et al., 2018]
Matteo Cognolato, Manfredo Atzori, and Henning Müller.
Head-mounted eye gaze tracking devices: An overview of modern devices and recent advances.
Journal of rehabilitation and assistive technologies engineering, 5:0 2055668318773991, 2018.

[Daunoraviciene et al., 2018]
Kristina Daunoraviciene, Ausra Adomaviciene, Agne Grigonyte, Julius Griškevičius, and Alvydas Juocevicius.
Effects of robot-assisted training on upper limb functional recovery during the rehabilitation of poststroke patients.
Technology and Health Care, 260 (S2):0 533–542, 2018.

[Desanghere and Marotta, 2015]
Loni Desanghere and Jonathan J Marotta.
The influence of object shape and center of mass on grasp and gaze.
Frontiers in psychology, 6:0 1537, 2015.

[Dixit and Tedla, 2019]
Snehil Dixit and Jaya Shanker Tedla.
Effectiveness of robotics in improving upper extremity functions among people with neurological dysfunction: a systematic review.
International Journal of Neuroscience, 1290 (4):0 369–383, 2019.

[Dobkin and Dorsch, 2011]
Bruce H Dobkin and Andrew Dorsch.
The promise of mhealth: daily activity monitoring and outcome assessments by wearable sensors.
Neurorehabilitation and neural repair, 250 (9):0 788–798, 2011.

[Durandau et al., 2019]
Guillaume Durandau, Dario Farina, Guillermo Asín-Prieto, Iris Dimbwadyo-Terrer, Sergio Lerma-Lara, Jose L Pons, Juan C Moreno, and Massimo Sartori.
Voluntary control of wearable robotic exoskeletons by patients with paresis via neuromechanical modeling.
Journal of neuroengineering and rehabilitation, 16:0 1–18, 2019.

[García-Vellisca et al., 2021]
MA García-Vellisca, Ana Matran-Fernandez, Riccardo Poli, and Luca Citi.
Hand-movement prediction from emg with lstm-recurrent neural networks.
In 2021 Global Medical Engineering Physics Exchanges/Pan American Health Care Exchanges (GMEPE/PAHCE), pages 1–5. IEEE, 2021.

[Gavas et al., 2018]
Rahul Dasharath Gavas, Sangheeta Roy, Debatri Chatterjee, Soumya Ranjan Tripathy, Kingshuk Chakravarty, and Aniruddha Sinha.
Enhancing the usability of low-cost eye trackers for rehabilitation applications.
PloS one, 130 (6):0 e0196348, 2018.

[Karamians et al., 2020]
Reneh Karamians, Rachel Proffitt, David Kline, and Lynne V Gauthier.
Effectiveness of virtual reality-and gaming-based interventions for upper extremity rehabilitation poststroke: a meta-analysis.
Archives of physical medicine and rehabilitation, 1010 (5):0 885–896, 2020.

[Kim et al., 2023]
Sung-Hoon Kim, Dong-Min Ji, In-Su Hwang, Jinwhan Ryu, Sol Jin, Soo-A Kim, and Min-Su Kim.
Three-dimensional magnetic rehabilitation, robot-enhanced hand-motor recovery after subacute stroke: A randomized controlled trial.
Brain Sciences, 130 (12):0 1685, 2023.

[Kothe et al., 2024]
Christian Kothe, Clemens Brunner, Joshua Podmore, Tim Mullen, Benjamin Blankertz, et al.
<https://github.com/sccn/labstreaminglayer>, 2024.
Accessed: March 13, 2024.

[Kwakkel et al., 2019]
G Kwakkel, Erwin EH van Wegen, Jane H Burridge, CJ Winstein, Liesjet Elisabeth Henriette Van Dokkum, Margit Alt Murphy, MF Levin, and JW Krakauer.
Standardized measurement of quality of upper limb movement after stroke: consensus-based core recommendations from the second stroke recovery and rehabilitation roundtable.
Neurorehabilitation and neural repair, 330 (11):0 951–958, 2019.

[Little et al., 2021]
Kieran Little, Bobby K Pappachan, Sibo Yang, Bernardo Noronha, Domenico Campolo, and Dino Accoto.
Elbow motion trajectory prediction using a multi-modal wearable system: A comparative analysis of machine learning techniques.
Sensors, 210 (2):0 498, 2021.

[Luo et al., 2010]
Zhiqiang Luo, Chee Kian Lim, Weiting Yang, Ke Yen Tee, Kang Li, Chao Gu, Kim Doang Nguen, I-Ming Chen, and Song Huat Yeo.
An interactive therapy system for arm and hand rehabilitation.
In 2010 IEEE Conference on Robotics, Automation and Mechatronics, pages 9–14. IEEE, 2010.

[Mochizuki et al., 2019]
George Mochizuki, Andrew Centen, Myles Resnick, Catherine Lowrey, Sean P Dukelow, and Stephen H Scott.
Movement kinematics and proprioception in post-stroke spasticity: assessment using the kinarm robotic exoskeleton.
Journal of neuroengineering and rehabilitation, 16:0 1–13, 2019.

[Mubin et al., 2019]
Omar Mubin, Fady Alnajjar, Nalini Jishtu, Belal Alsinglawi, and Abdullah Al Mahmud.
Exoskeletons with virtual reality, augmented reality, and gamification for stroke patients’ rehabilitation: systematic review.
JMIR rehabilitation and assistive technologies, 60 (2):0 e12010, 2019.

[Novak and Riener, 2013]
Domen Novak and Robert Riener.
Enhancing patient freedom in rehabilitation robotics using gaze-based intention detection.
In 2013 IEEE 13th international conference on rehabilitation robotics (ICORR), pages 1–6. IEEE, 2013.

[Park et al., 2019]
Jeong Hye Park, Han Jae Jeon, Eun-Cheon Lim, Ja-Won Koo, Hyo-Jeong Lee, Hyung-Jong Kim, Jung Seop Lee, Chang-Geun Song, and Sung Kwang Hong.
Feasibility of eye tracking assisted vestibular rehabilitation strategy using immersive virtual reality.
Clinical and experimental otorhinolaryngology, 120 (4):0 376, 2019.

[Saggio et al., 2015]
Giovanni Saggio, Francesco Riillo, Laura Sbernini, and Lucia Rita Quitadamo.
Resistive flex sensors: a survey.
Smart Materials and Structures, 250 (1):0 013001, 2015.

[Sharma et al., 2017]
Sagar Sharma, Simone Sharma, and Anidhya Athaiya.
Activation functions in neural networks.
Towards Data Sci, 60 (12):0 310–316, 2017.

[Tarnita et al., 2022]
Daniela Tarnita, Ionut Daniel Geonea, Doina Pisla, Giuseppe Carbone, Bogdan Gherman, Nicoleta Tohanean, Paul Tucan, Cristian Abrudan, and Danut Nicolae Tarnita.
Analysis of dynamic behavior of parreex robot used in upper limb rehabilitation.
Applied Sciences, 120 (15):0 7907, 2022.

[Thostenson et al., 2019]
Erik T Thostenson, Sagar Doshi, and Amit Chaudhari.
Flexible and wearable sensors for human motion analysis.
In Abstracts of Papers of the American Chemical Society, volume 258. AMER CHEMICAL SOC 1155 16TH ST, NW, Washington, DC 20036 USA, 2019.

[Trigili et al., 2019]
Emilio Trigili, Lorenzo Grazi, Simona Crea, Alessandro Accogli, Jacopo Carpaneto, Silvestro Micera, Nicola Vitiello, and Alessandro Panarese.
Detection of movement onset using emg signals for upper-limb exoskeletons in reaching tasks.
Journal of neuroengineering and rehabilitation, 16:0 1–16, 2019.

[Vibhuti et al., 2023]
Vibhuti, Neelesh Kumar, and Chitra Kataria.
Efficacy assessment of virtual reality therapy for neuromotor rehabilitation in home environment: a systematic review.
Disability and Rehabilitation: Assistive Technology, 180 (7):0 1200–1220, 2023.

[Walker et al., 2023]
Michael Walker, Thao Phung, Tathagata Chakraborti, Tom Williams, and Daniel Szafir.
Virtual, augmented, and mixed reality for human-robot interaction: A survey and virtual design element taxonomy.
ACM Transactions on Human-Robot Interaction, 120 (4):0 1–39, 2023.

[Wei and Wu, 2023]
Suyao Wei and Zhihui Wu.
The application of wearable sensors and machine learning algorithms in rehabilitation training: A systematic review.
Sensors, 230 (18):0 7667, 2023.

[Wonsick and Padir, 2020]
Murphy Wonsick and Taskin Padir.
A systematic review of virtual reality interfaces for controlling and interacting with robots.
Applied Sciences, 100 (24):0 9051, 2020.

[Yang et al., 2023]
Sibo Yang, Neha P Garg, Ruobin Gao, Meng Yuan, Bernardo Noronha, Wei Tech Ang, and Dino Accoto.
Learning-based motion-intention prediction for end-point control of upper-limb-assistive robots.
Sensors, 230 (6):0 2998, 2023.


#  Lowering PyTorch’s Memory Consumption for Selective Differentiation

Samarth Bhatia
IIT Delhi
<EMAIL_ADDRESS>Felix Dangel
Vector Institute
<EMAIL_ADDRESS>
###### Abstract

Memory is a limiting resource for many deep learning tasks. Beside the neural
network weights, one main memory consumer is the computation graph built up by
automatic differentiation (AD) for backpropagation. We observe that PyTorch’s
current AD implementation neglects information about parameter
differentiability when storing the computation graph. This information is
useful though to reduce memory whenever gradients are requested for a
parameter subset, as is the case in many modern fine-tuning tasks.
Specifically, inputs to layers that act linearly in their parameters (dense,
convolution, or normalization layers) can be discarded whenever the parameters
are marked as non-differentiable. We provide a drop-in, differentiability-
agnostic implementation of such layers and demonstrate its ability to reduce
memory without affecting run
time111Code:github.com/plutonium-239/memsave_torch..

##  1 Introduction & Motivation

The success of many deep learning applications is driven by scaling
computational resources [32]. One important resource is GPU memory,
specifically on low- and mid-end GPUs which usually offer between 6 to 16 GiB.
Therefore, down-scaling the computational demands of deep learning is an
important objective to widen its accessibility to researchers and
practitioners with fewer hardware resources.

Two major memory consumers are the neural network weights, and the computation
graph stored by the automatic differentiation (AD) engine. There exist various
approaches to reduce them; e.g., the parameters can be compressed with low-
precision data types (quantization [16, 19, 21]) or sparsified [12, 7], and
the computation graph can be off-loaded to CPU [26], compressed [3],
randomized [22], or partially recorded and re-computed (gradient checkpointing
[10, 5]).

Figure 1: _PyTorch’s AD is not agnostic to parameter differentiability._ We
consider a deep CNN that consists entirely of size-preserving convolutions and
measure the forward pass’s peak memory when processing a mini-batch of size
(256, 8, 256, 256), requiring 512 MiB storage. Memory increases linearly in
the number of layers when all parameters are marked differentiable and remains
constant when all parameters are marked non-differentiable. Surprisingly,
_when only one layer’s parameters are marked as differentiable the memory
increases as if all subsequent parameters were marked differentiable_. Our
drop-in solution stores layer inputs depending on the layer parameter’s
differentiability and reduces memory compared to the current PyTorch
implementation.

Here, we consider a special AD scenario we call _selective differentiation_
where gradients are requested only for a subset of the computation graph leafs
(the neural network inputs and parameters). This approach has gained a lot of
practical relevance in the era of large foundation models. For instance, fine-
tuning techniques for pre-trained models rely on training only a subset of
layers [36, 18], or extend some with parameter-efficient adapters [e.g. LoRA
15] which are then trained, keeping the original weights fixed. But it is also
common in other applications such as generating adversarial examples [9] and
neural style transfer [8] which optimize the input to a network with frozen
weights.

We find that PyTorch’s [23] current AD allows for additional, simple,
optimizations to further reduce memory consumption in the context of selective
differentiation:

  1. 1.

We observe that PyTorch’s AD neglects the differentiability of layer
parameters when storing the computation graph (Figure 1). This information is
useful though as it allows discarding inputs to linear layers whose parameters
are marked as non-differentiable.

  2. 2.

We provide a drop-in implementation of various layers that is agnostic to the
parameter differentiability and demonstrate on various convolutional neural
networks (CNNs) that it lowers on the default implementation’s memory
footprint without increasing run time.

This easy-to-use insight benefits many tasks with selective differentiation
and enables them to scale further. We hope it will stimulate further research
into AD optimizations.

##  2 Selective Differentiation in PyTorch

PyTorch users can specify whether they want to compute gradients w.r.t. a
tensor through its requires_grad attribute, which is dominantly inherited by
child tensors: If any input (parent) to an operation is marked differentiable,
the output (child) will also be differentiable and require the computation
graph to be stored for backpropagation. As we will see, PyTorch does not check
the differentiability of all parents, but rather starts storing the
computation graph as if all parents were differentiable once it encounters a
differentiable input. In the following example, we demonstrate this behavior,
and how it misses out on possible optimizations for linear operations with
non-differentiable parameters.

##### Experimental procedure:

We use PyTorch 2.2.1 and measure the forward pass’s peak memory of different
neural networks as a proxy for the computation graph size stored by the AD
engine. We chose peak memory as it is a reasonable proxy for the computation
graph size and the relevant metric for causing out-of-memory errors in
practise. On CPU, we use the memory_profiler library [24]. On GPU, we rely on
torch.cuda.max_memory_allocated(), the maximum memory allocated by CUDA. Each
measurement is performed in a separate Python session to avoid memory
(de)allocation leaks between consecutive runs.

###  2.1 A Simple Example

We start with a synthetic example to probe the internals of PyTorch’s AD that
are responsible for identifying and storing the computation graph during a
forward pass (summarized in Figure 2). We consider a deep CNN consisting
entirely of convolutions without bias. Each convolution preserves its input
size (kernel size 3, unit padding and stride) and number of channels and we
vary its depth as well as the parameter differentiability. As input, we choose
a mini-batch of size (256, 8, 256, 256). Storing this tensor, and each
intermediate output generated by a layer, requires 512 MiB in single
precision. Memory consumed by the convolution kernels of shape (8,8,3,3) is
negligible compared to the hidden features. Figure 1 summarizes our findings.

First, we investigate the computation graph size when marking all or no
parameters as differentiable. When all parameters are differentiable, all
layer inputs must be stored to compute gradients. Consequently, we observe a
linear relation with a slope corresponding to the 512 MiB consumed by each
intermediate input (Figure 1, _fully differentiable_). When no parameter is
differentiable, the memory consumption flattens after more than two layers
(Figure 1, _fully non-differentiable_). This is because during a forward pass
both input and output tensors of a layer are allocated in memory, in addition
to the network’s input. Hence, at most two tensors are in memory at a time for
a single layer, while at most three are allocated for two or more layers.

$\mathbf{W}$$\mathbf{x}$$\mathbf{z}$$\mathbf{W}$$\mathbf{x}$$\mathbf{z}$PyTorch$\mathbf{W}$$\mathbf{x}$$\mathbf{z}$$\mathbf{W}$$\mathbf{x}$$\mathbf{z}$Ours/:
Non-/Differentiable,: Stored by AD Figure 2: _PyTorch’s current behaviour of
storing the computation graph._ We illustrate this on a linear layer
$\mathbf{z}=\mathbf{W}\mathbf{x}$. PyTorch stores the layer input whenever it
is differentiable, although this is not necessary if the weight does not
require gradients. Our approach uses this information to discard the layer
input if possible.

Next, we observe the memory consumption in the context of selective
differentiation. We consider two scenarios: In the first, all parameters after
and including the fourth layer are differentiable, hence all layer inputs
after the third layer must be stored in memory. In the second, only the fourth
layer is differentiable and therefore only that layer’s input is necessary to
compute gradients. However, we observe that both scenarios exhibit the _same_
memory footprint (Figure 1, _layers 4+ differentiable_ & _layer 4
differentiable_)!

We conclude that PyTorch stores a layer input with requires_grad = True
although the layer’s parameters might be non-differentiable, and therefore not
require it to be stored. This information can be useful to reduce the
information stored for backpropagation, as we show with our implementation
(Figure 1, _layer 4 differentiable (ours)_).

| Memory [GiB]
---|---
Case | All | Input | Conv | Norm
(GPU) Default ResNet-101 | 7.82 (1.00) | 7.82 (1.00) | 7.82 (1.00) | 7.82 (1.00)
\+ swap Convolution | 7.82 (1.00) | 7.78 (0.99) | 7.82 (1.00) | 7.78 (0.99)
\+ swap BatchNorm, ReLU | 8.75 (1.12) | 1.61 (0.21) | 4.84 (0.62) | 5.44 (0.70)
(CPU) Default ResNet-101 | 8.24 (1.00) | 8.24 (1.00) | 8.24 (1.00) | 8.26 (1.00)
\+ swap Convolution | 8.35 (1.01) | 8.30 (1.01) | 8.33 (1.01) | 8.26 (1.00)
\+ swap BatchNorm, ReLU | 9.37 (1.14) | 2.52 (0.31) | 5.53 (0.67) | 6.37 (0.77)
| Time [s]
(GPU) Default ResNet-101 | 1.55 (1.00) | 1.39 (0.90) | 1.57 (1.01) | 1.40 (0.90)
\+ swap Convolution | 1.50 (0.97) | 1.37 (0.88) | 1.50 (0.97) | 1.36 (0.88)
\+ swap BatchNorm, ReLU | 1.46 (0.94) | 1.32 (0.85) | 1.45 (0.94) | 1.30 (0.84)
(CPU) Default ResNet-101 | 22.78 (1.00) | 16.81 (0.78) | 22.84 (1.00) | 16.66 (0.73)
\+ swap Convolution | 23.05 (1.01) | 17.08 (0.75) | 23.14 (1.02) | 16.90 (0.74)
\+ swap BatchNorm, ReLU | 25.34 (1.11) | 16.79 (0.74) | 23.68 (1.04) | 17.57 (0.77)
Table 1: _Peak memory and run time comparison between PyTorch and our memory-
saving layers_ on ResNet-101. Normalized values are relative to ‘All’ with
PyTorch’s default layers.

Bottleneck | (64, 56, 56)
---|---
— Conv2d | (64, 56, 56)
— BatchNorm2d | (64, 56, 56)
— ReLU | (64, 56, 56)
— Conv2d | (64, 56, 56)
— BatchNorm2d | (64, 56, 56)
— ReLU | (64, 56, 56)
— Conv2d | (64, 56, 56)
— BatchNorm2d | (256, 56, 56)
— Sequential | (64, 56, 56)
— — Conv2d | (64, 56, 56)
— — BatchNorm2d | (256, 56, 56)
— ReLU | (256, 56, 56)
Table 2: _Details of one residual block of ResNet101._ Shapes of layer inputs
(excluding the batch sizes) are shown on the right. The output of a ReLU layer
is stored by both the activation layer itself, and the convolution layer it
feeds into.

###  2.2 Implementation Details

Our implementation of memory-saving layers is straightforward and does not
require low-level code. For each layer we create a new torch.autograd.Function
AD primitive, and its associated torch.nn.Module layer. The primitive uses the
same forward and backward routines as the original operation (from
torch.nn.functional and torch.ops.aten), but has additional logic for the
information that is saved to the AD tape, which we describe below. Hence, our
implementation shares the performance of PyTorch’s. We also provide a
converter that replaces supported layers of a net with our memory-saving
equivalents.

####  2.2.1 Linear Layers

Consider a dense layer $\mathbf{z}=\mathbf{W}\mathbf{x}+\mathbf{b}$ with
vector input $\mathbf{x}$ and output $\mathbf{z}$, weight matrix $\mathbf{W}$,
and bias vector $\mathbf{b}$. Its input Jacobian depends on $\mathbf{W}$, the
weight Jacobian on $\mathbf{x}$, and the bias Jacobian has no dependency [e.g.
6, Chapter 2.3]. During the forward pass, we check the differentiability of
$\mathbf{W},\mathbf{x}$, and only store tensors required by the Jacobians that
will be applied during backprogation. The same dependency pattern holds for
other layers that process their inputs linearly w.r.t. their weight and add a
bias term, such as convolution [see 2] and normalization [1, 33, 34, 17]
layers. We implement them in exactly the same fashion.

####  2.2.2 Interaction with Other Layers

Real neural nets contain additional layers that are interleaved with linear
layers: activations, dropout, pooling, etc. The information stored by such
layers may overlap with the input to a linear layer, e.g. if an activation
layer feeds into a convolution. Depending on the implementation details, our
described optimizations might not apply in such cases because the linear
layer’s input tensor could still be stored by the preceding activation. We
encountered this effect for ReLU layers in real-world CNNs (Section 3), but
were able to overcome it using a customized ReLU implementation that saves a
boolean mask of the output instead. As PyTorch currently only supports 1-byte
booleans, this leads to a 4x reduction of the stored tensor’s size, which
could be further reduced to 32x with a 1-bit boolean implementation (which is
in the works).

This underlines an important challenge for saving memory in selective
differentiation whenever multiple layers use the same tensor for
backpropagation. Inputs to linear layers, even if their parameters are non-
differentiable, might still be stored by neighboring layers. We believe it
should often be possible to resolve this scenario—albeit through careful
implementation—e.g. whenever the Jacobian can be implemented by either storing
the layer input or output, as is the case for various activation functions,
and pooling layers.

##  3 Real-World Examples

| Memory [GiB] | Time [s]
---|---|---
Case | All | Input | Conv | Norm | All | Input | Conv | Norm
DeepLabv3 (RN101) [4] | 22.78 (1.00) | 22.78 (1.00) | 22.78 (1.00) | 22.78 (1.00) | 2.00 (1.00) | 1.76 (0.88) | 2.00 (1.00) | 1.76 (0.88)
\+ MemSave | 24.90 (1.09) | 4.28 (0.17) | 13.72 (0.55) | 15.17 (0.61) | 1.97 (0.99) | 1.76 (0.88) | 1.99 (1.00) | 1.77 (0.89)
EfficientNetv2-L [30, 31] | 26.83 (1.00) | 26.83 (1.00) | 26.83 (1.00) | 26.83 (1.00) | 1.79 (1.00) | 1.59 (0.89) | 1.79 (1.00) | 1.57 (0.88)
\+ MemSave | 26.83 (1.00) | 10.45 (0.39) | 18.62 (0.69) | 18.64 (0.69) | 1.66 (0.93) | 1.45 (0.87) | 1.63 (0.91) | 1.45 (0.87)
FCN (RN101) [20] | 22.19 (1.00) | 22.19 (1.00) | 22.19 (1.00) | 22.19 (1.00) | 1.94 (1.00) | 1.72 (0.89) | 1.94 (1.00) | 1.71 (0.88)
\+ MemSave | 24.41 (1.10) | 4.26 (0.17) | 13.42 (0.55) | 15.14 (0.62) | 1.90 (0.98) | 1.70 (0.88) | 1.91 (0.98) | 1.68 (0.87)
MobileNetv3-L [28, 14] | 2.80 (1.00) | 2.80 (1.00) | 2.80 (1.00) | 2.80 (1.00) | 1.46 (1.00) | 1.35 (0.92) | 1.48 (1.01) | 1.35 (0.92)
\+ MemSave | 2.94 (1.05) | 0.87 (0.30) | 1.90 (0.65) | 1.91 (0.65) | 1.44 (0.99) | 1.30 (0.90) | 1.44 (0.99) | 1.31 (0.90)
ResNeXt101-64x4d [35] | 15.18 (1.00) | 15.18 (1.00) | 15.18 (1.00) | 15.18 (1.00) | 1.71 (1.00) | 1.53 (0.89) | 1.68 (0.98) | 1.53 (0.89)
\+ MemSave | 16.75 (1.10) | 2.46 (0.15) | 9.28 (0.55) | 9.89 (0.59) | 1.59 (0.93) | 1.44 (0.84) | 1.61 (0.94) | 1.43 (0.84)
VGG-16 [29] | 4.93 (1.00) | 4.93 (1.00) | 4.93 (1.00) | N/A | 1.50 (1.00) | 1.33 (0.89) | 1.49 (0.99) | N/A
\+ MemSave | 4.30 (0.87) | 2.60 (0.60) | 4.30 (0.87) | N/A | 1.45 (0.97) | 1.35 (0.90) | 1.45 (0.97) | N/A
Faster-RCNN (RN101) [27] | 6.62 (1.00) | 6.56 (0.99) | 6.57 (0.99) | 6.53 (0.99) | 1.58 (1.00) | 1.47 (0.93) | 1.61 (1.02) | 1.46 (0.92)
\+ MemSave | 7.21 (1.09) | 1.90 (0.26) | 4.55 (0.63) | 4.43 (0.61) | 1.51 (0.96) | 1.39 (0.88) | 1.54 (0.97) | 1.38 (0.87)
SSDLite (MobileNetv3-L) [28] | 0.53 (1.00) | 0.53 (1.00) | 0.53 (1.00) | 0.53 (1.00) | 1.54 (1.00) | 1.39 (0.90) | 1.53 (0.99) | 1.40 (0.91)
\+ MemSave | 0.55 (1.04) | 0.26 (0.48) | 0.40 (0.73) | 0.41 (0.74) | 1.48 (0.96) | 1.35 (0.88) | 1.50 (0.97) | 1.34 (0.88)

Table 3: _Peak memory and run time comparison between PyTorch and our memory-
saving layers_ on GPU on different scenarios and CNNs.

Here we measure the effect of our memory-saving layers on real-world CNNs. We
consider four different scenarios:

All:

All network parameters are differentiable. This serves as reference to
establish similar performance of our layers in the absence of selective
differentiation.

Input:

Only the input to the neural net is differentiable. This situation resembles
constructing adversarial examples [9] or style-transfers [8] by optimizing
noisy inputs.

Conv:

Only convolution layers are differentiable. This situation is similar to
surgical fine-tuning [18], which splits a network into different blocks, each
containing a subset of layers, which are then trained one at a time.

Norm:

Only normalization layers are differentiable, resembling layer norm fine-
tuning in LLMs [36].

We report GPU results on an NVIDIA RTX A6000 with 48 GiB of VRAM, and CPU
results for an Apple M2 with 16 GiB of RAM. All networks are fed inputs of
size (64, 3, 224, 224) and we measure according to the procedure described in
Section 2. For object detection models, the batch size is 4 and 2 boxes are
predicted per input image.

###  3.1 ResNet-101

For ResNet-101, we take a detailed look at our memory-saving layer’s effects.
It is a decently powerful, modern model and frequently used as a backbone for
other architectures such as CLIP [25] and LAVA [11]. It contains dense,
convolution, batch normalization, ReLU, and max/average pooling layers. Table
2 summarizes our findings.

##### PyTorch is unaware of selective differentiation:

In the upper part of Table 2, we see that the default PyTorch implementation
uses the same amount of memory for any scenario. In fact, marking only the
input to a neural net as differentiable consumes as much memory as marking all
parameters, although the former does not require storing the inputs to linear
layers. This confirms the findings on the toy model from Section 2.

##### Layer interactions diminish memory savings:

To gradually investigate the effect of our layers, we only swap out the
convolutions; _without_ observing any effects. At first, this seems counter-
intuitive. However, a closer look at the architecture (Table 2) reveals that
the convolutions are preceded by ReLU activations which unconditionally store
their outputs and render our convolution layer’s optimizations ineffective as
described in Section 2.2.2.

After swapping out ReLU with our mask-based custom implementation, we observe
substantial memory savings. E.g., memory consumption for ‘Input’ dropped by a
factor of almost 5x on GPU. Only when all parameters are differentiable, we
see a slight increase in memory after swapping out ReLUs. This is to be
expected as the mask stored by our custom implementation requires additional
storage.

##### Run time remains unaffected:

The bottom half of Table 2 shows that run time for a fixed case remains equal
up to measurement noise. This confirms that our memory-saving layers share the
default implementation’s performance.

###  3.2 More results

To further solidify our findings, we now evaluate our layers on other popular
and commonly used CNNs, including ResNet-18 [13], VGG-16 [29],
ResNeXT101-64x4d [35], EfficientNetv2-L [30, 31], MobileNetv3-L [28, 14], FCN
with a ResNet-101 backbone [20], DeepLabv3 with a ResNet-101 backbone [4],
Faster-RCNN with a ResNet-50 backbone [27], and SSDLite with a MobileNetv3-L
backbone [28]. Table 3 summarizes the comparison.

On this large repertoire of networks, we observe the same effects as on
ResNet-101 from the previous section: Due to our customized ReLU
implementation, memory is slightly higher in the absence of selective
differentiation. Run time is unaffected by swapping in our layers, while the
selective differentiation scenarios consistently show lower memory consumption
with reduction factors up to 6x, underlining the usefulness of our approach in
this context.

##  4 Conclusion

We have shown how to improve the memory consumption of PyTorch’s automatic
differentiation in the context of selective differentiation where gradients
are only requested for a subset of variables—a common situation in modern
fine-tuning tasks. To overcome this, we provide a drop-in implementation which
takes into account the differentiability of all tensors when storing the
computation graph. Empirically, we demonstrated the effectiveness of our
approach to reduce memory in multiple selective differentiation cases without
affecting run time.

Our method is easy to use, requiring only a single call to a converter
function that replaces all supported layers with our equivalents. Next, we
plan to apply it to parameter-efficient LLM fine-tuning or layer norm tuning,
and other modern architectures like vision transformers. We expect this to
require additional efforts due to the interaction with other transformer
layers like softmax and GeLU.

## Acknowledgements

Resources used in preparing this research were provided, in part, by the
Province of Ontario, the Government of Canada through CIFAR, and companies
sponsoring Vector Institute.

## References

  * Ba et al. [2016] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.  Layer normalization.  2016\.
  * Chellapilla et al. [2006] Kumar Chellapilla, Sidd Puri, and Patrice Simard.  High performance convolutional neural networks for document processing.  In _International Workshop on Frontiers in Handwriting Recognition_ , 2006.
  * Chen et al. [2021] Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael W Mahoney, and Joseph E Gonzalez.  Actnn: Reducing training memory footprint via 2-bit activation compressed training.  In _International Conference on Machine Learning (ICLR)_ , 2021.
  * Chen et al. [2017] Liang-Chieh Chen, George Papandreou, Florian Schroff, and Hartwig Adam.  Rethinking atrous convolution for semantic image segmentation, 2017.
  * Chen et al. [2016] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin.  Training deep nets with sublinear memory cost, 2016.
  * Dangel [2023] Felix Julius Dangel.  Backpropagation beyond the gradient.  2023\.
  * Frantar and Alistarh [2022] Elias Frantar and Dan Alistarh.  Optimal brain compression: A framework for accurate post-training quantization and pruning.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Gatys et al. [2016] Leon A Gatys, Alexander S Ecker, and Matthias Bethge.  Image style transfer using convolutional neural networks.  In _IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2016.
  * Goodfellow et al. [2015] Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy.  Explaining and harnessing adversarial examples.  In _International Conference on Learning Representations (ICLR)_ , 2015.
  * Griewank and Walther [2008] Andreas Griewank and Andrea Walther.  _Evaluating derivatives: principles and techniques of algorithmic differentiation_.  SIAM, 2008.
  * Gurram et al. [2022] Sumanth Gurram, David Chan, Andy Fang, and John Canny.  LAVA: Language audio vision alignment for data-efficient video pre-training.  In _First Workshop on Pre-training: Perspectives, Pitfalls, and Paths Forward at ICML 2022_ , 2022.
  * Hassibi and Stork [1992] Babak Hassibi and David Stork.  Second order derivatives for network pruning: Optimal brain surgeon.  In _Advances in Neural Information Processing Systems (NIPS)_ , 1992.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2016.
  * Howard et al. [2019] A. Howard, M. Sandler, B. Chen, W. Wang, L. Chen, M. Tan, G. Chu, V. Vasudevan, Y. Zhu, R. Pang, H. Adam, and Q. Le.  Searching for mobilenetv3.  In _2019 IEEE/CVF International Conference on Computer Vision (ICCV)_ , pages 1314–1324, Los Alamitos, CA, USA, 2019. IEEE Computer Society.
  * Hu et al. [2022] Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.  LoRA: Low-rank adaptation of large language models.  In _International Conference on Learning Representations (ICLR)_ , 2022.
  * Hubara et al. [2018] Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran El-Yaniv, and Yoshua Bengio.  Quantized neural networks: Training neural networks with low precision weights and activations.  _Journal of Machine Learning Research (JMLR)_ , 2018.
  * Ioffe and Szegedy [2015] Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In _International Conference on Machine Learning (ICML)_ , 2015.
  * Lee et al. [2023] Yoonho Lee, Annie S Chen, Fahim Tajwar, Ananya Kumar, Huaxiu Yao, Percy Liang, and Chelsea Finn.  Surgical fine-tuning improves adaptation to distribution shifts.  In _The Eleventh International Conference on Learning Representations_ , 2023.
  * Li et al. [2017] Hao Li, Soham De, Zheng Xu, Christoph Studer, Hanan Samet, and Tom Goldstein.  Training quantized nets: A deeper understanding.  _Advances in Neural Information Processing Systems (NeurIPS)_ , 2017.
  * Long et al. [2015] J. Long, E. Shelhamer, and T. Darrell.  Fully convolutional networks for semantic segmentation.  In _2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 3431–3440, Los Alamitos, CA, USA, 2015. IEEE Computer Society.
  * Nagel et al. [2021] Markus Nagel, Marios Fournarakis, Rana Ali Amjad, Yelysei Bondarenko, Mart van Baalen, and Tijmen Blankevoort.  A white paper on neural network quantization, 2021.
  * Oktay et al. [2021] Deniz Oktay, Nick McGreivy, Joshua Aduol, Alex Beatson, and Ryan P Adams.  Randomized automatic differentiation.  In _International Conference on Learning Representations (ICLR)_ , 2021.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  PyTorch: An Imperative Style, High-Performance Deep Learning Library.  In _Advances in Neural Information Processing Systems 32_ , pages 8024–8035. Curran Associates, Inc., 2019.
  * [24] Fabian Pedregosa.  memory_profiler: Monitor memory usage of python code.
  * Radford et al. [2021] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever.  Learning transferable visual models from natural language supervision.  In _Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event_ , pages 8748–8763. PMLR, 2021.
  * Ren et al. [2021] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He.  Zero-offload: Democratizing billion-scale model training, 2021.
  * Ren et al. [2015] Shaoqing Ren, Kaiming He, Ross B. Girshick, and Jian Sun.  Faster R-CNN: towards real-time object detection with region proposal networks.  In _Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada_ , pages 91–99, 2015.
  * Sandler et al. [2018] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen.  Mobilenetv2: Inverted residuals and linear bottlenecks.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2018.
  * Simonyan and Zisserman [2015] Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  In _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_ , 2015.
  * Tan and Le [2019] Mingxing Tan and Quoc V. Le.  Efficientnet: Rethinking model scaling for convolutional neural networks.  In _Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA_ , pages 6105–6114. PMLR, 2019.
  * Tan and Le [2021] Mingxing Tan and Quoc V. Le.  Efficientnetv2: Smaller models and faster training.  In _Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event_ , pages 10096–10106. PMLR, 2021.
  * Thompson et al. [2020] Neil C. Thompson, Kristjan Greenewald, Keeheon Lee, and Gabriel F. Manso.  The computational limits of deep learning, 2020.
  * Ulyanov et al. [2016] Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky.  Instance normalization: The missing ingredient for fast stylization.  2016\.
  * Wu and He [2019] Yuxin Wu and Kaiming He.  Group normalization.  _International Journal of Computer Vision_ , 2019.
  * Xie et al. [2017] Saining Xie, Ross B. Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He.  Aggregated residual transformations for deep neural networks.  In _2017 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2017, Honolulu, HI, USA, July 21-26, 2017_ , pages 5987–5995. IEEE Computer Society, 2017.
  * Zhao et al. [2024] Bingchen Zhao, Haoqin Tu, Chen Wei, Jieru Mei, and Cihang Xie.  Tuning layernorm in attention: Towards efficient multi-modal LLM finetuning.  In _International Conference on Learning Representations (ICLR)_ , 2024.