RiteeshGrandhi/sciclaimeval-shared-task huggingface dataset:
Memory-Efficient Training.
The increase in model size has necessitated the development of smart methods that make training more memory efficient.
Gradient checkpointing [ 8 ] significantly lowers the memory requirements during model training by recomputing activations for the backward pass instead of storing them during the forward pass. However, doing so increases the training time from the need to re-compute gradients.
Adafactor [ 40 ] and its followup [ 1 ] lowers the memory by working with the row-column outer-product factorized moments of adaptive optimizers.
LOMO [ 29 ] was developed for large models and works by fusing the gradient computation and the parameter update in one step to reduce memory usage, effectively only saving the current layer gradients in memory.
Recently, GaLore [ 53 ] proposed projecting the gradients onto a lower-dimensional space [ 9 , 6 ] , and can reduce the memory during both pre-training and finetuning. However, they store all the full intermediate activations to compute gradient updates. Their memory advantage is derived from computing the first and second-order statistics of the gradients in a lower-dimensional space, thus limiting it to second-order optimizers only. Furthermore, GaLore needs an expensive SVD operation that introduces some significant overheads in terms of both memory and computation costs.
Unlike these methods, VeLoRA does not introduce any large computation overhead while at the same time comes with a significant memory reduction. Furthermore, VeLoRA is in principle independent of the underlying optimizer.
Parameter-Efficient Fine-Tuning (PEFT)
is an emerging field that focuses on fine-tuning a large model with a minimal number of trainable parameters.
This typically involves freezing and then augmenting the original model with adapter modules.
LoRA (Low-Rank Adaptation) [ 17 ] is a technique that optimizes a few rank-decomposed weight matrices during fine-tuning, rather than updating the entire set of pre-trained weights for each attention layer.
This approach substantially reduces the number of trainable parameters, thereby accelerating the fine-tuning process and making it more memory-efficient.
The method was later extended to also work with multi-layer perceptrons in Transformers [ 20 , 13 ] .
Several other methods built upon these works improving the capacity or performance of the model [ 20 , 26 , 7 , 41 , 34 , 51 , 48 , 25 , 44 ] . These works can be well-complemented with quantization, further reducing the memory while keeping the performance [ 10 , 11 , 23 ] .
Our memory-efficient algorithm is complementary to PEFT and can be used to provide additional memory efficiency in the fine-tuning regime.
Subspace training
In [ 21 , 15 ] , the authors show that most of the learning process occurs within a significantly low-rank parameter space and that model weights can be effectively optimized within this lower-dimensional subspace. These subspace learning techniques have been widely adopted in various machine learning domains, including meta-learning [ 22 ] and continual learning [ 5 ] . However, unlike VeLoRA, resource-efficient training/fine-tuning is not the focus of these methods, therefore, often resulting in an overhead to meet other requirements.
Gradient Sparsification.
Recently, there has been a surge in interest for memory-efficient training methods.
In [ 49 ] only a sparse subset of the gradient vector components are stored zeroing out the remaining components. Different criteria have been proposed for selecting which gradient components to retain, such as Top-K SGD [ 42 ] which keeps only the top-k largest components, Sparsified-SGD [ 43 ] and various other sparsification methods [ 36 , 37 , 39 , 24 , 27 , 16 ] .
More recently, techniques combining quantization and sparsification have been proposed for resource-efficient training. Examples include TernGrad [ 50 ] , Qsparse-local-SGD [ 3 ] , sparse ternary compression [ 38 ] , and the sparse-signSGD [ 31 ] method which combine sparsity with quantizing gradients to just the sign. A key difference is how VeLoRA compresses the intermediate activations that are used to compute gradients. Our compression algorithm is fully dense-to-dense without any pruning or sparsification of the activations. This prevents accuracy degradation issues associated with sparse updates and facilitates memory-efficient training.
Method
The task. In this work, we propose a memory-efficient modification to the back-propagation algorithm. Our
primary motivation is to reduce the GPU memory footprint during training without resorting to any hardware-specific quantization schemes [ 11 ] and without trading compute for memory as is done with gradient checkpointing [ 8 ] .

