PaperlayerView exact revision on arXiv

Source-linked AI summary

Flora: Low-Rank Adapters Are Secretly Gradient Compressors

Yongchang Hao, Yanshuai Cao, Lili Mou

arXiv:2402.03293v2cs.LGcs.AIstat.ML

TL;DR

Large neural networks require substantial memory for optimization states, while LoRA reduces those states but restricts weight updates to low rank. FLORA uses resampled random projections to support high-rank updates with sublinear-memory accumulation and momentum. Across evaluated tasks, it achieves performance similar to uncompressed full-matrix updates while outperforming compression techniques such as LoRA.

  • Problem

    Large neural networks require substantial memory for optimization states, and LoRA reduces that burden while restricting weight updates to low rank.

    Source: §1 · p. 1

  • Method

    FLORA compresses gradients with random projections that it resamples, storing compressed gradient accumulation and momentum to achieve sublinear memory use.

    Source: §1 · pp. 1–2

  • Results

    FLORA achieves performance similar to uncompressed full-matrix updates while largely outperforming compression techniques such as LoRA.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Resampling projection matrices addresses LoRA’s low-rank limitation while retaining sublinear memory for optimization states.

    Source: §1 · p. 2, §5 · p. 9

    The largest tested model is 7B; the estimated 2.08% optimization-state memory for GPT-3 with r = 256 remains empirically unverified.

    Source: §5 · p. 9

Abstract

from arXiv · show

Despite large neural networks demonstrating remarkable abilities to complete different tasks, they require excessive memory usage to store the optimization states for training. To alleviate this, the low-rank adaptation (LoRA) is proposed to reduce the optimization states by training fewer parameters. However, LoRA restricts overall weight update matrices to be low-rank, limiting the model performance. In this work, we investigate the dynamics of LoRA and identify that it can be approximated by a random projection. Based on this observation, we propose Flora, which is able to achieve high-rank updates by resampling the projection matrices while enjoying the sublinear space complexity of optimization states. We conduct experiments across different tasks and model architectures to verify the effectiveness of our approach.

1. Introduction

Modern optimization states can add substantial memory, while LoRA reduces trainable state by constraining updates to low rank. FLORA uses resampled random projections for high-rank updates with sublinear-memory accumulation and momentum, reporting performance similar to full-matrix updates and better than compression alternatives such as LoRA.

1. Introduction

  • Momentum and gradient accumulation each require an additional memory buffer equal to the model size.Source: §1 · p. 1
  • Adam stores first- and second-moment estimates, tripling memory per parameter; for 175-billion-parameter GPT-3, its optimization states require an additional 1.4TB.Source: §1 · p. 1
  • LoRA reduces trainable parameters and optimizer-state sizes through low-rank patches, but its low-rank restriction limits the optimization space of model parameters.Source: §1 · p. 1
  • Adafactor reduces second-moment memory for an n×m matrix from O(nm) to O(n + m), but drops momentum and does not reduce gradient-accumulation memory.Source: §1 · p. 1
  • FLORA applies resampled random projections to support high-rank updates while storing compressed gradient accumulation and momentum, using sublinear memory and performing similarly to full-matrix updates while outperforming compression techniques such as LoRA.Its space complexity is in the same order as LoRA but has a smaller constant in practice, leading to less memory usage than LoRA.Source: §1 · pp. 1–2

2. Approach

The approach interprets LoRA updates as gradient compression, then uses resampled projections to obtain high-rank updates while keeping optimization states sublinear in model size. FLORA applies this idea to gradient accumulation and momentum.

2.1. Dynamics of low-rank adaptation (LoRA)

  • With small learning rates, the change in B dominates LoRA's final weight update because the third term becomes negligible.The stated condition is η ≪ 1/L.Source: §2.1 · p. 2

2.2. Random projection of gradients

  • With approximations, LoRA down-projects gradients and then up-projects them, making its updates equivalent to random projection.This interpretation frames a LoRA update as gradient compression and decompression.Source: §2.2 · p. 3, §2.3 · p. 3

2.3. Our interpretation of LoRA

  • Theorem 2.4 states that rank r need only scale logarithmically to preserve element-wise reconstruction error; the projection also approximately preserves row norms.The stated logarithmic rank scaling is presented as efficient in computation and memory when m is large.Source: §2.3 · pp. 3–4
  • LoRA(B) performs close to LoRA and follows the same curve as random projection in the Fashion-MNIST pilot study.The comparison supports the interpretation that LoRA's update is dominated by compression and decompression.Source: Figure 1 · p. 3, §2.3 · p. 4

