LoRA+ (separate learning rates for A and B), explained

LoRA+ is LoRA with one change: the two adapter matrices get different learning rates. In LoRA's update $W_0+\frac{\alpha }{r}BA$, the matrix $A$ reads from the layer input and $B$ writes to the output. Hayou, Ghosh and Yu (2024) showed with an infinite-width scaling analysis that training both at the same learning rate cannot make both matrices contribute efficiently as models get wider, and that the fix is to set ${\eta }_B=\lambda {\eta }_A$ with $\lambda \gg 1$. Nothing else about LoRA changes: the same adapters, the same memory, the same merge, the same serving. In practice it is a two-group optimiser configuration and one extra hyperparameter, the ratio $\lambda$. It applies wherever LoRA does, under SFT, DPO or RL objectives, and combines with initialisations such as PiSSA or with QLoRA's quantised base.

The problem and where it came from#

Hu et al.'s original LoRA used one learning rate for both adapter matrices, a choice inherited by most libraries. The scaling-theory literature (the line of work behind $\mu$P) had already shown that, as network width grows, different parameter tensors need learning rates that scale differently with width, otherwise some layers learn too slowly or blow up. Hayou, Ghosh and Yu (UC Berkeley, 2024, ICML) applied that reasoning to LoRA. Their result: with a shared learning rate, either $A$ or $B$ is effectively frozen in the large-width limit, so LoRA operates in a regime where only one side of the adapter learns. Setting ${\eta }_A=\Theta (n^{-1})$ and ${\eta }_B=\Theta (1)$ in width $n$ restores "efficient feature learning" where both sides contribute. Empirically they reported 1% to 2% improvements and up to about 2 times faster convergence at the same compute, with larger gains on harder tasks: on Llama-7B fine-tuned on a 100K subset of Flan-v2, MMLU improved by about 1.3% over the best equal-rate setting, while on MNLI (easy for Llama) the ratio made no difference.

The update rule and the scaling argument#

For a frozen $W_0\in {\mathbb{R}}^{d\times k}$, LoRA trains $B\in {\mathbb{R}}^{d\times r}$ and $A\in {\mathbb{R}}^{r\times k}$ with scale $\gamma =\alpha /r$. LoRA+ runs the same optimiser (typically AdamW) with two parameter groups:

where $u_t^A$ and $u_t^B$ are the optimiser's processed gradients (for Adam, the bias-corrected first moment divided by the root of the second). You tune ${\eta }_A$ on a one-dimensional grid as before, with $\lambda$ fixed.

Why one learning rate is wrong

Follow the change in the adapter's output $z=BAx$ for an input $x$ over one step. It splits into three pieces:

Efficient learning, in the paper's definition, means both ${\delta }^1$ and ${\delta }^2$ are $\Theta (1)$ as width $n$ grows: both matrices move the features by a non-vanishing amount. Now count orders. $A$ multiplies an $n$-dimensional input, so an Adam update of $\Theta ({\eta }_A)$ per entry, aligned with $x$, changes $Ax$ by $\Theta ({\eta }_{A}n)$: the update sums over $n$ coordinates. $B$ multiplies an $r$-dimensional vector $Ax$ whose entries are $\Theta (1)$, so an update of $\Theta ({\eta }_B)$ per entry changes $z$ by $\Theta ({\eta }_B)$. Keeping $B$ at $\Theta (1)$ scale, the two terms are

With ${\eta }_A={\eta }_B=\eta$, making ${\delta }^2$ order one requires $\eta =\Theta (1)$, which makes ${\delta }^1=\Theta (n)$ and the training unstable; keeping ${\delta }^1$ stable requires $\eta =\Theta (1/n)$, which makes ${\delta }^2=\Theta (1/n)\to 0$, so $B$ is effectively frozen. Only ${\eta }_A=\Theta (1/n)$ with ${\eta }_B=\Theta (1)$ satisfies both, which means $\lambda ={\eta }_B/{\eta }_A=\Theta (n)$. The paper proves this for a linear toy model with gradient descent (Proposition 2) and states it for Adam-type optimisers in general networks (Theorem 1, under an alignment assumption). The constant inside $\Theta$ is not given by the theory, so the ratio must be set empirically.