To formalize the problem statement, let us take a step back and look at the components needed to implement back-propagation.
Firstly, during the forward pass, each trainable layer in the neural network needs to store two key tensors in memory - the input activations received by that layer, and the layer’s trainable weight parameters. Retaining these tensors is required for performing the parameter update rules and computing the input gradients.
More specifically, during the backward pass, the previously stored input activations and model weights are used to calculate the gradients with respect to the weights and the input via the chain rule of differentiation (see Fig. 1 ).
Storing both these sets of tensors comes with a significant memory overhead which scales with the model size.
We focus on optimizing the cost of storing the input activations. We do this by compressing intermediate activation vectors and then reconstructing them when the original activations are needed for gradient calculations.
This is orthogonal to PEFT [ 17 ] methods which address the overhead of saving the full-precision weights in memory by introducing cheaper trainable adapters.
Further, in Section 4 , we show how to combine our method with PEFT methods to achieve state-of-the-art results.
Comparison with the state-of-the-art
In this section, we thoroughly evaluate VeLoRA and its individual components. In section 4.2 we demonstrate the complementary nature of VeLoRA in conjunction with many other existing PEFT methods. Section 4.4 then scales up these results to the LLaMA models, whereby we achieve a significant memory reduction when coupled with 4-bit scalar quantisation. Finally, section 4.5 extends VeLoRA to the pre-training setting where we see competitive performance alongside a real reduction for the on-device GPU memory usage.

For the VTAB-1k experiments, we applied VeLoRA to all the down projections in the trainable adapters, while for SSF we applied it to the input scale and shift layers only. For all the other experiments we simply applied VeLoRA to the value projection layer and the down projection layer of the MLP blocks.
Implementation details
We performed all the vision experiments on a subset of the VTAB-1k [ 52 ] benchmark for a combination of 16 different vision datasets. We finetuned a ViT-B [ 12 ] model pre-trained on ImageNet-21K using
the AdamW optimizer with a learning rate of 5e-4 and a weight decay of 1e-4.
We performed the language models experiments on the GLUE [ 47 ] benchmark using the encoder-decoder RoBERTa transformer [ 28 ] .
We then scaled our model to large language models causal transformers using the LLama [ 2 ] family of models, finetuned in Alpaca dataset [ 35 ] and reported results on the MMLU benchmark [ 35 ] .
We finally applied our method to the pre-training stage using some smaller LLama [ 2 ] models on the C4 dataset [ 33 ] .
All our models were trained using the AdamW optimizer with a learning rate of 1e-3 and a weight decay of 0.
We give all the other major hyper-parameters to replicate these experiments in the Appendix and also in the code.

Memory (MB)

Caltech101

Cifar100

DTD

Flower102

Pets

SVHN

Sun397

Camelyon

EuroSAT

Resisc45

Retinopathy

Clevr-Count

DMLab

KITTI-Dist

sNORB-Azim

sNORB-Ele

