Why LoRA works: fine-tuning a giant model with a tiny matrix

Rank r = 1

Loading the weight update…

Trained adapter (r = 64): 87.0% on SST-2 (base model 60.9%). This is the trained model's score, so it stays the same as you move the slider.

How is the error computed?

\begin{aligned} \text{error}(r) &= \frac{\lVert \Delta W - \Delta W_r \rVert_F}{\lVert \Delta W \rVert_F} \\ &= \sqrt{\frac{\sum_{i > r} \sigma_i^2}{\sum_i \sigma_i^2}} \end{aligned}

The rebuild keeps the r largest singular values of the full 768 × 768 update, which is the best rank-r approximation there is. The error comes from the singular values it drops.

Top-left 64 × 64 corner of GPT-2 small, layer 6, query matrix W_q, fine-tuned on SST-2 sentiment.

New to LoRA? Start with the one-minute version.

Why does LoRA work?

GPT-3 has 175 billion parameters. When the LoRA authors fine-tuned it, they trained 4.7 million numbers and left every original weight frozen. It scored about the same as fine-tuning the whole model, and better on two of the three tasks they tested. That is one trained number for every 37,000 frozen ones.

So how do you fine-tune a giant model while training almost none of it? The slider above shows it on one weight matrix from a real GPT-2 fine-tune.

The slider sets how many pieces the rebuild may use. On the LoRA update, a handful of pieces already captures most of the pattern. Switch to the full fine-tune and the rebuild stays blurry much longer, because that update changed every direction a little. The LoRA adapter here was trained with room for 64 pieces and scored 87.0% on this sentiment task, against 88.9% for the full fine-tune and 60.9% before training. Most of its change sits in a few pieces, and a separate adapter trained with room for just one piece still reached 83.6%.

I have set r=8 in a LoRA config and moved on more than once. It worked, so I never asked why eight was enough for a matrix with millions of entries. This essay is me finally asking. The fine-tune behind the slider is in ml-explainers/lora and reruns on a laptop in minutes.

The change a fine-tune needs can be far smaller than the matrix it changes.

What is LoRA? The one-minute version

If you have never fine-tuned a model, here is LoRA before any math. Press play or drag the timeline.

Step 1 of 7

A trained model is a giant grid of numbers

Every square is one weight. GPT-3 has 175 billion, all set during pretraining.

175,000,000,000 numbers in GPT-3

The big grid never moves. All the learning happens in two thin strips, and their product can only paint a striped change, where every row is the same pattern scaled up or down. The rest of this essay is about why that is enough, starting with rank: how many strip pairs you use.

Terms you'll see

Fine-tuning:: extra training so a pretrained model does a new task.
Adapter:: the trained pair of strips for one task (B and A, later on).
Rank r:: how many strip pairs an adapter uses.
SST-2 and WikiSQL:: two benchmarks, one for movie-review sentiment and one for turning questions into SQL.
W_q and W_v:: two weight matrices inside each attention layer that LoRA often adapts.

What is the rank of a matrix?

In the determinant essay, the plane collapsed to a line when the two columns lined up and the determinant hit zero. There is a name for the number of directions the outputs can still reach.

Rank

The number of independent directions a matrix's outputs can reach. A 2 × 2 matrix has rank 2 when its columns point different ways, rank 1 when they line up, and rank 0 only when every entry is zero.

Drag the two column arrows below and watch the rank drop as they line up.

\operatorname{rank} \begin{bmatrix} 1.5 & 0.5 \\ 0.5 & 1.2 \end{bmatrix} = 2

Rank

2

The columns point in different directions, so the outputs fill the plane.

det = 1.55

Orange is the first column, blue is the second. The faint orange grid is where the matrix sends the plane.

Try this

Get rank 1 without any zero entries in the matrix.

Rank 1 means every output lands on one line, however big the matrix is. A 4,096 × 4,096 matrix can have rank 1 too: almost 17 million entries, one direction of output. LoRA is built on that gap.

Low-rank matrices: one column times one row

Weight matrices are rarely square, so from here on a matrix has d rows and k columns. The simplest way to build one with rank 1 is to multiply a column by a row. Take a column b with d entries and a row a^T with k entries. Their product b a^T is a d × k grid where the entry in row i, column j is b_i times a_j.

Watch the grid fill in square by square. Then change b or a, or tap any square.

b1a1

Tap a square to see which numbers made it.

Rank

1

Numbers stored

9

to produce 20 entries

The left column is b and the top row is a. Each square is its row's b times its column's a. Orange is positive, blue is negative.

Every column of b a^T is b scaled by one entry of a, so every column points the same way. That is rank 1. Add a second pair and the columns can point two ways, so the rank goes up to 2. Keep going and a rank-r matrix is a sum of r column-times-row pairs. Put the r columns side by side as B and the r rows on top of each other as A, and the sum is the product BA:

