Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Fast Forwarding Low-Rank Training

Adir Rahamim, Naomi Saphra, Sara Kangaslahti, Yonatan Belinkov

TL;DR

Fast Forward presents a simple line-search-based acceleration for low-rank finetuning by alternating conventional Adam SGD steps with Fast Forward stages that extrapolate along the most recent weight update direction. By computing $ΔW = W_t - W_{t-1}$ and updating weights as $W_t + τ ΔW$ until the tiny validation loss stops improving, and by triggering these stages every $T_{interval} = 6$ steps, the method achieves substantial FLOPs and training-time reductions without harming final performance. Across three finetuning tasks and four models (1.4B–8B parameters), Fast Forward yields 41–87% FLOPs savings and 40–81% training-time reductions while preserving accuracy and standard benchmark results. The paper also analyzes why Fast Forward is ineffective for full-rank training, attributing it to the projection structure introduced by LoRA and suggesting directions for future optimizers and dynamic scheduling tailored to low-rank regimes.

Abstract

Parameter efficient finetuning methods like low-rank adaptation (LoRA) aim to reduce the computational costs of finetuning pretrained Language Models (LMs). Enabled by these low-rank settings, we propose an even more efficient optimization strategy: Fast Forward, a simple and effective approach to accelerate large segments of training. In a Fast Forward stage, we repeat the most recent optimizer step until the loss stops improving on a tiny validation set. By alternating between regular optimization steps and Fast Forward stages, Fast Forward provides up to an 87\% reduction in FLOPs and up to an 81\% reduction in train time over standard SGD with Adam. We validate Fast Forward by finetuning various models on different tasks and demonstrate that it speeds up training without compromising model performance. Additionally, we analyze when and how to apply Fast Forward.

Fast Forwarding Low-Rank Training

TL;DR

Fast Forward presents a simple line-search-based acceleration for low-rank finetuning by alternating conventional Adam SGD steps with Fast Forward stages that extrapolate along the most recent weight update direction. By computing and updating weights as until the tiny validation loss stops improving, and by triggering these stages every steps, the method achieves substantial FLOPs and training-time reductions without harming final performance. Across three finetuning tasks and four models (1.4B–8B parameters), Fast Forward yields 41–87% FLOPs savings and 40–81% training-time reductions while preserving accuracy and standard benchmark results. The paper also analyzes why Fast Forward is ineffective for full-rank training, attributing it to the projection structure introduced by LoRA and suggesting directions for future optimizers and dynamic scheduling tailored to low-rank regimes.

Abstract

Parameter efficient finetuning methods like low-rank adaptation (LoRA) aim to reduce the computational costs of finetuning pretrained Language Models (LMs). Enabled by these low-rank settings, we propose an even more efficient optimization strategy: Fast Forward, a simple and effective approach to accelerate large segments of training. In a Fast Forward stage, we repeat the most recent optimizer step until the loss stops improving on a tiny validation set. By alternating between regular optimization steps and Fast Forward stages, Fast Forward provides up to an 87\% reduction in FLOPs and up to an 81\% reduction in train time over standard SGD with Adam. We validate Fast Forward by finetuning various models on different tasks and demonstrate that it speeds up training without compromising model performance. Additionally, we analyze when and how to apply Fast Forward.
Paper Structure (20 sections, 2 equations, 14 figures, 3 tables)

This paper contains 20 sections, 2 equations, 14 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Background: Low rank adaptors
  3. Fast Forward
  4. Review of related approaches
  5. Experiments
  6. Models and data.
  7. Training and Evaluation Procedure
  8. Results
  9. FF does not harm long-term accuracy.
  10. FF does not harm performance on a standard benchmark.
  11. When does Fast Forward work?
  12. Why can't we Fast Forward at full rank?
  13. Fast Forward works better as we increase the rank of LoRA.
  14. Fast Forward doesn't work in full rank settings even when restricted to the attention layer.
  15. Conclusions and Future Work
  16. ...and 5 more sections

Figures (14)

  • Figure 1: Fast Forward algorithm. We alternate between SGD and Fast Forward, exiting the Fast Forward stage when the loss on a tiny validation set stops improving.
  • Figure 2: The percentage of FLOPs saved during (\ref{['fig:lora_results']}) LoRA and (\ref{['fig:dora_results']}) DoRA finetuning with Fast Forward to match test loss after 5 epochs of regular Adam SGD training. Fast Forward saves 41--87% FLOPs, depending on the task.
  • Figure 3: The percentage of train time saved during LoRA finetuning with Fast Forward to match test loss after 5 epochs of regular Adam SGD training. Fast Forward saves 40--81% of the training time, depending on the task.
  • Figure 4: Training Pythia-6.9B on the chat tuning task, with other models in Appx \ref{['app:all-models']}. Red dots represent SGD steps and green dots represent Fast Forward steps. The blue line shows vanilla Adam SGD training.
  • Figure 5: Test loss on the plane intersecting the pretrained model $\mathbf{W}_0$ and the models trained with Adam SGD $\mathbf{W}_{\textrm{SGD}}$, and with Fast Forward $\mathbf{W}_{\textrm{FF}}$. Axis scale corresponds to the norm of differences $\|\mathbf{W}_{\textrm{FF}} - \mathbf{W}_0\|_2$.
  • ...and 9 more figures