anthology-of-the-sota

SOTA-426: When you merge and reinitialise a low-rank adapter mid-training, prune the optimizer state and re-warm the learning rate from zero — both, not either

4 min read

Properties6
StatusProposed
Consensusunreplicated
Tagstraining-optimization, adaptation-and-tuning, model-stability
SourceLIT-104 — ReLoRA: High-Rank Training Through Low-Rank Updates
Introduced byLIT-104 — ReLoRA: High-Rank Training Through Low-Rank Updates
ExtendsSOTA-184 — Adapt a pretrained model by training a low-rank update to each weight matrix, not the matrix itself

Graph View

Source

Lialin et al. (2023), LIT-104 — ARXIV-2307.05695, read in full (v4, December 2023) as NOTE-024. The evidence is Table 6 and the paragraph under “Adding restarts and optimizer resets” in §4.2; the recipe is Algorithm 1 and Figure 2.

When this applies

You are training a low-rank adapter (SOTA-184’s W + s·W_A W_B) with Adam, and at some point during training you fold the adapter into the frozen weight and start a fresh one — W ← W + s·W_A W_B, W_A re-drawn (Kaiming), W_B zeroed — so that the next adapter can learn a direction the last one could not. ReLoRA does this every 2K–5K steps of pretraining, so that a sum of rank-r updates reaches a rank no single one has.

This practice is about what the restart needs. It does not say to restart; the next section but one says what the evidence for that is, and it is much thinner.

Do this, at every restart

  1. Prune the optimizer state of the new adapter. ReLoRA zeroes 99% of the Adam moments by magnitude. The reason is specific to Adam: with β₁, β₂ at 0.9–0.999 the moments accumulated for the old W_A steer the new one back into the same subspace, which is exactly what the restart was meant to escape. Keep the state and the restart is a no-op.
  2. Drop the learning rate to zero and re-warm it. ReLoRA’s “jagged cosine” (Figure 2) sets the rate to zero at each reset and warms back onto the underlying cosine over 50–100 steps; the 1.3B run used 50. Reset the state without this and the run diverges.

The ablation that makes this one practice rather than two (Table 6, 130M, no warm start, perplexity):

restartsoptimizer resetre-warmppl
–––34.17 (plain LoRA)
✓––34.25
✓✓–diverged
✓–✓34.29
✓✓✓29.77

Either half alone is worthless or worse. That is the negative half of this practice, and it carries its citation in the same table — the shape ADR-042 asks a “do not do this” to have.

What the evidence does not cover

  • The ablation is from random initialisation. The method as recommended starts from a full-rank warm start, and there the reset and re-warm were never ablated separately: from a warm start the paper reports only the full recipe against no restarts at all. The divergence finding is established for one configuration, at one scale, apparently one run per arm.
  • “Optimizer reset” in the ablation is not fully specified. The method prunes 99% of the state; the prose calls the diverging arm “a naive optimizer reset”, which could mean a full one. Appendix A reports no significant dependence on the pruning fraction above 90%, and that higher fractions can bring “possible loss instabilities during the reset” — which is the re-warm’s job. Algorithm 1 as printed prunes the moments of W_A only.
  • The 50-step re-warm is anecdote. The paper says the partial reset plus the jagged schedule allowed “as low as 50 steps, instead of hundreds” for an optimizer initialised from scratch, from “initial experiments” with no table.
  • The rationale is Adam’s. The argument for the reset is about momentum and second-moment state; the paper does not test another optimizer.

Whether to restart at all, which this practice does not say

The evidence that restarting pays is real and small, and the paper’s own ablation is the reason to read it carefully:

  • Most of ReLoRA’s gain is the warm start. At 130M, LoRA plus the warm start alone reaches 25.46; the full restart recipe on top reaches 25.04, and regular training 23.65. Across 60M–1.3B the restarts beat warm-started LoRA by 0.27–0.96 perplexity (Tables 2 and 4), and with only a 2K-step warm start by 1.4 (Appendix E).
  • More restarts is not better. “Online ReLoRA”, merging every ~100 steps while resetting every 2–5K, is worse at 250M and 1.3B (Table 5).
  • In fine-tuning it does not help. On GLUE with T5-base and T5-large, ReLoRA does not outperform plain LoRA (Appendix B, Table 8).

So the practice is conditional on a decision this record does not make for you. If you are restarting, restart this way.

Known implementations

  • The authors’ reference implementation, github.com/guitaricet/relora.

Backlinks