Demo: Finetuning Accuracy vs. LoRA Rank

Pretrain a small network on upright MNIST digits, then hand it a task the pretraining never covered: the same digits, rotated \(90^\circ\). We compare full finetuning against LoRA at several ranks \(r\), using only a few hundred rotated examples, which is the realistic setting where finetuning data is scarce.

import copy
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import scienceplots
from sklearn.datasets import fetch_openml

plt.style.use(["science", "no-latex"])
TEAL, CARDINAL, GRAY = "#009090", "#9c1b33", "#c9c9c9"
torch.manual_seed(17)
rng = np.random.default_rng(17)

X, y = fetch_openml("mnist_784", version=1, return_X_y=True, as_frame=False, parser="auto")
y = y.astype(int)
idx = rng.choice(len(X), size=15000, replace=False)
X, y = X[idx] / 255.0, y[idx]
n_train = 6000
Xtr, ytr = X[:n_train], y[:n_train]
Xte, yte = X[n_train:], y[n_train:]

Xtr_t = torch.tensor(Xtr, dtype=torch.float32)
ytr_t = torch.tensor(ytr, dtype=torch.long)
print(f"pretraining on {n_train} upright digits")

pretraining on 6000 upright digits

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)
    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

model = MLP()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
lossfn = nn.CrossEntropyLoss()
for epoch in range(15):
    perm = torch.randperm(len(Xtr_t))
    for i in range(0, len(Xtr_t), 128):
        b = perm[i:i + 128]
        opt.zero_grad()
        loss = lossfn(model(Xtr_t[b]), ytr_t[b])
        loss.backward()
        opt.step()

Xte_t = torch.tensor(Xte, dtype=torch.float32)
yte_t = torch.tensor(yte, dtype=torch.long)
with torch.no_grad():
    acc_upright = (model(Xte_t).argmax(1) == yte_t).float().mean().item()
print(f"pretrained accuracy, upright digits: {acc_upright:.1%}")

pretrained accuracy, upright digits: 92.2%

Now rotate the digits \(90^\circ\) and check the same pretrained model, with no adaptation at all.

def rotate_batch(X):
    imgs = X.reshape(-1, 28, 28)
    return np.rot90(imgs, k=1, axes=(1, 2)).copy().reshape(-1, 784)

Xte_rot = rotate_batch(Xte)
with torch.no_grad():
    acc_rot_noadapt = (model(torch.tensor(Xte_rot, dtype=torch.float32)).argmax(1) == yte_t).float().mean().item()
print(f"same model, rotated digits, no finetuning: {acc_rot_noadapt:.1%}")

n_ft = 300
Xft_t = torch.tensor(Xte_rot[:n_ft], dtype=torch.float32)
yft_t = torch.tensor(yte[:n_ft], dtype=torch.long)
Xrottest_t = torch.tensor(Xte_rot[n_ft:], dtype=torch.float32)
yrottest_t = torch.tensor(yte[n_ft:], dtype=torch.long)
print(f"finetuning on just {n_ft} rotated examples")

same model, rotated digits, no finetuning: 11.8%
finetuning on just 300 rotated examples

Full finetuning

Unfreeze every parameter and train on the small rotated set.

def full_finetune(Xft, yft, Xtest, ytest, epochs=30, lr=1e-3):
    m = copy.deepcopy(model)
    opt = torch.optim.Adam(m.parameters(), lr=lr)
    for _ in range(epochs):
        opt.zero_grad()
        loss = lossfn(m(Xft), yft)
        loss.backward()
        opt.step()
    with torch.no_grad():
        acc = (m(Xtest).argmax(1) == ytest).float().mean().item()
    return acc, sum(p.numel() for p in m.parameters()), m

acc_full, params_full, model_full = full_finetune(Xft_t, yft_t, Xrottest_t, yrottest_t)
print(f"full finetuning: accuracy {acc_full:.1%}, trainable parameters {params_full:,}")

full finetuning: accuracy 56.8%, trainable parameters 101,770

How much rank did that update actually use?

Full finetuning was free to move all \(10 \times 128\) entries of the output layer’s weight matrix independently, in any of the \(10\) singular directions available to it. Before we restrict it, let’s see how many of those directions it really used: take the singular values of \(\Delta\mathbf{W} = \mathbf{W}^{\mathrm{finetuned}} - \mathbf{W}\) and check how much of its Frobenius energy the top few carry.

dW = (model_full.fc2.weight - model.fc2.weight).detach().numpy()
sv = np.linalg.svd(dW, compute_uv=False)
energy = np.cumsum(sv ** 2) / np.sum(sv ** 2)

print("singular values of dW: " + ", ".join(f"{v:.3f}" for v in sv))
for k in [1, 2, 4, 8]:
    print(f"top {k:2d} of {len(sv)} directions hold {energy[k - 1]:.1%} of dW's energy")

singular values of dW: 0.382, 0.324, 0.276, 0.215, 0.173, 0.145, 0.120, 0.109, 0.092, 0.077
top  1 of 10 directions hold 31.4% of dW's energy
top  2 of 10 directions hold 54.0% of dW's energy
top  4 of 10 directions hold 80.2% of dW's energy
top  8 of 10 directions hold 96.9% of dW's energy