Average
Vision experiments
We conduct experiments evaluating the performance of VeLoRA for full-tuning and how it complements other PEFT methods.
In Tab. 1 we reproduce a large set of results for LoRA [ 17 ] , SSF [ 26 ] , and Hydra [ 20 ] on a subset of the VTAB-1K benchmark, where the sub-token size for each experiment is given in the Appendix.
Unlike what is more common in the PEFT literature [ 19 , 20 ] , we do not perform any task-specific hyperparameter tuning that would change the memory, such as batch size and rank, and to also avoid any potential overfitting to the specific task. For all experiments we used the authors provided implementations for the adapters and integrated them into the same training framework for a fair comparison.
We observe that VeLoRA improves the performance compared to full-tuning by 1.5 percentage points ( pp ), while lowering the memory requirements.
We also observe that when combined with PEFT methods, VeLoRA can come with improvement in memory and performance.
VeLoRA lowers the memory requirement of SSF [ 26 ] by 16% with only a minor degradation (0.1 pp ) in accuracy.
It lowers the memory requirement of Hydra [ 20 ] by 7% while improving the accuracy by 0.1 pp .
Finally, it lowers the memory requirements of LoRA [ 17 ] by 4% while improving the accuracy by 0.6 pp
Roberta experiments
We now evaluate our method with M=16 on various language tasks, using RoBERTa-Base [ 28 ] in the GLUE benchmark, and compare it with full fine-tuning, LoRA [ 17 ] and GaLore [ 53 ] , presenting the results in Tab. 2 .
We observe that both GaLore and LoRA lower the memory requirements compared to fine-tuning from 4.64GB to 4.04GB, respectively to 2.71GB, at a cost of accuracy degradation with GaLore performance lowered by 0.39 pp , while LoRA accuracy drops by 0.67 pp .
Our method further reduces the memory needed for training to 2.23GB, an improvement of 18% compared to LoRA, and 45% compared to GaLore, while still reaching higher results than either of them.
More impressively, VeLoRA reduces the memory by half compared to full fine-tuning with an accuracy degradation of only 0.37 pp , reaching the best tradeoff between memory and accuracy.
Scaling up to LLaMA
We now scale our method to large language models, demonstrating the effectiveness of VeLoRA in finetuning them.
We do comparisons with LoRA on both BFloat16 and Float4, in addition to the recent method of QLoRA [ 11 ] which is widely used for fine-tuning LLMs with very low memory budget.
We aim to further lower this budget, showing in the process that VeLoRA is also complementary to QLoRA, resulting in a much lower memory consumption.
We present the results in Tab. 4.4 using M=32 for 7B and M=128 for 13B.
We can see that our method outperforms QLoRA by 0.5 pp in the Llama model, while reaching a massive performance increase compared to LoRA models.
Furthermore, we reach this performance improvement, while at the same time further reducing the memory.
In particular, we reduce the memory for 0.89GB, a relative improvement of 15.4% from QLoRA.
We observe that this performance improvement is maintained on the larger model of 13B parameters, where again our method outperforms QLoRA by 0.5 pp and lowers the memory requirements by 1.43GB, a relative improvement of 14.4%.
Pre-training on C4
We now perform an experiment where we use VeLoRA to train language models from scratch in the large C4-English dataset, presenting the results in Tab. 4 .
We use M=128 for both the 60M and 130M, while following the same training pipeline and evaluation as GaLore and comparison with LoRA.
However, unlike in the GaLore paper [ 53 ] , which estimates the memory usage using the optimizer weights and memory alone, we choose to compute the real on-device memory. This quantity would take into account the cost of additionally storing the intermediate activations and also highlight the benefits of LoRA in terms of memory since the base weights will be frozen.
In contrast to other experiments, our use of VeLoRA here is not with any additional adapters and is simply compressing the input activations for the original trainable base layers.
We observe that our method significantly outperforms the other methods, reaching 1.08 pp lower perplexity than GaLore.
We observe that our method outperforms GaLore in Llama-130M too.
Convergence properties
We observed that the rank-1 projections would encourage much higher levels of gradient sparsity (see Fig. 2 b).
A natural question to ask from this observation is if the gradient sparsity will come at the cost of the model’s convergence.
In other words, we want to verify if our model needs to be trained longer to compensate for the gradient compression.
To do so, we evaluate the performance of our model at the end of each epoch and compare it with the performance of a competing model, the QLoRA.
As shown in Tab. 5 a, VeLoRA and QLoRA improve at the same rate.
Our model outperforms QLoRA by 0.3 pp by the end of the first epoch, and keeps this improvement rate, outperforming QLoRA by 0.4 pp at the end of the training.
In this way, we verify that the additional compression of input activations does not affect the model’s convergence.
Sub-token size
We provide an ablation on the impact on the size of each sub-token and the model’s performance, showing the results in Tab. 5 b. We can see that there is a sweet spot for which the rank-1 projections of sub-tokens using an average initialisation is very effective both in terms of memory and accuracy. However, if the size of the sub-tokens is too large (i.e. when M=D/8 ), the gradients will become too sparse, which will hinder performance (see also figure 2 b). In contrast, if the size of each sub-token is too small, for example with D/64 , there is a more significant memory compression but at the cost of model performance.
Initialisation strategy
In Tab. 5 c, we show the performance after training with various ways of initialising the vectors for each group. To avoid exceeding memory requirements we sub-sampled the tokens for SVD and we also consider the case of instance-wise initialisation. Although we would have expected better performance since the vector will always align with each incoming batch, we found that it did not lead to any performance improvement. In contrast, doing SVD initialization comes with a drop in performance. This result further confirms that the performance improvement from VeLoRA is not specifically correlated with a lower reconstruction error of the input activations.
Choice of layers
A key design decision for many Parameter-Efficient Fine-Tuning (PEFT) methods, including LoRA, is where to insert the adapters within the model layers.
Since VeLoRA shares this core architectural choice, we aim to provide stronger intuition on VeLoRA’s suitability by performing a thorough ablation study analyzing the trade-offs between memory consumption and accuracy when considering different layer placements.
We present the results in Tab. 6 .
We observe that we achieve memory improvement in all cases where we adopt VeLoRA.
However, to improve the accuracy, VeLoRA must be used in the MLP down-projection.
A possible explanation is that this layer might suffer from overfitting on the training data or forgetting [ 4 ] from the pre-trained data.
Overall, we conclude that applying VeLoRA to the value and down projection appears to be the best choice.
We strengthen this claim by using this setting for all other experiments.
Conclusion
In this work, we proposed VeLoRA, a novel framework that enables the training of networks, including large language models in a highly memory-efficient manner.
Our approach compresses intermediate activations during the forward pass and coarsely reconstructs them during backpropagation.
VeLoRA complements PEFT methods and is able to significantly reduce memory requirements while improving
the performance.
VeLoRA is effective when tested in both moderately-sized vision transformers as well as in large language models.
We performed experiments to demonstrate the method’s effectiveness on VTAB-1K, MMLU, GLUE, and C4 benchmarks outperforming state-of-the-art methods such as LoRA, QLoRA or GaLore.
Limitations and Broader Impact
Limitations. We performed all experiments on Transformer models. Although Transformers have become dominant in machine learning and computer vision, there are other important deep learning networks such as CNNs, RNNs and SSMs. It remains unclear whether our methods can be extended to non-Transformer-based models and how such an extension could be accomplished.
Moreover, although our method is computationally more efficient than competing methods, its primary advantage lies in the substantial reduction of GPU memory.
However, the issue of training time still persists.

