Fast Forwarding Low-Rank Training (2024)

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 SGD 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 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.

## introduction modern optimizers provide a spectacular array of tweaks to stabilize training trajectories and accelerate stochastic gradient descent (sgd). yet even with every optimization hack in the modern machine learning toolkit, the expense of training accumulates. in this paper, we ask: what if we just keep going in the same direction until it stops helping? applying this exceedingly simple approach, which we call fast forward, to low-rank training (section 2) allows enormous speedups over standard optimization. we alternate between adam sgd for burn-in and accelerating by line search with a tiny validation set, which provides an ad-hoc optimal step size much larger than that determined by the learning rate. (see figure 1 and section 3.) we experiment with fast forward on three finetuning tasks and four language models ranging from 1.4 to 6.9 billion parameters (section 4). in all cases, we find consistent efficiency gains with fast forward (section 5), reaching the performance of regular low-rank finetuning 41-87% faster. while fast forward works incredibly well in low-rank finetuning, it fails to improve full-rank standard fine-tuning (section 6). we investigate possible causes, showing that the low-rank loss surface is smooth and that fast forwarding along a given direction reduces the role of that direction later in training. we also provide evidence that fast forward could be an even more versatile tool: although we cannot fast forward full-rank standard finetuning, this limitation is not caused solely by decreased effectiveness on non-attention parameters or in higher dimensional subspaces. ## background: low rank adaptors low rank adaptation (lora) @xcite is a parameter-efficient finetuning method that freezes the lm weights and injects trainable low-rank decompositions into each updated matrix, reducing the number of trainable parameters. given a pretrained parameter setting w 0 ∈ r d×k , lora updates the weight with a low-rank decomposition weight decomposed low-rank adaptation (dora) @xcite decomposes the pretrained weight matrix into magnitude and direction, then updates the direction matrix using lora. ## fast forward our proposal, fast forward, is a procedure for accelerating training at regular intervals by selecting an optimal step size with a line search. following warmup, we apply fast forward every t @xmath0, as seen in figure 1 . during a fast forward stage, for each trainable parameter, the difference between weights in the current and previous timesteps is calculated: the direction ∆ w is used to iteratively update w t . in the τ -th fast forward step, the updated weight matrix is given by w t + τ ∆ w . the recursive updates continue until the model's loss on a small validation set l(x val ) stops improving. when a fast forward step causes this validation loss to increase, the fast forward stage concludes, and regular optimization resumes for the next t interval steps before reapplying fast forward. ## results fast forward accelerates training across low rank methods, in all three datasets and all four models. as figure 2a shows, fast forwarding lora cuts 41-66% of finetuning flops for our largest model (lllama-3 8b) and 65-86% for our smallest (pythia 1.4b). fast forwarding dora, meanwhile, cuts 42-69% of finetuning flops for lllama-3 8b and 66-85% for pythia 1.4b (figure 2b ). in terms of train time, fast forwarding lora cuts 63-78% of pythia 1.4b training time and 41-65% of llama-3 8b training time (figure 3 ). although fast forward is more effective on pythia than llama, we see no clear trend as to whether it is generally more effective at smaller scales. as seen in figure 4 , fast forward accelerates segments of training by simulating predicted sgd steps. although it requires more total steps than vanilla training (counting sgd interval steps and simulated fast forward steps), the efficiency of fast forward leads to substantial reductions in computation cost-and appendix d suggests that we could see even greater gains from fast forwarding even more often. fast forward is more effective earlier in training (see appendix b for details), but below we find that even after training converges, we see substantial savings. ## conclusions and future work we have presented fast forward, a simple approach to accelerating low-rank finetuning. our method reduces cost by 41-87% flops in matching 5-epoch adam sgd loss and 56% at loss convergence. a variety of tweaks on our approach are likely to convey further benefits. rather than using a fixed tiny validation set throughout training, we might introduce new ways of sampling that introduce little overhead but avoid the possibility of overfitting. other future work may schedule the sgd interval lengths dynamically or predict the optimal duration for fast forward, reducing the required number of inference runs on the validation set. one reading of our findings is that current optimizers like adam are poorly designed for lowrank approaches. future optimizers might improve these standard approaches by aligning momentum with the known low-rank basis or applying other methods that select better step sizes at low dimensions. as one example of an issue with momentum alignment, @xcite proves that lora's uniform learning rate is suboptimal and proposes a fix; however, their modified optimizer requires substantial tuning and may not significantly increase training efficiency for a full training run.