Source-linked AI summary
ReLoRA: High-Rank Training Through Low-Rank Updates
Vladislav Lialin, Namrata Shivagunde, Sherin Muckatira, Anna Rumshisky
TL;DR
Training increasingly overparameterized models is costly, while the need for their scale remains poorly understood. The paper introduces ReLoRA, which accumulates low-rank updates to train high-rank networks, and finds comparable performance to regular training with reported efficiency gains.
Problem
The necessity of training models with far more parameters than examples remains poorly understood, while scaling makes training costs prohibitive and rarely democratizes access.
Method
ReLoRA adapts parameter-efficient low-rank training by repeatedly merging updates, resetting optimizer state, re-warming the learning rate, and using a full-rank warm start.
Results
ReLoRA achieves performance similar to regular training, reduces wall-clock time by 9% on 8xA100s and by 20–40% on cheaper hardware, and consistently outperforms LoRA.
Takeaways & Limitations
Parameter-efficient training can be adapted for pre-training large language models and may improve the efficiency of multi-billion-parameter training.
Takeaways & Limitations
GLUE results are expected to be far from state of the art because the models were pre-trained on roughly 20 times less data than T5 or BERT.
Abstract
from arXiv · showhide
Despite the dominance and effectiveness of scaling, resulting in large networks with hundreds of billions of parameters, the necessity to train overparameterized models remains poorly understood, while training costs grow exponentially. In this paper, we explore parameter-efficient training techniques as an approach to training large neural networks. We introduce a novel method called ReLoRA, which utilizes low-rank updates to train high-rank networks. We apply ReLoRA to training transformer language models with up to 1.3B parameters and demonstrate comparable performance to regular neural network training. ReLoRA saves up to 5.5Gb of RAM per GPU and improves training speed by 9-40% depending on the model size and hardware setup. Our findings show the potential of parameter-efficient techniques for large-scale pre-training.
1 Introduction
Scaling has produced increasingly large networks whose training costs are prohibitive, while the need for extreme overparameterization remains poorly understood. Parameter-efficient methods improve accessibility, motivating investigation of whether they can also support pre-training.
- Hundreds of billions of parameters have made training large networks prohibitively expensive for most research groups.
- The theoretical necessity of models with far more parameters than training examples remains poorly understood.
- Alternative scaling approaches offer trade-offs but rarely democratize training or clarify why overparameterized models are needed.
- Parameter-efficient methods have made billion-scale fine-tuning possible on consumer hardware, raising whether they can benefit pre-training.
- ReLoRA trains a high-rank network through aggregated low-rank updates and is evaluated on transformer models up to 1.3B parameters.
2 Method
ReLoRA adapts LoRA’s low-rank updates for high-rank training by repeatedly merging updates, resetting parameters and optimizer state, and warming the learning rate. These mechanisms preserve parameter efficiency while allowing accumulated updates to achieve higher rank.
- A sum of low-rank matrices can have higher rank than either matrix, motivating high-rank training through repeated low-rank updates.
- LoRA represents a weight update as a rank-r product of two trainable matrices, limiting each individual update’s rank.
- ReLoRA periodically merges LoRA matrices into the main weights, reinitializes them, and thereby increases the total rank of accumulated updates.
- Because Adam retains gradient moments, ReLoRA partially resets optimizer state after merges to avoid guiding new updates toward the previous subspace.
- ReLoRA sets the learning rate to zero and warms it back up after optimizer resets, while a short full-rank warm start is used from random initialization.
- Training fewer parameters reduces optimizer memory and distributed bandwidth requirements, and permits larger batch sizes and low-precision frozen weights.
3 Experiments
The experiments train transformer language models across five sizes using compute-optimal, single-epoch data budgets and compare ReLoRA with full-rank and low-rank baselines. The setup applies ReLoRA broadly to attention and fully connected parameters while retaining full-rank embeddings and normalization layers.
- ReLoRA is evaluated on C4 transformer language models ranging from 60M to 1.3B parameters.
- All experiments use a single epoch and at least a compute-optimal amount of data estimated from Chinchilla scaling laws.
- The models use a transformer architecture closely resembling LLaMA, with pre-normalization, RMSNorm, SwiGLU activations, and rotary embeddings.
- ReLoRA uses rank r = 128, bfloat16 operations, and FlashAttention; rank comparisons are additionally performed for the 1.3B model.
- ReLoRA replaces attention and fully connected parameters while embeddings and normalization layers remain full-rank; a matched-trainable-parameter Control baseline is included.
- Training initializes ReLoRA from full-rank checkpoints, resets every 5,000 steps, prunes 99% of optimizer state, and applies 100-iteration loss warm-ups.
- The 1.3B experiments compare ReLoRA runs initialized from 2K, 5K, and 10K checkpoints against a model trained from scratch on 23B tokens.
4 Results
ReLoRA achieves performance close to full-rank training while outperforming LoRA and producing higher-rank updates through repeated low-rank updates. At 1.3B scale, it also reduces training time and GPU memory, although speedups depend on hardware.
- ReLoRA significantly outperforms LoRA in parameter-efficient pre-training and achieves similar performance to full-rank training on upstream and downstream tasks.
- High-rank training through low-rank updates: ReLoRA produces fewer near-zero singular values and a distribution more like full-rank training than LoRA, demonstrating a high-rank learned update.
- Scaling up to 1.3B: At 1.3B scale, ReLoRA outperforms LoRA throughout training, with the gap increasing from 0.56 at 15K steps to 0.96 at 30K steps.Its final perplexity is 17.24, only 0.44 higher than full-rank training.
- Training speedup: ReLoRA took 440 A100-hours, saving 56 A100-hours versus full-rank training, and reduced GPU RAM from 27.8Gb to 22.3Gb at equal microbatch size.In an 8xA100 setup, total wall-clock time fell from 93.5 hours to 86 hours, a 9% relative speed improvement.
- Training speedup: Hardware-dependent speedups reached 42% on 2xRTX3090 and 21% on 6xA6000 Ada, which the authors attribute to GPU memory speed.The 6xA6000 Ada comparison estimates 152 hours for full-rank training and 119 hours for ReLoRA.
- Ablation studies: Restarts and warm starts are essential for good performance, while combining restarts and optimizer resets without a jagged schedule causes divergence.Ablations evaluate restarts, jagged scheduling, optimizer resets, and warm starts.
5 Related work
Related work examines alternatives to scaling, parameter-efficient fine-tuning, and low-rank neural network training. These approaches offer efficiency benefits, but prior methods often leave overparameterization unexplained or do not scale to large transformers.
- Alternative scaling strategies introduce trade-offs but generally do not explain why overparameterized models are needed or democratize their training.
- Parameter-efficient fine-tuning: Parameter-efficient fine-tuning enables billion-scale language and diffusion model fine-tuning on consumer hardware and motivates applying such methods to pre-training.
- Low-rank neural network training: Prior low-rank training methods target compression, regularization, or efficient training, but many are CNN-specific, do not scale well, or lack evaluation on large transformers.
6 Conclusion
ReLoRA adapts parameter-efficient fine-tuning to pre-training by combining multiple low-rank updates into a high-rank network. It consistently outperforms LoRA and achieves performance similar to regular training while reducing training time.
- 6 Conclusion: ReLoRA trains high-rank networks through multiple low-rank updates and periodically merges them into the main parameters.The method also uses optimizer resets and learning-rate re-warmup during training.
- 6 Conclusion: ReLoRA consistently outperforms LoRA for training large transformer models.
- 6 Conclusion: 9% wall-clock time reduction was achieved in an 8xA100 setup, with 20-40% speed improvements on cheaper hardware.
- 6 Conclusion: ReLoRA achieves similar performance to regular training, supporting its potential for improving the efficiency of large model training.
A A Practical guide to ReLoRA
The practical guide identifies settings that reliably support ReLoRA training. Ranks of 64 or 128 work well up to 1B parameters, while learning rate and reset choices affect performance and stability.
- A A Practical guide to ReLoRA: Ranks r ∈ {64, 128} work well for networks up to 1B parameters.
- A A Practical guide to ReLoRA: ReLoRA and LoRA require a 1.5-2 times larger learning rate than regular training or fine-tuning to achieve similar performance.
- A A Practical guide to ReLoRA: Optimizer pruning above 90% showed no significant performance dependence, although larger pruning rates can cause slightly better performance and possible reset-time loss instability.
- A A Practical guide to ReLoRA: A 2K-iteration reset rate performed consistently well in both pre-training and fine-tuning experiments.
B ReLoRA for fine-tuning
ReLoRA fine-tuning was evaluated on T5 models and GLUE, with additional QNLI ablations. The ablations favor rank 64 and show robustness to extensive optimizer-state resets, while ReLoRA did not outperform LoRA overall in fine-tuning.
- B ReLoRA for fine-tuning: ReLoRA fine-tuning was applied to T5-base and T5-large on the GLUE benchmark, with additional T5-large ablations on QNLI.
- B ReLoRA for fine-tuning: ReLoRA fine-tuning does not outperform LoRA according to Table 8.
- B ReLoRA for fine-tuning: Rank 64 provided the best performance in the fine-tuning ablations.
- B ReLoRA for fine-tuning: Performance remained largely unaffected when 99% of optimizer states were reset.
- B ReLoRA for fine-tuning: Two resets were adequate for QNLI classification accuracy with reset rate 4000.
C Learning curves of models pre-trained in the study
The supplementary material provides additional training-loss plots across the study’s pre-trained model sizes. It also documents the layout and variables of the ReLoRA fine-tuning ablation table.
- C Learning curves of models pre-trained in the study: Additional training-loss plots are provided for the 60M, 130M, 250M, 350M, and 1.3B models.
- C Learning curves of models pre-trained in the study: Table 9 varies LoRA rank, optimizer-state pruning percentage, and ReLoRA reset frequency for T5-large fine-tuning on QNLI.
- C Learning curves of models pre-trained in the study: Reset rate denotes the number of iterations between ReLoRA resets.
D Ranks of 130M models
Figures 11 and 12 examine the spectral properties of the learned updates in 130M models.
- Figures 11 and 12 present spectral analyses for 130M models.
- The analysis covers singular-value spectra of learned updates.
- The section focuses on rank-related properties of the learned updates.
E Smaller warm start period
Starting ReLoRA from a shorter warm-start period improves its performance relative to LoRA, although it remains below full-rank training in the reported 2K-step setting.
- ReLoRA consistently outperforms the warmed-started LoRA baseline.
- 1.4 ppl points separate ReLoRA from LoRA after a 2K-step warm start: ppl 23.64 vs 25.08.
- ReLoRA’s absolute performance remains lower than full-rank training in this 2K-step context.
- The experiments validate the hypothesis that LoRA restarts positively impact performance.
- The section includes training-loss experiments for 60M, 130M, 250M, 350M, and 1.3B models.
- Additional analyses compare ReLoRA with LoRA through singular-value spectra and counts of small singular values in 130M models.