LM8-T2 — Low-Rank Adaptation (LoRA)

Learning objectives

After this lesson, you should be able to:

• Write the LoRA layer y = Wx + (α/r)BAx with the shapes of every tensor.
• Count LoRA parameters for a layer and for a model, and compare them with the full weights.
• Explain the zero initialisation of B, the scale α/r and the gradient at the first step.
• Merge an adapter into the base weights and check that the function is unchanged.

The idea: the change is simpler than the weights

Fine-tuning replaces a weight matrix W by W + ΔW. The hypothesis behind LoRA is that ΔW, the change a given task needs, has a low intrinsic rank: it can be written as the product of a tall thin matrix and a short wide one. If so, we can parameterise ΔW = BA directly with two thin matrices and train only them. The big matrix W stays frozen.

The pieces

• A (r × in) projects the input down to r numbers. B (out × r) projects back up to the output width. Their product has the shape of W but rank at most r.
• Scale α/r. The update is multiplied by α/r, with α a constant you choose (often 1 to 2 times r). Dividing by r keeps the size of the update comparable when you change the rank, so the learning rate need not be retuned.
• Initialisation. A gets small random values; B starts at zero. Then BA = 0 and the adapted model is exactly the base model at the start of training, so training begins from a good point instead of from a perturbed one.
• Where. The original LoRA paper adapted the attention matrices; later work found that adapting all linear layers (attention and feed-forward) usually helps at the same total number of parameters. Embeddings and norms are typically left frozen.

Merge

Since Wx + (α/r)B(Ax) = (W + (α/r)BA)x, the trained adapter can be added into W once, and the merged model has exactly the original architecture and speed. Keeping the adapter separate instead allows swapping adapters on a shared base. Merging into quantised weights needs care (Lesson LM8-T4).

Why the first gradient to A is zero

With B = 0, the output does not depend on A at all, so the gradient with respect to A is zero at the first step. B receives a gradient and moves away from zero; from then on A receives gradients too. Training therefore starts by moving B, which is why a small non-zero random A and a zero B work, and why starting both at zero would never learn anything: with both zero, neither gets a gradient.

The mathematics: the update, the gradients and the count

For a frozen weight W ∈ ℝ^out×in, the adapted layer is

x

input vector of width in (or a batch, shape (…, in)).

A, B

the trainable matrices; A is initialised N(0, σ²) with small σ, B with zeros.

r

the rank: a small integer such as 4, 8, 16 or 64.

α

scaling constant (a hyperparameter); α/r multiplies the update.

W

the frozen weight; it receives no gradient and no optimizer state.

Parameters. The adapter has

so the fraction trained is r(in + out)/(in · out). For a square layer of width d this is 2r/d: for d = 4,096 and r = 8, 0.39%. For the four weight matrices of a GPT block (Q/K/V with in = d, out = 3d; output projection d × d; feed-forward d × 4d and 4d × d) the adapter has 16dr parameters per block. GPT-2 small, r = 8: 12 · 16 · 768 · 8 = 1,179,648; r = 16 doubles it to 2,359,296.

Gradients. With loss L and the upstream gradient δ = ∂L/∂y, and s = α/r:

Since B = 0 at the start, ∂L/∂A = 0 and ∂L/∂B ≠ 0, as stated above. Gradients exist only for A and B; the input gradient passed to earlier layers still includes the frozen Wᵀδ.

Rank. The product BA has rank at most r, so the update cannot represent more than r independent directions (Lab LM8-L2 Stage 1). Larger r adds capacity and parameters; returns diminish quickly for many tasks, and Hu et al. found small ranks sufficient for the tasks they studied.

Merge. The merged weight is

(same shape as W); unmerging subtracts the same term.

A worked example: a LoRA layer in PyTorch

The code defines the layer, shows that it equals the base layer at initialisation, that only B gets a gradient at first, counts the parameters, and checks the merge numerically.

import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    def __init__(self, base, r, alpha):
        super().__init__()
        self.base, self.scale = base, alpha / r
        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.02)
        self.B = nn.Parameter(torch.zeros(base.out_features, r))
    def forward(self, x):
        return self.base(x) + self.scale * (x @ self.A.T) @ self.B.T