Broader Impact. As compute power grows exponentially, model sizes grow even faster, making it challenging for smaller institutions, especially academic ones, to conduct high-quality research. This work aims to democratize AI research, particularly in large language models, by significantly reducing the memory needed for training, enabling researchers with limited compute resources to train networks and contribute to their research. However, the democratization of AI is controversial, with leading institutions like OpenAI, Anthropic, and Google DeepMind becoming more closed due to the potential risks of LLMs in the wrong hands. We acknowledge this concern and do not endorse misuse of our research.
Memory Estimates
In this section, we provide a detailed explanation of the memory estimates presented in the paper. For the ablation experiments in Tables 5 b and 6 , we only estimate the memory consumption allocated for storing the intermediate activations. We do this because all the other memory allocations, such as the weights and optimizer states, would be the same for each experiment. The reported memory estimates will highlight the impact of each design choice or parameter on the memory of the intermediate activations only, which is one of the focuses of this paper.

Additionally, for the QLoRA Llama experiments in Tab. 3 , we also use the estimated memory, but we compute this from three components: (i) Saving the frozen base weights in memory, these weights are what are quantized; (ii) The trainable adapter weights, which are stored in fp16 format for all methods; and (iii) The memory cost of storing the input features for each trainable layer. We acknowledge that an efficient implementation of our method with 4-bit quantization is left for future work. However, for all other experiments using fp16, we do report the real on-device GPU memory usage.
Roberta Experiments
We fine-tune the pre-trained RoBERTa-Base model on the GLUE benchmark. For the training parameters, we follow the same settings as GaLore [ 53 ] , except that we maintain the same hyperparameters for all experiments. These hyperparameters are given in table 7 .
VTAB-1K Experiments
We reproduce all the reported PEFT methods under a fixed rank and batch-size setting across all tasks. This is to ensure that the memory usage is fixed for all the tasks. However, for Hydra [ 20 ] we do use the optimal scale and dropout parameters provided the authors (see table 8 ). The rank v.s. memory for each reported training run is provided in table 9 . Here we were able to use a higher rank for LoRA to get better performance, while also reducing the total memory. We found that Hydra did not noticeably improve performance when increasing the rank and so for these we maintained the same rank for with and without the addition of VeLoRA.

Caltech101

Cifar100

DTD

Flower102

Pets

SVHN

Sun397

Camelyon

EuroSAT

Resisc45

Retinopathy

Clevr-Count

DMLab

KITTI-Dist

sNORB-Azim

sNORB-Ele
Implementation Details
All of the experiments in sections 4.2 and 4.5 were performed using 8 NVIDIA V100 GPUs with the fp16 data type. For the LLaMA experiments in section 4.4 , we trained on 4 NVIDIA A100 GPUs using the QLoRA nf4 data type. The LLaMA experiments were based on the alpaca-lora repository 1 1 1 https://github.com/tloen/alpaca-lora , while the RoBERTa and C4 experiments was based on the GaLore repository 2 2 2 https://github.com/jiaweizzhao/GaLore .
Gradient Similarity
Using sim(\mathbf{z}_{i},\mathbf{z}_{j})=|\mathbf{z}_{i}|~{}|\mathbf{z}_{j}|~{}cos~{}\theta_{ij} and proj_{\mathbf{v}}(\mathbf{z}_{i})=|\mathbf{z}_{i}|~{}cos~{}\theta_{iv} we wish to see how much the gradient similarity is being preserved under various assumptions about the distributions of both \mathbf{z}_{i} and \mathbf{z}_{j} . We can introduce a measure for this divergence as follows:

 d(\mathbf{z}_{i},\mathbf{z}_{j};\mathbf{v}) =\left|sim(proj_{\mathbf{v}}(\mathbf{z}_{i}),~{}proj_{\mathbf{v}}(\mathbf{z}_{j}))-sim(\mathbf{z}_{i},\mathbf{z}_{j})\right| (4) =|sim\left((\mathbf{v}\cdot\mathbf{z}_{i})\mathbf{v},~{}(\mathbf{v}\cdot\mathbf{z}_{j})\mathbf{v}\right)-|\mathbf{z}_{i}|~{}|\mathbf{z}_{j}|~{}cos~{}\theta_{ij}| (5)