BA = \sum_{i=1}^{r} b_i\, a_i^\top

\underbrace{B}_{d \times r}\;\underbrace{A}_{r \times k}

The grid above made 20 entries from 9 stored numbers. In general, B and A hold r(d + k) numbers, against d · k for the full matrix.

d \cdot k \quad \text{vs} \quad r(d + k)

Try this

Add a second pair that does not raise the rank.

How LoRA works: freeze W, train BA

A layer in a neural network multiplies its input by a weight matrix W, as the matrix essay showed. Full fine-tuning changes every entry of W. LoRA leaves W frozen and learns the change as a low-rank product, $\Delta W = BA$ , so the layer computes:

h = Wx + BAx

In the one-minute version, B was the side strip and A was the top strip.

LoRA

Low-rank adaptation. Keep the pretrained weight matrix W frozen and learn its update as the product of two thin matrices, B (d × r) and A (r × k), with r much smaller than d and k.

Training adjusts B and A and never touches W. The widget shows two details that make this practical. Drag x, then press Reset B to zero.

B (2 × 1)

A (1 × 2)

Wx (frozen)

(1.35, 1.4)

BAx (update)

(0.93, -0.62)

h = Wx + BAx

(2.28, 0.78)

W is fixed. Drag x, change B and A, then reset B to zero or merge the two.

First, B starts at zero, so BA is zero and h equals Wx. The model starts exactly as it was.

B = 0 \;\Rightarrow\; h = Wx

Second, after training, W + BA merges into one matrix, so the fine-tuned model runs as fast as the original.

W' = W + BA, \qquad W'x = Wx + BAx

Try this

Set B to zero and move x anywhere. Can you make h differ from Wx?

Hu et al. scale BA by a constant $\alpha / r$ , and libraries follow them. They differ on how A starts out, but PEFT and most other libraries start B at zero.

Why a small LoRA rank is enough

Why would a few directions be enough? Aghajanyan, Zettlemoyer and Gupta tuned RoBERTa-Large to 90% of full fine-tuning on a paraphrase task while training about 200 numbers, projected back into the full model. Their reading is that pretraining leaves a model only a few directions away from many tasks.

Hu et al. looked inside the update too: it amplifies directions already in W that pretraining had left weak.

This has limits. On code and math, Biderman et al. found LoRA underperforms full fine-tuning in both instruction tuning and continued pretraining, with a smaller gap for instruction tuning. The full fine-tunes learned updates with a rank 10 to 100 times higher than typical LoRA settings, and LoRA learned less but forgot less. Hu et al. add that a task in a language the model never saw could need close to full rank.

For many tasks the useful change lives in a few directions, so a small r loses very little. For large, hard tasks it can lose more.

LoRA parameter count and memory

Here is r(d + k) at real model sizes. Pick a model, the matrices to adapt, and a rank.

Model

Layers

Matrices per layer

Rank r

Trainable parameters

18,874,368

0.07% of each adapted matrix

Adapter size (FP16)

37.7 MB

2 bytes per number

Update, per matrix per token

98,304 multiply-adds

vs 150,994,944 for a full d × d update

Layer, per matrix per token

151,093,248 unmerged

150,994,944 after merging, the same as the base model

Hu et al. report about 35 MB for GPT-3 at r = 4 on W_q and W_v. Counting 2 bytes per number gives 37.7 MB, which is 36 MiB. Biderman et al. found that adapting the MLP layers beats attention alone. This calculator counts attention only.

The GPT-3 checkpoint shrinks from 350 GB to about 35 MB with r = 4 on W_q and W_v, so 100 task adapters plus one base model take about 354 GB, against 35 TB for 100 full copies. Training needed 1.2 TB of GPU memory with full fine-tuning and 350 GB with LoRA.

Applying the update as B(Ax) costs r(d + k) multiply-adds, far less than a full d × k update. Unmerged, that work comes on top of Wx, so the adapter is slightly slower than the base model. Merged, it is one matrix multiply again.

B(Ax):\; r(d + k) \text{ multiply-adds}

What the rank slider does not show

The slider rebuilt each update from its strongest directions, without showing what those directions are. It counts how many a fine-tune needs, and what each one does is still an open question.

One question to take with you: if a fine-tune only needs a few directions, which directions does your task need?

The code that trained and measured these updates is in ml-explainers/lora. The next essay in this series is on probability and loss functions.

Why does LoRA work?

What is LoRA? The one-minute version

What is the rank of a matrix?

Low-rank matrices: one column times one row

How LoRA works: freeze W, train BA

Why a small LoRA rank is enough

LoRA parameter count and memory

What the rank slider does not show

Sources