torch.manual_seed(0)
base = nn.Linear(768, 768)
for p in base.parameters():
    p.requires_grad = False
lora = LoRALinear(base, r=8, alpha=16)
x = torch.randn(4, 768)

print(torch.allclose(lora(x), base(x)))                       # True: B = 0, so the layer starts as the base layer
print(sum(p.numel() for p in lora.parameters() if p.requires_grad), 8 * (768 + 768))   # 12288 12288
print(768 * 768, round(100 * 12288 / (768 * 768), 2))         # 589824 2.08  (percent of the full matrix)

lora(x).sum().backward()
print(float(lora.A.grad.abs().sum()), float(lora.B.grad.abs().sum()) > 0)   # 0.0 True

# after some training B is no longer zero; simulate it, then merge
with torch.no_grad():
    lora.B.copy_(torch.randn_like(lora.B) * 0.05)
ybefore = lora(x)
with torch.no_grad():
    merged = nn.Linear(768, 768)
    merged.weight.copy_(base.weight + lora.scale * lora.B @ lora.A)
    merged.bias.copy_(base.bias)
print(torch.allclose(merged(x), ybefore, atol=1e-5))         # True: the merged layer computes the same function

The percentage printed is for a single 768 × 768 layer at rank 8: 2.08%. The rank-8 adapter of the whole block adds about 1.9% of the weights of the block it adapts.

Choices that matter in practice

Engineering implications and limitations

Freeze before you wrap. Set requires_grad = False on the base model, then add the adapters (their parameters default to trainable). A frozen embedding and norm stay frozen.

Save only the adapter. Store A and B (megabytes) with the base model’s identity, not the whole model. Libraries such as PEFT automate this, but the contents are just these matrices.

Initialisation matters. A zero B keeps the start identical to the base model. Initialising B randomly perturbs the model from step 0.

Test equality at initialisation and after merge. Both are one-line assertions and catch most bugs (the notebook has them).

Dtype. Keep adapter parameters in float32 even if the base is 16-bit or quantised; small updates vanish in low precision.

Rank search. Try 8 and 32 on the same data and look at the validation loss. If the two are equal, the smaller one is enough.

Watch overfitting. A LoRA with a high rank on a few hundred examples can still memorise them. Keep a validation set and stop early.

Serving many adapters. Keep one base in memory and apply a different (A, B) per request; batching requests with different adapters needs support from the serving framework.

Common misconceptions

Engineering takeaways

1. y = Wx + (α/r) B A x with A (r × in) small random, B (out × r) zero: the model starts as the base model.
2. An adapter has r·(in + out) parameters instead of in·out; per GPT block with the four weight matrices, 16·d·r.
3. ∂L/∂A = 0 at the first step because B = 0; B moves first and then A learns.
4. Merge with W′ = W + (α/r)·B·A for no inference overhead; unmerge by subtracting.
5. Rank bounds what the update can represent; pick it by validation, adapt all linear layers, use a higher learning rate than for full fine-tuning.

References and further reading

See Appendix E of the book (parameter-efficient finetuning with LoRA) for an implementation of the LoRA layer and its use in a GPT model.

• Hu, E. J. et al. (2022). LoRA: Low-Rank Adaptation of Large Language Models. arXiv:2106.09685.
• Hu et al. (2022), LoRA (research paper)
• Build a Large Language Model (From Scratch), Appendix E: Parameter-efficient Finetuning with LoRA (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: the memory and compute of training in detail, with the arithmetic for weights, gradients, optimizer state and activations.

Related content

Nothing below is required to finish this lesson.

• Before
  LM8-T1: Full Versus Parameter-Efficient Fine-Tuning

  Compare the cost and quality trade-offs of updating all weights or a small subset.

• Lab
  LM8-L2: LoRA Low-Rank Update Numerics

  Build a low-rank update, check its rank and parameter count, and see how rank limits what can be learned.

• Lab
  LM8-L3: LoRA Fine-Tuning Experiment

  Add LoRA layers to a small GPT, freeze the base weights and compare quality, parameters and memory with full fine-tuning.