Learning Rate Scaling across LoRA Ranks and Transfer to Full Finetuning

TL;DR

Introduces Maximal-Update Adaptation (µA), the LoRA analogue of µP: a theoretical framework that derives how the optimal learning rate should scale with both model width $n$ and LoRA rank $r$ to give stable, non-vanishing feature updates in the joint limit $n, r \to \infty$ . The analysis separates LoRA design choices into two regimes — rank-dependent and rank-invariant — depending on which factor is randomly initialized (Init[A] vs Init[B]) and how the multiplier $\alpha$ scales with $r$ . The headline practical claim: under Init[B] (random $B$ , $A=0$ ) with constant $\alpha=1$ , the optimal LoRA learning rate is rank-invariant and coincides with the optimal FFT learning rate, so an LR tuned for LoRA transfers directly to full finetuning. Verified across SFT (Llama-3.2-1B, Qwen2.5-3B, RoBERTa-large, ViT-H/14, Qwen3-VL-2B), RLVR, and Stable Diffusion finetuning.

Key claims

Under Init[A] with constant $\alpha=1$ the optimal LR scales as $r^{-1/2}$ at fixed width: doubling rank requires halving the LR; observed across Llama-3.2-1B, RoBERTa-large, Qwen2.5-3B, ViT-H/14, Qwen3-VL-2B [§4.1 Corollary 4.4, §5.2 “Init[A] with $\alpha=1$ ”, Fig. 2a, Fig. 3].
Under Init[A] with $\alpha=1/r$ the optimal LR is rank-invariant but the width scaling is different from FFT’s, and the stable LR range is narrow — small overshoots diverge, especially on Qwen2.5-3B and RoBERTa-large [§4.1, §5.2 “Init[A] with $\alpha=1/r$ ”, Fig. 2b, Fig. 4].
Under Init[B] with constant $\alpha=1$ the optimal LR is rank-invariant and matches FFT’s optimal LR, enabling direct LR transfer from LoRA → FFT [§4.2 Corollary 4.6, §5.2 “Init[B] with $\alpha=1$ ”, Fig. 2c, Fig. 5].
The decomposition $\Delta h = T_1 + T_2 + T_3$ of the LoRA feature increment isolates contributions from training $A$ only, $B$ only, and the bilinear $\Delta A \Delta B$ cross-term; under each stable scaling exactly one $T_i$ is $\Theta(1)$ and the others vanish, giving a “learning-efficiency imbalance” diagnostic between the two factors [§3.1 Eq. 1, §4.1, §4.2].
Under Init[A], stable scaling forces $T_1 = \Theta(1)$ while $T_2, T_3 \to 0$ — i.e. learning happens almost entirely through updates to $B$ and the adapter behaves like a random projection with $A$ effectively frozen at init [§4.1 “Learning efficiency”].
Under Init[B] with $\alpha=1$ the imbalance flips — learning happens primarily through $A$ when $r \ll n$ , and both factors contribute meaningfully only when $r$ is on the order of $n$ [§4.2 “Learning efficiency”].
The intermediate-feature scale (the $r$ -dimensional bottleneck activation) is suppressed under Init[A] with $\alpha=1$ as rank grows but not under $\alpha=1/r$ , which the authors offer as the mechanistic reason Init[A] with $\alpha=1/r$ exhibits a narrower stable LR range [§4.1 “Intermediate feature scaling”].
µA scaling rules transfer beyond SFT: in RLVR they match the SFT trends, and crucially, suboptimal LRs increase wall-clock training time by 4–6× — so the µA prescription pays off in time, not just final reward [§5.3, Fig. 6c].
For text-to-image diffusion finetuning (Stable Diffusion), the same Init[B] + $\alpha=1$ optimum coincides with FFT’s; FFT is more LR-sensitive on UNets (sharp FID rise above optimum), making LoRA-tuned LR transfer especially valuable as a low-memory proxy [§5.4, Fig. 7].
Theoretical assumptions: SignSGD as a tractable proxy for Adam, single-sample loss, joint limit $n, r \to \infty$ with $r/n \to 0$ , only one LoRA layer trainable to isolate its contribution [§3.2, §4 Assumptions 4.1–4.2].
Practical workflow advocated: sweep LR with Init[B] + $\alpha=1$ on a mid-tier GPU (LoRA uses ~half the GPU memory of FFT), then transfer the optimum directly to FFT on higher-capacity hardware with minimal re-tuning [§5.3, Fig. 6d].

Method

The analysis parameterizes the LoRA layer as $W + \alpha B A$ with $B \in \mathbb{R}^{n \times r}$ , $A \in \mathbb{R}^{r \times n}$ , takes the joint asymptotic limit $n, r \to \infty$ , and asks when the LoRA-induced feature increment $\Delta h_t = h_t^{\text{LoRA}} - h_0^{\text{LoRA}}$ remains $\Theta(1)$ — bounded but non-vanishing — at every step $t$ of a fixed-horizon training. This is the LoRA analogue of µP’s “stable feature learning” condition. The feature increment decomposes as $T_1 + T_2 + T_3$ (one-sided $A$ -update, one-sided $B$ -update, bilinear cross-term) and the authors require each $T_i = O(1)$ with at least one $= \Theta(1)$ .