2.4. Our method: FLORA

  • FLORA resamples projection matrices to remove the low-rank constraint on total weight change while compressing optimization states at LoRA's asymptotic rate with a lower constant.The method targets high-rank updates and sublinear memory for gradient accumulation and momentum.Source: §2.4 · p. 4
  • In the pilot study, RRP largely recovers full-matrix SGD performance and significantly surpasses LoRA and its fixed-projection approximation.Figure 1 compares LoRA, LoRA(B), RP, RRP, and full-model SGD.Source: Figure 1 · p. 3, §2.4 · p. 4
  • For gradient accumulation, FLORA stores accumulated gradients in down-projected space during each cycle, then decompresses and reuses gradient memory.The cycle's projection is resampled when accumulation ends.Source: §2.4 · p. 4
  • For momentum, FLORA addresses historical-state distortion by holding projections longer or transferring compressed momentum between projections, reducing storage from O(nm) to O(nr).The approach can extend to other EMA-based statistics.Source: §2.4 · p. 5

3. Experiments

Experiments compare FLORA with naive accumulation, LoRA, and baselines across Transformer models and language tasks. FLORA reduces memory while recovering or surpassing full-state performance at larger projection ranks, with analyses testing resampling frequency and linear-memory optimizers.

3.1. Experiment setup

  • Experiments use T5 and GPT-2 models on XSum summarization and IWSLT-2017 German-English translation, comparing None, naive accumulation or momentum, LoRA, and FLORA.For LoRA and FLORA, projections are applied to attention and feed-forward layers; other layers use the naive procedure.Source: §3.1 · pp. 5–6

3.2. Main results

  • For T5-3B and GPT-2-XL, FLORA at r = 256 uses 30% of naive accumulation’s memory overhead with performance on par, and reduces memory on all benchmarks.At reasonably large ranks, it recovers performance of full-matrix accumulation and surpasses the no-accumulation baseline.Source: Table 1 · p. 6, §3.2 · p. 7
  • FLORA uses less memory than naive momentum and recovers or surpasses its performance as r increases, including when training starts from random initialization.LoRA does not save memory with Adafactor’s already-sublinear state and has significantly lower performance from scratch under its low-rank update constraint.Source: Table 2 · p. 7, §3.2 · p. 7

3.3. In-depth analyses

  • With T5-small on summarization, ROUGE scores rise as κ increases below 1000, then decrease beyond a threshold; κ = 1000 balances preserved information and momentum rank.κ controls how frequently FLORA resamples the random-projection matrix.Source: §3.3 · pp. 7–8
  • With a linear-memory optimizer, FLORA is more memory-efficient than LoRA at r = 256 and outperforms it by 2–3 ROUGE points in all tested settings.LoRA saves more memory at ranks below 128.Source: Table 4 · p. 8, §3.3 · p. 8

3.4. Additional experiments

  • Preliminary experiments report FLORA performing well with ViT on images, outperforming GaLore in memory reduction and model quality, and combining with activation checkpointing and layer-wise updating to reduce peak memory.Source: §3.4 · p. 8

4. Related work and discussion

The discussion contrasts FLORA with parameter-efficient fine-tuning and matrix-compression methods, and describes how it can complement memory-efficient optimization and back-propagation techniques.

Parameter-efficient fine-tuning

  • Unlike LoRA’s low-rank updates, FLORA can start from scratch and achieve full training performance while maintaining sublinear complexity throughout.ReLoRA periodically reinitializes its LoRA patch to pursue high-rank updates but requires full-weight pre-training, increasing peak memory linearly with model size.Source: §4 · p. 8

Matrix compression

  • FLORA uses random projection rather than SVD-based compression, avoiding the latter’s computational expense and difficulty with parallelization; its decompression procedure is justified by Theorem 2.4.Source: §4 · p. 8

Memory-efficient optimizers

  • FLORA is orthogonal to memory-efficient optimizers and can compress their momentum or gradient accumulation; it can also combine with back-propagation memory techniques, though further exploration is left to future work.Source: §4 · p. 8

5. Conclusion

FLORA uses resampled random projections to address LoRA’s low-rank limitation while keeping optimization-state memory sublinear. Experiments show significant memory reduction with maintained model performance; empirical testing extends only to 7B models.

5. Conclusion

  • Experiments demonstrate significant memory reduction with maintained model performance.The conclusion frames these results as evidence of random projection’s potential in deep learning.Source: §5 · p. 9
  • Resampling projection matrices addresses LoRA’s low-rank limitation while supporting sublinear memory for gradient accumulation and momentum.FLORA is based on random projections.Source: §5 · p. 9
  • The largest empirically tested model is 7B; the estimated 2.08% optimization-state memory for GPT-3 at r = 256 remains unverified.The authors leave empirical verification of the GPT-3 estimate to future work.Source: §5 · p. 9

Impact statement