Initialisation changes the right ratio

LoRA needs $BA=0$ at step 0, so one matrix starts at zero. The paper studies both choices: Init[1] is the standard one ($B=0$, $A$ random with variance $\Theta (1/n)$), used by PEFT and the original LoRA; Init[2] is the reverse ($A=0$, $B$ random with variance $\Theta (1)$). Their main experiments found $\lambda \approx 2^4$ best with Init[2] on RoBERTa and GPT-2; with Init[1] the best ratio was smaller, of order $2^2$ to $2^3$; and for Llama experiments a ratio of order $2^1$ to $2^2$ appeared best. Because the commonly quoted default of 16 comes from the Init[2] setting, it is worth knowing which initialisation your library uses before copying it.

Interaction with alpha, rank and the base learning rate

LoRA's scale $\gamma =\alpha /r$ multiplies both factors' contributions, so with Adam it acts like a common multiplier on both ${\delta }^1$ and ${\delta }^2$ and does not change their balance. The ratio $\lambda$ is therefore a separate axis from $\alpha$, and changing $\alpha$ or switching to rsLoRA's $\alpha /\sqrt{r}$ should be followed by retuning ${\eta }_A$, not $\lambda$. Rank enters the counting only through the size of $Ax$, so to first order the best ratio does not depend on $r$; width does the work. Two consequences are practical. A ratio tuned on a 1B model should be raised, not copied, for a model several times wider, since $\lambda$ scales with $n$ in the theory. And because ${\eta }_B$ is now the larger rate, it is the one that sets the stability limit: when you scale up $\lambda$, expect the best ${\eta }_A$ to come down, which is exactly what our toy showed (the best ${\eta }_A$ fell tenfold for each fourfold increase in width while ${\eta }_B$ stayed roughly constant).

Next to its siblings

LoRA+ and rsLoRA address different axes (width and rank) and can be combined. LoRA+ also composes with PiSSA, DoRA and QLoRA, since it only touches the optimiser's parameter groups.

Worked example: balancing the two terms at width 4096#

Use the counting above with Adam's sign-like updates, which move every entry by about the learning rate. Take width $n=4096$, an input with entries of size about 1, and $Ax$ of order 1 (as it is with $A$ initialised at variance $1/n$). Early in training under the standard initialisation, $B$ starts at zero and after one step has entries of order ${\eta }_B$. Then ${\delta }^1\approx {\eta }_B{\eta }_{A}n$, because $A$'s change is summed over $n$ input coordinates and then read out through $B$, and ${\delta }^2\approx {\eta }_B|Ax|\approx {\eta }_B$:

The last row is the theory's prescription taken literally, $\lambda =n$. Real networks put constants in front of every $\Theta$, which is why measured best ratios are in the range 2 to 16 rather than 4,096, but the direction is the same: as width grows, $B$ needs a larger step relative to $A$. The worked numbers also explain the most common LoRA symptom this fixes. If you raise a shared learning rate until training moves, the $A$ side becomes unstable first; if you keep it safe, the $B$ side barely moves.

A toy run: the best ratio grows with width#

We reproduced the paper's style of toy in PyTorch (seed 0). A frozen random two-layer ReLU network maps 8 inputs through a width-$n$ hidden layer $W\in {\mathbb{R}}^{n\times n}$ to one output; the target network is the same with $W$ shifted by a fixed rank-4 matrix. A rank-4 LoRA on $W$ is trained for 200 Adam steps (batch 128, 1,024 training points), and we report the training loss relative to the starting loss. For each width $n\in \{64,256,1024\}$ and ratio $\lambda$ we took the best ${\eta }_A$ over a log grid ($10^{-5}$ to $10^{-1.5}$).