Under each configuration (Init[A] vs Init[B], and choice of $\alpha \in \{1, 1/r, 1/\sqrt{r}\}$ ) this yields a closed-form scaling rule for the learning rate $\eta$ as a joint power of $n$ and $r$ . The optimizer is abstracted as SignSGD (tractable proxy for Adam) under a single-sample loss; signals are assumed bounded forward and backward (Assumption 4.1). Full proofs are in Appendix A; the body lists prescriptions in Table 1 (Init[A]) and Corollary 4.6 (Init[B] with $\alpha=1$ ).

Experiments use a unified AdamW recipe (weight decay 0.1, global-norm gradient clip 1.0, linear warmup over 3% of steps to peak LR, cosine decay) across: SFT on Llama-3.2-1B/Tulu3, Qwen2.5-3B/OpenThoughts-114k, RoBERTa-large/ANLI, ViT-H/14/ImageNet-1K, Qwen3-VL-2B/LLaVA-Instruct-Mix; RLVR with verifiable rewards; and Stable Diffusion finetuning. For each (config, rank) pair, peak LR is swept on a log₂ grid with step 0.5 and the optimum selected by EMA-smoothed final training loss (also validation-metric-based sweeps in appendices).

Results

Init[A], $\alpha=1$ , rank sweep: optimal $\log_2(\eta)$ decreases by ~1 unit per doubling of $r$ (matching the predicted $r^{-1/2}$ ) across Llama-3.2-1B, RoBERTa-large, Qwen2.5-3B, ViT-H/14, Qwen3-VL-2B [Fig. 3, §5.2].
Init[A], $\alpha=1/r$ , rank sweep: optimal LR falls within a ~0.5-unit log₂ band across ranks (rank-invariant); the stable range is narrow and overshoots diverge, especially on Qwen2.5-3B [Fig. 4].
Init[B], $\alpha=1$ , rank sweep: optimal LR is tightly clustered across ranks and aligns with the FFT optimum on every model except RoBERTa-large, where FFT prefers a slightly smaller LR (attributed to the jointly-trained classification head) [Fig. 5].
RLVR (Init[A] $\alpha=1/r$ ): rank-invariant optimum confirmed; suboptimal LRs cost 4–6× wall-clock time with no reward gain [Fig. 6a, Fig. 6c].
RLVR (Init[B] $\alpha=1$ ): rank-invariant optimum confirmed and aligns with FFT optimum; FFT uses ~2× the GPU memory of LoRA (gradient buffers + optimizer states), making the LoRA→FFT transfer workflow concretely cheaper [Fig. 6b, Fig. 6d].
Stable Diffusion finetuning: Init[A] $\alpha=1/r$ optimum is rank-invariant; Init[B] $\alpha=1$ optimum is rank-invariant and matches FFT; FFT is sharply LR-sensitive (FID spikes at higher LRs) while LoRA is more forgiving [Fig. 7].
Scaling-rule fit quality: across all SFT runs the rank-invariant configs cluster within one log₂ grid step; rank-dependent configs shift monotonically by one grid step per doubling of $r$ , matching the $r^{-1/2}$ prediction qualitatively.

Why it’s interesting

This is the third paper on the wiki to attack the where does the right LR live? question and the first to push the analysis to finetuning rather than from-scratch pretraining. How to Set the Learning Rate for Large-Scale Pre-training? argues µP’s stability motivation has eroded for LLM pretraining; Controlled LLM Training on Spectral Sphere argues µP needs stronger enforcement via spectral-sphere optimizers. µA stakes out a third position: µP-style asymptotic analysis still has bite, but the right place to apply it is the LoRA configuration choice (initialization + $\alpha$ -schedule), where the bilinear $B A$ structure introduces a new scalable dimension (rank) on top of width. The “Init[B] + $\alpha=1$ transfers to FFT” finding is the kind of cheap, falsifiable hypothesis that’s directly checkable in Luma’s pretraining-adjacent finetuning runs.

The Slack post is explicit about the immediate Luma-internal question: Luma’s default LoRA init is $A=0, B \sim \mathcal{N}, \alpha=1/r$ — which µA’s Table 1 / Appendix A.5 does cover (it’s Init[B] with $\alpha=1/r$ , not one of the three “named” configs in the tweet) but is not the configuration the paper’s recommended LR-transfer workflow uses. The cleanest internal action is: (a) switch the default to Init[B] + $\alpha=1$ for finetuning, (b) tune LR on a small-rank LoRA, (c) transfer directly to FFT. The diffusion-side experiment (§5.4) is the closest analogue to Luma image/video model finetuning and the rule held there too.

Note also the learning-efficiency imbalance claim: under both Init[A] and Init[B] (small $r$ ), one of the two LoRA factors contributes essentially nothing to feature learning. Under Init[A], updates to $A$ are noise; under Init[B] with $r \ll n$ , updates to $B$ are noise. This is a separate question from LR transfer and suggests that the standard LoRA factorization may be over-parameterized for the rank regime in actual use.