The authors anticipate that reducing model-training resources could benefit the environment, while expecting no direct consequences for the general public.

Impact statement

  • The authors anticipate environmental benefits from reduced model-training resources, but no direct consequences for the general public.The impact statement presents the environmental effect as an anticipated positive impact.Source: Impact statement · p. 9

A. Proof of Theorem 2.1

The proof derives iterative forms for the LoRA update factors and uses induction to establish a bound on the A-factor norm. It also states that the model remains within a finite Euclidean ball during training.

A. Proof of Theorem 2.1

  • The proof establishes that the model stays within a finite Euclidean ball during training.The passage states this follows from an upper bound by L for every T.Source: A. Proof of Theorem 2.1 · p. 13
  • The proof gives the forms and bounds for fA(t) and fB(t) that describe the dynamics of A_t and B_t.The stated bound includes ∥fA(t)∥₂ ≤ 1−η²L² for every t.Source: A. Proof of Theorem 2.1 · p. 13
  • The update factors fA(t) and fB(t) start at zero and are obtained iteratively for t ≥ 1.The lemma states fA(0) = fB(0) = 0.Source: A. Proof of Theorem 2.1 · p. 13
  • An induction argument establishes ∥fA(t)∥ ≤ a_t for every t.The proof uses the Frobenius-norm bound and concludes by taking square roots.Source: A. Proof of Theorem 2.1 · pp. 13–14

B. Proof of Theorem 2.4

The theorem gives a rank condition for a Gaussian matrix and states a probability guarantee for an all-index bound. Its proof uses chi-square bounds and a union bound.

B. Proof of Theorem 2.4

  • For Gaussian A, the theorem sets r = c log(2m/δ)ϵ⁻² as a sufficient rank condition for its all-i,j guarantee.The theorem specifies ϵ, δ ∈ (0, 1] and a constant c.Source: B. Proof of Theorem 2.4 · p. 15
  • The all-i,j guarantee is stated to hold with probability at least 1 − δ.The theorem also states a further guarantee with confidence at least 1 − δ.Source: B. Proof of Theorem 2.4 · pp. 15–16
  • The proof uses chi-square bounds for diagonal terms and a union bound over the derived inequalities.It identifies the diagonal statistic as chi-square distributed and invokes Laurent–Massart bounds.Source: B. Proof of Theorem 2.4 · pp. 15–16

C. Additional results

On CIFAR-100 image classification with ViT, FLORA saves training memory without sacrificing accuracy, supporting its use beyond text-generation settings.

C.1. Beyond the text modality

  • The evaluation tests FLORA on image classification with a Vision Transformer, extending beyond the paper’s main text-generation experiments.Source: §C.1 · p. 16
  • The results show that FLORA saves as much as 32.4% of training memory without sacrificing accuracy, consistent with the main experiments and supporting generality across model architectures and domains.Source: §C.1 · p. 16
  • The authors conclude that these results demonstrate generality across model architectures and domains.Source: §C.1 · p. 16

C.2. Comparison with GaLore

On C4 language modeling, FLORA uses memory similar to GaLore while achieving better perplexity. With activation checkpointing and LOMO enabled for a 7B model, peak allocated memory is 21.2 GiB.

C.2. Comparison with GaLore

  • FLORA generates projection matrices on the fly, whereas GaLore performs SVD operations and stores its projection matrices on-device.Source: §C.2 · pp. 16–17
  • For the 7B model, Table 6 reports memory usage only, using batch size 16 because training takes months to complete.Source: Table 6 · p. 17
  • FLORA achieves better perplexity than GaLore, with similar memory usage and a small overhead for GaLore.Table 6 reports C4 language-modeling results; lower token-level perplexity is better.Source: §C.2 · p. 17, Table 6 · p. 17
  • With activation checkpointing and LOMO enabled, FLORA uses 21.2 GiB peak memory for the 7B model, surpassing GaLore’s memory saving in the same setting.Source: §C.2 · p. 17

C.3. Profiling results

Memory profiling compares Adam, LoRA, and FLORA over four T5 training steps, with and without activation checkpointing and LOMO. Both adapter methods have negligible optimizer-state memory, and their peak usage is similar.

C.3. Profiling results

  • The profiling uses four training steps of T5 with sequence length 512 and batch size 4, and also evaluates activation checkpointing and LOMO.Source: §C.3 · p. 17
  • LoRA and FLORA have negligible optimization-state memory; despite FLORA’s larger gradient storage, activation-dominated peak memory is similar.Source: §C.3 · p. 17, Figure 2 · p. 18
  • With activation checkpointing and LOMO, the gradient-storage difference becomes less significant and LoRA and FLORA show similar memory-usage patterns.Source: §C.3 · p. 17, Figure 2 · p. 18
Loading 2402.03293v2…