With Init[2] ($A=0$, $B$ random) the effect was large and followed the theory closely. The best ratio was 64 at width 64, 256 at width 256 and 1,024 at width 1,024 (the top of our grid), growing in proportion to width. The best equal-rate loss was 0.155, 0.240 and 0.167 for the three widths; the best LoRA+ loss was 0.046, 0.051 and 0.052, three to five times lower. The best ${\eta }_A$ fell from 3.2e-3 to 3.2e-4 to 3.2e-5 as width grew, while ${\eta }_B=\lambda {\eta }_A$ stayed between 0.03 and 0.2, which is the ${\eta }_A=\Theta (1/n)$, ${\eta }_B=\Theta (1)$ pattern.

With Init[1] ($B=0$, $A$ random, the PEFT default) the picture was different. The best ratio was 1 at widths 64 and 256 and 4 at width 1,024, and the gain at 1,024 was small (0.0339 against 0.0361 with equal rates). Equal-rate training with Init[1] already reached about the loss that Init[2] needed LoRA+ to reach. That matches the paper's report that the optimal ratio is smaller under Init[1], and it is a useful warning: on a standard LoRA setup, expect LoRA+ to be a modest, task-dependent improvement, and tune $\lambda$ over something like $\{1,2,4,8,16\}$ rather than assuming 16.

The algorithm and training loop#

model = add_lora(base_model, r, alpha)                    # freeze base, add A (random), B (zero)
groups = [
    {"params": [all lora_A matrices],               "lr": lr},
    {"params": [all lora_B matrices and 1-D extras], "lr": lr * ratio},   # ratio = lambda
    {"params": [lora embeddings, if any],           "lr": lr_embedding},
]
optimizer = AdamW(groups, weight_decay=0.0)
scheduler = cosine_with_warmup(optimizer)                # scales every group's lr together

for step in range(total_steps):                          # unchanged LoRA training loop
    loss = objective(model(batch)); loss.backward()
    clip_grad_norm_(trainable, 1.0)
    optimizer.step(); scheduler.step(); optimizer.zero_grad()

That is the whole method. The schedule must apply multiplicatively to both groups so the ratio is preserved; standard schedulers do this. Gradient clipping by global norm is fine, because the gradients are unchanged; only the step sizes differ. Data requirements are those of the objective and of LoRA.

Hyperparameters that matter#

Tune in two passes. First fix $\lambda =1$ and find the best LoRA learning rate. Then try $\lambda \in \{2,4,8,16\}$ with ${\eta }_A$ at that rate and at half of it. If the best $\lambda$ is at the top of your grid, extend it; if it is 1 or 2, your setup is not width-limited and LoRA+ will not help much.

Compute and memory cost#

None beyond LoRA's. The trainable parameters, optimiser state, activations and FLOPs are identical: memory is about $2P+16P_a+M_{\text{act}}$ for $P$ base and $P_a$ adapter parameters in bf16 with AdamW. The only real cost is the extra hyperparameter, which at most doubles or triples a small learning-rate sweep. When the paper reports "up to 2 times speedup", it means reaching the same accuracy in about half the steps (RoBERTa-base on MNLI with $\lambda =2^4$), not a faster step.

Measuring whether one side is frozen#

You can check the theory's diagnosis directly instead of inferring it from loss curves. On a fixed probe batch, every few hundred steps, record for a handful of layers the two contributions to the change in adapter output since the previous checkpoint: $\Vert B_{\text{old}}(A_{\text{new}}-A_{\text{old}})x\Vert$ for the $A$ side and $\Vert (B_{\text{new}}-B_{\text{old}})A_{\text{old}}x\Vert$ for the $B$ side. If one is consistently one or two orders of magnitude smaller than the other, that matrix is effectively frozen and the ratio is wrong for your width and initialisation. Under the standard $B=0$ initialisation with equal rates the usual pattern is the opposite of what people expect early on: $B$ moves first, because $A$ has zero gradient at step 0, and the $A$ side catches up only once $B$ is non-zero. A healthy LoRA+ run shows the two contributions within a factor of a few of each other after warmup. The measurement costs two extra adapter forward passes per probe and needs only the saved adapter weights, so it can be run offline on checkpoints.

Failure modes and how to detect them#

Instability from a large ${\eta }_B$