LoRA finetuning

Now freeze every pretrained weight and let the output layer move only inside a rank-\(r\) subspace: replace its update by \(\Delta\mathbf{W} = \mathbf{B}\mathbf{A}\) with \(\mathbf{B} \in \mathbb{R}^{d \times r}\) and \(\mathbf{A} \in \mathbb{R}^{r \times k}\), here \(d = 10\) and \(k = 128\). As in the reading, \(\mathbf{B}\) starts at zero and \(\mathbf{A}\) starts small and random, so training begins exactly at the pretrained model.

class LoRALayer(nn.Module):
    def __init__(self, base: nn.Linear, r):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False
        k, d = base.in_features, base.out_features   # W is d x k, as in the reading
        self.A = nn.Parameter(torch.randn(r, k) * 0.01)   # small and random
        self.B = nn.Parameter(torch.zeros(d, r))          # B = 0, so BA = 0 at init
    def forward(self, x):
        return self.base(x) + (x @ self.A.T) @ self.B.T

class LoRAMLP(nn.Module):
    def __init__(self, base_model, r):
        super().__init__()
        self.fc1 = base_model.fc1
        for p in self.fc1.parameters():
            p.requires_grad = False
        self.fc2 = LoRALayer(base_model.fc2, r)
    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

def lora_finetune(r, Xft, yft, Xtest, ytest, epochs=60, lr=1e-2):
    m = LoRAMLP(copy.deepcopy(model), r)
    trainable = [p for p in m.parameters() if p.requires_grad]
    opt = torch.optim.Adam(trainable, lr=lr)
    for _ in range(epochs):
        opt.zero_grad()
        loss = lossfn(m(Xft), yft)
        loss.backward()
        opt.step()
    with torch.no_grad():
        acc = (m(Xtest).argmax(1) == ytest).float().mean().item()
    return acc, sum(p.numel() for p in trainable)

ranks = [1, 2, 4, 8, 16]
results = [lora_finetune(r, Xft_t, yft_t, Xrottest_t, yrottest_t) for r in ranks]
for r, (acc, n_params) in zip(ranks, results):
    print(f"LoRA r={r:2d}: accuracy {acc:.1%}, trainable parameters {n_params:,}")

LoRA r= 1: accuracy 22.8%, trainable parameters 138
LoRA r= 2: accuracy 38.0%, trainable parameters 276
LoRA r= 4: accuracy 53.3%, trainable parameters 552
LoRA r= 8: accuracy 61.3%, trainable parameters 1,104
LoRA r=16: accuracy 69.8%, trainable parameters 2,208

accs = [acc for acc, _ in results]

fig, ax = plt.subplots(figsize=(7, 3))
ax.plot(ranks, accs, color=TEAL, linewidth=1.6, marker="o", markersize=4, label="LoRA")
ax.axhline(acc_full, color=CARDINAL, linestyle="--", linewidth=1.4,
           label=f"Full finetuning ({params_full:,} params)")
ax.set_xscale("log", base=2)
ax.set_xticks(ranks); ax.set_xticklabels(ranks)
ax.set_xlabel("LoRA rank $r$")
ax.set_ylabel("Accuracy on rotated digits")
ax.legend(frameon=False)
plt.show()

What if we had more finetuning data?

Set n_ft_new = 3000 below (ten times more rotated examples). Does full finetuning’s disadvantage shrink, and does LoRA still keep up at small \(r\)?

n_ft_new = 3000    # change me!

Xft_new = torch.tensor(Xte_rot[:n_ft_new], dtype=torch.float32)
yft_new = torch.tensor(yte[:n_ft_new], dtype=torch.long)
Xtest_new = torch.tensor(Xte_rot[n_ft_new:], dtype=torch.float32)
ytest_new = torch.tensor(yte[n_ft_new:], dtype=torch.long)

acc_full_new, _, _ = full_finetune(Xft_new, yft_new, Xtest_new, ytest_new)
print(f"full finetuning with {n_ft_new} examples: {acc_full_new:.1%}")
for r in ranks:
    acc_r, _ = lora_finetune(r, Xft_new, yft_new, Xtest_new, ytest_new)
    print(f"LoRA r={r:2d} with {n_ft_new} examples: {acc_r:.1%}")

full finetuning with 3000 examples: 61.0%
LoRA r= 1 with 3000 examples: 25.1%
LoRA r= 2 with 3000 examples: 34.2%
LoRA r= 4 with 3000 examples: 57.3%
LoRA r= 8 with 3000 examples: 63.9%
LoRA r=16 with 3000 examples: 72.1%

Punchline: with only a few hundred finetuning examples, LoRA at rank \(16\) (about \(2\%\) of full finetuning’s parameter count) matches or beats full finetuning outright. The frozen weights do nothing clever here; a rank-\(16\) adapter simply has far less room to overfit a small finetuning set than the whole network does.