Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

VeLoRA: Memory Efficient Training using Rank-1 Sub-Token Projections

Roy Miles, Pradyumna Reddy, Ismail Elezi, Jiankang Deng

TL;DR

This paper identifies and characterises the important components needed for effective model convergence using gradient descent and finds that the intermediate activations used to implement backpropagation can be excessively compressed without incurring any degradation in performance.

Abstract

Large language models (LLMs) have recently emerged as powerful tools for tackling many language-processing tasks. Despite their success, training and fine-tuning these models is still far too computationally and memory intensive. In this paper, we identify and characterise the important components needed for effective model convergence using gradient descent. In doing so we find that the intermediate activations used to implement backpropagation can be excessively compressed without incurring any degradation in performance. This result leads us to a cheap and memory-efficient algorithm for both fine-tuning and pre-training LLMs. The proposed algorithm simply divides the tokens up into smaller sub-tokens before projecting them onto a fixed 1-dimensional subspace during the forward pass. These features are then coarsely reconstructed during the backward pass to implement the update rules. We confirm the effectiveness of our algorithm as being complimentary to many state-of-the-art PEFT methods on the VTAB-1k fine-tuning benchmark. Furthermore, we outperform QLoRA for fine-tuning LLaMA and show competitive performance against other memory-efficient pre-training methods on the large-scale C4 dataset.

VeLoRA: Memory Efficient Training using Rank-1 Sub-Token Projections

TL;DR

This paper identifies and characterises the important components needed for effective model convergence using gradient descent and finds that the intermediate activations used to implement backpropagation can be excessively compressed without incurring any degradation in performance.

Abstract

Large language models (LLMs) have recently emerged as powerful tools for tackling many language-processing tasks. Despite their success, training and fine-tuning these models is still far too computationally and memory intensive. In this paper, we identify and characterise the important components needed for effective model convergence using gradient descent. In doing so we find that the intermediate activations used to implement backpropagation can be excessively compressed without incurring any degradation in performance. This result leads us to a cheap and memory-efficient algorithm for both fine-tuning and pre-training LLMs. The proposed algorithm simply divides the tokens up into smaller sub-tokens before projecting them onto a fixed 1-dimensional subspace during the forward pass. These features are then coarsely reconstructed during the backward pass to implement the update rules. We confirm the effectiveness of our algorithm as being complimentary to many state-of-the-art PEFT methods on the VTAB-1k fine-tuning benchmark. Furthermore, we outperform QLoRA for fine-tuning LLaMA and show competitive performance against other memory-efficient pre-training methods on the large-scale C4 dataset.
Paper Structure (29 sections, 15 equations, 2 figures, 10 tables)

This paper contains 29 sections, 15 equations, 2 figures, 10 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Memory-Efficient Training.
  4. Parameter-Efficient Fine-Tuning (PEFT)
  5. Subspace training
  6. Gradient Sparsification.
  7. Method
  8. VeLoRA
  9. Insights and properties of VeLoRA
  10. Comparison with the state-of-the-art
  11. Implementation details
  12. Vision experiments
  13. Roberta experiments
  14. Scaling up to LLaMA
  15. Pre-training on C4
  16. ...and 14 more sections

Figures (2)

  • 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.
  • Figure 2: (a) Stable rank for the input activations using a different number of groups, with $= 1$ indicating no sub-division of the tokens into smaller sub-tokens. (b) Approximate probability of the feature similarity diverging by at least $k$. (c) visualisation the rank-1 projection of sub-tokens.