Loss spikes or diverges early, especially with high $\alpha /r$. The effective step on the output side has grown by $\lambda$. Lower ${\eta }_A$ (and so ${\eta }_B$) together, or reduce $\lambda$. Watch $\Vert B{\Vert }_F$ per layer in the first few hundred steps; it should rise smoothly.

No gain at all

The task is easy for the model (the paper's Llama on MNLI case), or the standard $B=0$ initialisation already balances the two sides at your width, as in our toy. Not a failure, just a signal to stop tuning $\lambda$.

Parameter-group mistakes

The higher rate lands on the wrong matrix. In PEFT naming, lora_B is the output-side matrix that is initialised to zero and should get the higher rate; lora_A is the input side. Print the learning rate of each group and a parameter name from it before training. Custom LoRA code that names the matrices the other way round inverts the method.

Ratio lost to a scheduler or wrapper

Some training wrappers rebuild the optimiser from a single learning rate, silently erasing the groups. Log the actual learning rates of both groups at step 0 and mid-run.

When to use LoRA+, and what to use instead#

LoRA+ is nearly free, so the question is less whether to use it than how much tuning effort to spend on it. Treat $\lambda$ as one more axis of the learning-rate sweep you already run. It pays most on wide models, hard tasks that need the adapter to learn new features, and setups that use the $A=0$ initialisation; it pays least on easy tasks, short alignment stages and the standard $B=0$ initialisation at moderate width. If a sweep says $\lambda =1$, believe it.

Minimal implementation#

Two parameter groups are all you need. PEFT ships a helper that builds them for you.

import torch

def loraplus_param_groups(model, lr, ratio=8.0, lr_embedding=1e-6, weight_decay=0.0):
    a, b, emb = [], [], []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if "lora_embedding" in name:
            emb.append(p)
        elif "lora_B" in name or p.ndim == 1:        # output side (zero-initialised) and 1-D extras
            b.append(p)
        else:                                        # lora_A: input side
            a.append(p)
    return [{"params": a, "lr": lr, "weight_decay": weight_decay},
            {"params": b, "lr": lr * ratio, "weight_decay": 0.0},
            {"params": emb, "lr": lr_embedding, "weight_decay": weight_decay}]

# opt = torch.optim.AdamW(loraplus_param_groups(model, lr=1e-4, ratio=8.0))
# with PEFT:
# from peft.optimizers import create_loraplus_optimizer
# opt = create_loraplus_optimizer(model=peft_model, optimizer_cls=torch.optim.AdamW,
#                                 lr=1e-4, loraplus_lr_ratio=8.0)
# then pass optimizers=(opt, None) to a Hugging Face Trainer

In POFE every LoRA recipe exposes the ratio as a single field next to the learning rate, defaulting to 1, so that a sweep over $\lambda$ is one extra axis rather than custom optimiser code.

Further reading#

Hayou, Ghosh and Yu (2024) is the paper; Section 3 has the intuitive toy-model argument and Section 4 the general theorem, with ratio studies in Section 5 and Appendix C. Hu et al. (2021) is LoRA itself. Yang and Hu (2020) on feature learning in infinite-width networks is the theory LoRA+ builds on. Kalajdzievski (2023) is rsLoRA, the companion fix for the rank axis. Meng et al. (2024, PiSSA) is a different answer to slow early LoRA training, through initialisation rather than step sizes.

Frequently asked questions#

References#

1. Soufiane Hayou, Nikhil Ghosh, Bin Yu (2024). LoRA+: Efficient Low Rank Adaptation of Large Models. arxiv.org/abs/2402.12354
2. Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen (2021). LoRA: Low-Rank Adaptation of Large Language Models. arxiv.org/abs/2106.09685
3. Greg Yang, Edward J. Hu (2020). Feature Learning in Infinite-Width Neural Networks. arxiv.org/abs/2011.14522
4. Damjan Kalajdzievski (2023). A Rank Stabilization Scaling Factor for Fine-Tuning with LoRA. arxiv.org/abs/2312.03732
5. Fanxu Meng, Zhaohui Wang, Muhan Zhang (2024). PiSSA: Principal Singular Values and Singular Vectors Adaptation of Large Language Models. arxiv.org/abs/2404.02948

Related reading