Since \theta_{vv}=0 and |\mathbf{v}|=1 , we can simplify this expression as follows:

=\left||\mathbf{z}_{i}||cos~{}\theta_{iv}||\mathbf{z}_{j}||cos~{}\theta_{jv}|-|\mathbf{z}_{i}|~{}|\mathbf{z}_{j}|~{}cos~{}\theta_{ij}\right| (6) =\left||\mathbf{z}_{i}||\mathbf{z}_{j}|(|cos~{}\theta_{iv}~{}cos~{}\theta_{jv}|-cos~{}\theta_{ij})\right| (7)

Without loss in generality, we can assume |\theta_{jv}|>|\theta_{iv}| and use \theta_{ij}=\theta_{jv}-\theta_{iv} .

=\left||\mathbf{z}_{i}||\mathbf{z}_{j}|(|cos~{}\theta_{iv}~{}cos~{}\theta_{jv}|-cos~{}(\theta_{jv}-\theta_{iv}))\right| (8)

For simplicity, consider the case where both the input activations \mathbf{z}_{i} and \mathbf{z}_{j} are also normalised to unit length. In practise, these magnitudes will simply be scaling the divergence linearly.

Let both \theta_{i} and \theta_{j} be normally distributed with standard deviation \sigma and have the vector v be appropriately initialised such that their means are 0 .
To understand how our projection degrades the gradient similarity, we will consider the probability of d(\cdot) exceeding some scalar k , i.e:

 Pr(\left|sim(proj_{v}(\mathbf{z}_{i}),~{}proj_{v}(\mathbf{z}_{j}))-sim(\mathbf{z}_{i},\mathbf{z}_{j})\right|>k) (9)
Small \sigma approximation
For normally distributed \theta_{i} and \theta_{j} , cos(\theta_{i}) and cos(\theta_{j}) will not be uniformly distributed and instead follow a distribution that is more concentrated around their mean values. Thus, we can use a first-order approximation of cos(\cdot) :

 cos(\theta_{i})\approx 1,\;\;\;\;cos(\theta_{j})\approx 1,\;\;\;\;cos(\theta_{i}-\theta_{j})\approx 1-\frac{(\theta_{i}-\theta_{j})^{2}}{2} (10)

Therefore, we can express d(\cdot) and Pr(d(\cdot)>k) as follows:

 d(\mathbf{z}_{i},\mathbf{z}_{j};\mathbf{v}) \approx\frac{(\theta_{i}-\theta_{j})^{2}}{2} (11) \rightarrow Pr(d(\mathbf{z}_{i},\mathbf{z}_{j};\mathbf{v})>k) \approx Pr(\frac{(\theta_{i}-\theta_{j})^{2}}{2}>k) (12) =Pr((\theta_{i}-\theta_{j})^{2}>2k) (13) =Pr(|\theta_{i}-\theta_{j}|>\sqrt{2k}) (14)

Since \theta_{j}-\theta_{i} is normally distributed with variance 2\sigma^{2} we have:

=2Pr(\theta_{j}-\theta_{i}>\sqrt{2k}) (15)

Finally, using the cumulative distribution function (CDF) of the normal distribution:

 Pr(\theta_{j}-\theta_{i}>\sqrt{2k})-1-\Phi\left(\frac{\sqrt{2k}}{\sqrt{2}\sigma}\right) =1-\Phi\left(\frac{\sqrt{k}}{\sigma}\right) (16) \rightarrow Pr(d(\mathbf{z}_{i},\mathbf{z}_{j};\mathbf{v})>k) \approx 2\left(1-\Phi\left(\frac{\sqrt{k}}{\sigma}\right)\right)\;\;\;\;\;\blacksquare (17)
Figure 1 : The memory overhead for backpropagation on a single layer consists of storing the intermediate activations and the weights. (a) demonstrates that PEFT methods can reduce the memory by using cheap low-rank adapters. (b) VeLoRA additionally compresses the saved intermediate activations to further reduce the memory usage.; The memory overhead for backpropagation on a single layer consists of storing the intermediate activations and the weights. (a) demonstrates that PEFT methods can reduce the memory by using cheap low-rank adapters. (b) VeLoRA additionally compresses the saved intermediate activations to further reduce the memory usage.