The frozen checkpoint problem
Real robot policies ship as frozen checkpoints. pi0, GR00T, SmolVLA — gigabytes of weights someone else pretrained on data you will never have. You want to teach one a new skill. You have two honest options. Full fine-tuning: unfreeze every weight and train. It works, but you now own a full-size copy of the model per skill, and — as we will measure — it quietly overwrites what the model knew. Or you freeze the whole thing and bolt on a tiny trainable adapter. That second option, in its most-used form, is LoRA — Low-Rank Adaptation — and it is about forty lines of code. This chapter builds it from scratch and measures what it actually buys.
The subject is deliberately small and state-based — the ch3.8 routing task, no rendering, so the
whole run is bitwise-reproducible on a CPU in seconds. An instruction token selects a skill, and
each action dimension reads a distinct state coordinate the skill picks out. Open lora.py. Six
regions: setup, model, lora, pretrain, sweep, report.
The base we will freeze
# The base policy we will FREEZE and adapt. It is a compact, state-based conditioned MLP# (the ch3.8 routing task, no attention needed here): an instruction token picks a "skill",# we concatenate its embedding with the raw state, and a small MLP maps that to an action.# `head` — the action head's OUTPUT projection — is what LoRA will wrap. The trunk (fc1, fc2)# stays frozen too; the adapter must reconstruct the held-out skill from the features the# trunk already computes, which is exactly why rank matters.class TinyPolicy(nn.Module): def __init__(self, dim: int, hidden: int) -> None: super().__init__() self.tok_embed = nn.Embedding(NUM_SKILLS, dim) # skill id -> a conditioning vector self.fc1 = nn.Linear(dim + STATE_DIM, hidden) # action-head trunk, layer 1 self.fc2 = nn.Linear(hidden, hidden) # action-head trunk, layer 2 self.head = nn.Linear(hidden, ACT_DIM) # action-head readout (LoRA-wrapped) def forward(self, skill: torch.Tensor, state: torch.Tensor) -> torch.Tensor: cond = torch.cat([self.tok_embed(skill), state], dim=-1) # fuse instruction + state h = F.silu(self.fc1(cond)) h = F.silu(self.fc2(h)) return self.head(h) def make_batch(gen: torch.Generator, n: int, skills: torch.Tensor): """Synthetic routing task (ch3.8, widened): each example draws a skill from `skills`; each of the ACT_DIM action dims reads a DISTINCT coordinate the skill selects (coordinate (skill + j) mod STATE_DIM, with an alternating sign). No noise — the rule is deterministic, so a good fit is a real R^2 near 1. Pretraining on skills {0,1,2} never touches skill 3's coordinates, so the held-out skill is genuinely unseen (poor zero-shot).""" idx = torch.randint(0, len(skills), (n,), generator=gen) skill = skills[idx] state = torch.randn(n, STATE_DIM, generator=gen) rows = torch.arange(n) cols = (skill[:, None] + torch.arange(ACT_DIM)) % STATE_DIM # (n, ACT_DIM) selected coordinates signs = torch.where(torch.arange(ACT_DIM) % 2 == 0, 1.0, -1.0) action = state[rows[:, None], cols] * signs # the deterministic action rule return skill.to(device), state.to(device), action.to(device)TinyPolicy is a compact conditioned MLP: embed the skill token, concatenate it with the raw state,
and map that through a two-layer action head to an action. Nothing exotic — the point is what we do
to it, not what it is. We pretrain it on three skills and hold one out. The held-out skill's
token embedding never receives a gradient and the trunk never sees its coordinates, so the reloaded
base genuinely cannot do it. head — the action head's output projection — is the layer LoRA will
wrap.
The base here is ~78K parameters, not ch3.8's ~7K. That is deliberate: LoRA's real-world hook is "you train ~1% of the weights," and that fraction is only meaningful if the base is wide enough that a rank-r adapter is genuinely ~1% of it. It is still a tiny, CPU-instant, bitwise MLP.
LoRA, in forty lines
# LoRA, from scratch. Wrap a FROZEN nn.Linear with a thin trainable low-rank bypass:# y = W x + (alpha / r) * B (A x), A: (r, in) B: (out, r)# W (and its bias) never train — only A and B do, r*(in+out) numbers instead of in*out.# The init is the whole trick (this mirrors peft's reset_lora_parameters): A gets the usual# nn.Linear kaiming init, B is ZEROED — so B(A x) = 0 at step 0 and the wrapped layer is# BITWISE the frozen one. Training then grows the bypass FROM the pretrained model, not from# noise. --break rand_init_B zeroes nothing: B starts random and the frozen prior is already# corrupted before the first gradient step (the deliberate failure the exercise measures).class LoRALinear(nn.Module): def __init__(self, base: nn.Linear, r: int, alpha: float, rand_init_B: bool = False) -> None: super().__init__() self.base = base # the frozen pretrained projection self.base.weight.requires_grad_(False) if self.base.bias is not None: self.base.bias.requires_grad_(False) self.r = r self.scaling = alpha / r if r > 0 else 0.0 if r > 0: self.A = nn.Parameter(torch.empty(r, base.in_features)) self.B = nn.Parameter(torch.empty(base.out_features, r)) nn.init.kaiming_uniform_(self.A, a=math.sqrt(5)) # A: the default nn.Linear init if rand_init_B: # BREAK: the natural mistake — init B like any nn.init.kaiming_uniform_(self.B, a=math.sqrt(5)) # other Linear instead of zeroing it. Now B(A x) else: # != 0 at step 0: the adapter is NOT a no-op and the nn.init.zeros_(self.B) # frozen policy is perturbed before any training. # zero-init B: B(A x) = 0 => adapted == frozen at step 0 def forward(self, x: torch.Tensor) -> torch.Tensor: y = self.base(x) if self.r > 0: # W x + (alpha/r) * B(A x) y = y + self.scaling * F.linear(F.linear(x, self.A), self.B) return y def build_arm(state_dict: dict, arm: str, rank: int) -> nn.Module: """Load the frozen pretrained weights into a fresh policy, then configure ONE arm: * "frozen"/"lora" with rank 0 -> nothing trains (the zero-shot base), * "full" -> every weight trains (full fine-tuning), * "lora" with rank > 0 -> W frozen, only the head's A/B adapter trains.""" model = TinyPolicy(args.dim, args.hidden).to(device) model.load_state_dict(state_dict) if arm == "full": return model # all params keep requires_grad=True for p in model.parameters(): # freeze the whole base p.requires_grad_(False) if arm == "lora" and rank > 0: # wrap the action head's output projection model.head = LoRALinear(model.head, rank, args.alpha, RAND_B) return model.to(device) # the fresh A/B adapter follows the base's deviceHere is the entire idea. Freeze a linear layer y = W x. Add a thin trainable bypass made of two
skinny matrices — A projects the input down to a tiny rank r, B projects it back up:
y = W x + (alpha / r) · B (A x)
W never trains. A is (r, in) and B is (out, r), so together they hold r·(in + out) numbers
instead of in·out — for our head, a rank-4 adapter is 4·(256 + 6) = 1048 numbers against 77,958,
about 1.3%. The scaling alpha/r is peft's convention, and --break will show you why the last
detail matters more than any of them: B is initialized to zero. So B(A x) = 0 at step 0, and
the adapted layer is bitwise identical to the frozen one. You begin fine-tuning from the exact model
you paid to pretrain, and the adapter grows out of it rather than shoving it around.
Pretrain, freeze, adapt three ways
# Pretrain the base on skills {0,1,2}, HOLDING OUT skill 3. Skill 3's token embedding never# gets a gradient and the trunk never sees its coordinates, so the reloaded base cannot do it# zero-shot — that gap is what the three arms then try to close.def r2(pred: np.ndarray, target: np.ndarray) -> float: resid = float(((target - pred) ** 2).sum()) total = float(((target - target.mean(0)) ** 2).sum()) return 1.0 - resid / total if total > 0 else 0.0 @torch.no_grad()def evaluate(model: nn.Module, skill: torch.Tensor, state: torch.Tensor, action: torch.Tensor): model.eval() pred = model(skill, state).cpu().double().numpy() tgt = action.cpu().double().numpy() return r2(pred, tgt), float(((tgt - pred) ** 2).mean()) def fit(model: nn.Module, skill, state, action, steps: int, tag: str) -> None: params = [p for p in model.parameters() if p.requires_grad] if not params: # r=0 / frozen: nothing to train return opt = torch.optim.Adam(params, lr=args.lr) for step in range(steps): loss = ((model(skill, state) - action) ** 2).mean() opt.zero_grad() loss.backward() opt.step() if args.rerun: rr.set_time("step", sequence=step) rr.log(f"train/{tag}/mse", rr.Scalars([loss.item()])) pretrain_gen = torch.Generator().manual_seed(args.seed) # the pretraining pile (skills 0..2)train_skills = torch.tensor([s for s in range(NUM_SKILLS) if s != HELD_OUT])pt_skill, pt_state, pt_action = make_batch(pretrain_gen, args.train_examples, train_skills) ckpt_path = args.out / "base_policy.pt"torch.manual_seed(args.seed) # reproducible initbase = TinyPolicy(args.dim, args.hidden).to(device)fit(base, pt_skill, pt_state, pt_action, args.pretrain_steps, "pretrain")torch.save(base.state_dict(), ckpt_path) # "ship" the checkpoint...frozen_state = torch.load(ckpt_path, weights_only=True) # ...and reload it, as any released policybase_total = sum(p.numel() for p in base.parameters())print(f"pretrained + reloaded base ({base_total:,} params); held out skill {HELD_OUT}") # The two held-out eval piles, fixed across every arm: the UNSEEN skill (the target of# adaptation) and an in-distribution skill TASK_A (the retention / forgetting probe).eval_gen = torch.Generator().manual_seed(args.seed + 100)ho_skill, ho_state, ho_action = make_batch(eval_gen, args.eval_examples, torch.tensor([HELD_OUT]))ta_skill, ta_state, ta_action = make_batch(eval_gen, args.eval_examples, torch.tensor([TASK_A]))# The adaptation pile: skill 3 ONLY (what a practitioner has for the new skill).adapt_gen = torch.Generator().manual_seed(args.seed + 1)ad_skill, ad_state, ad_action = make_batch(adapt_gen, args.adapt_examples, torch.tensor([HELD_OUT]))We pretrain on skills {0, 1, 2}, torch.save the weights, and reload them from disk — the move
you make with any released checkpoint you did not train. Then, from that one frozen state, we build
three arms and adapt each to the held-out skill: frozen (nothing trains — the zero-shot baseline),
full-FT (every weight trains), and LoRA (W frozen, only the head's A/B train). Each arm
prints its trainable-vs-total parameter count, which is the whole economic argument in one line.
The rank dial: the elbow
# One rank dial, three regimes. For each rank we rebuild the arm from the SAME frozen weights,# adapt on skill-3 data, and read two fits: the held-out skill (did we LEARN it?) and TASK_A# (did we FORGET an old one?). r=0 is the frozen zero-shot; "full" is full fine-tuning. The# elbow lives in held_out_r2 vs rank; the (honest) forgetting lives in task_a_r2.def trainable_count(model: nn.Module) -> int: return sum(p.numel() for p in model.parameters() if p.requires_grad) def run_arm(arm: str, rank: int) -> dict: torch.manual_seed(args.seed + 7) # same adapter/full init draw across ranks model = build_arm(frozen_state, arm, rank) trainable = trainable_count(model) # Step 0, BEFORE any update: for zero-init-B LoRA this is bitwise the frozen policy; the # break (and full-FT) moves it. We record how far step-0 already is from the frozen output. with torch.no_grad(): step0 = model(ho_skill, ho_state).cpu().double().numpy() frozen_out = base(ho_skill, ho_state).cpu().double().numpy() step0_gap = float(np.abs(step0 - frozen_out).max()) step0_r2 = r2(step0, ho_action.cpu().double().numpy()) fit(model, ad_skill, ad_state, ad_action, args.adapt_steps, f"{arm}{rank}") ho_r2, ho_mse = evaluate(model, ho_skill, ho_state, ho_action) ta_r2, ta_mse = evaluate(model, ta_skill, ta_state, ta_action) return {"arm": arm, "rank": rank, "trainable": trainable, "trainable_pct": round(100.0 * trainable / base_total, 4), "held_out_r2": round(ho_r2, 6), "held_out_mse": round(ho_mse, 6), "task_a_r2": round(ta_r2, 6), "task_a_mse": round(ta_mse, 6), "step0_heldout_r2": round(step0_r2, 6), "step0_frozen_gap": round(step0_gap, 6)} sweep = [run_arm("lora", r) for r in SWEEP_RANKS] + [run_arm("full", args.hidden)]frozen = next(r for r in sweep if r["arm"] == "lora" and r["rank"] == 0)lora = next(r for r in sweep if r["arm"] == "lora" and r["rank"] == HEADLINE_RANK)full = next(r for r in sweep if r["arm"] == "full")print("\nrank dial (adapt the held-out skill; held_out fit rises then plateaus; watch task_A):")print(f" {'arm':6s} {'rank':>5s} {'trainable':>10s} {'%':>7s} {'held_out R2':>12s} {'task_A R2':>11s}")for row in sweep: name = "frozen" if (row["arm"] == "lora" and row["rank"] == 0) else row["arm"] rk = "full" if row["arm"] == "full" else str(row["rank"]) print(f" {name:6s} {rk:>5s} {row['trainable']:>10,d} {row['trainable_pct']:>6.2f}% " f"{row['held_out_r2']:>12.3f} {row['task_a_r2']:>11.3f}")Now turn the dial. We sweep the LoRA rank r and read the held-out fit at each setting. The measured
result, seed 0 (and the shape holds on every seed):
| rank | % of weights trained | held-out fit (R²) |
|---|---|---|
| 0 (frozen) | 0.00% | −0.44 |
| 1 | 0.34% | −0.08 |
| 2 | 0.67% | 0.25 |
| 4 | 1.34% | 0.77 |
| 8 | 2.69% | 1.00 |
| 16 | 5.38% | 1.00 |
| full fine-tune | 100% | 1.00 |
Read the shape, not the decimals. The frozen base cannot do the held-out skill (a negative R² is worse than guessing the mean). Then the fit rises with rank and plateaus onto the full-fine-tune ceiling. A rank-4 adapter, training ~1.3% of the weights, recovers ~84% of full fine-tuning's fit; by rank 8 the curve sits on the full-FT line. Past that knee, more trainable parameters buy almost nothing. That knee is the elbow, and it kills a tempting intuition: fewer trainable parameters must mean a worse fit. They don't. Most of what full fine-tuning learns here lives in a handful of directions, and a low-rank bypass finds them.
The honest twist: freezing W is not freezing behavior
There is a second intuition, just as tempting, and it is wrong — so we measure it rather than
assert it. Surely, the story goes, LoRA can't forget the skills the base already knew: its W is
frozen, so the old mapping is untouched. We watch an in-distribution skill (task_A, one the base was
pretrained on) while we adapt to the held-out one:
# The headline, stated as the reproducible DIRECTION (mechanism, not a third decimal):# * held_out fit: full-FT is the ceiling; a small-rank LoRA recovers MOST of it (the elbow),# training ~1% of the weights — "fewer trainable params" does NOT mean "worse fit".# * forgetting (the honest twist): freezing W does NOT protect TASK_A. Its fit collapses under# LoRA as it does under full-FT, because the additive low-rank adapter fires on EVERY input# and cannot gate itself off for the old skill. LoRA buys parameter efficiency, not memory.recovered = (lora["held_out_r2"] - frozen["held_out_r2"]) / max( 1e-9, full["held_out_r2"] - frozen["held_out_r2"]) # fraction of full-FT's held-out gain LoRA recoversmetrics = { "alpha": args.alpha, "base_total_params": int(base_total), "break_mode": args.break_mode or "none", "frozen_heldout_r2": frozen["held_out_r2"], # zero-shot on the unseen skill: poor "frozen_task_a_r2": frozen["task_a_r2"], # the pretrained retention baseline: high "full_heldout_r2": full["held_out_r2"], # full-FT ceiling on the held-out skill "full_task_a_r2": full["task_a_r2"], "full_task_a_forget": round(frozen["task_a_r2"] - full["task_a_r2"], 6), # >0: full-FT forgot TASK_A "full_trainable_pct": full["trainable_pct"], "headline_rank": HEADLINE_RANK, "held_out_skill": HELD_OUT, "hidden": args.hidden, "lora_heldout_r2": lora["held_out_r2"], # small-rank LoRA held-out fit "lora_recovered_frac": round(recovered, 6), # fraction of full-FT's held-out gain LoRA recovers (headline) "lora_step0_frozen_gap": lora["step0_frozen_gap"], # 0.0 clean (adapted==frozen@step0); >0 under --break "lora_step0_heldout_r2": lora["step0_heldout_r2"], # == frozen_heldout_r2 clean; worse under --break "lora_task_a_r2": lora["task_a_r2"], "lora_task_a_forget": round(frozen["task_a_r2"] - lora["task_a_r2"], 6), # >0: LoRA ALSO forgot TASK_A "lora_trainable_pct": lora["trainable_pct"], "seed": args.seed, "smoke": bool(args.smoke), "sweep_heldout_r2": [row["held_out_r2"] for row in sweep], "sweep_ranks": [("full" if row["arm"] == "full" else row["rank"]) for row in sweep], "sweep_task_a_r2": [row["task_a_r2"] for row in sweep], "sweep_trainable_pct": [row["trainable_pct"] for row in sweep],}(args.out / "metrics.json").write_text(json.dumps(metrics, indent=2, sort_keys=True) + "\n") # demo/vizdata.json: the rank dial. Two synced readouts per rank — params-trained %# (climbing) and held-out fit (rising then PLATEAUING, the elbow) — plus the TASK_A trace# (collapsing under adaptation) and the full-FT reference line.viz_rows = [r for r in sweep if r["arm"] == "lora"] # the LoRA rank dial (r=0..16)vizdata = { "held_out_skill": HELD_OUT, "task_a_skill": TASK_A, "headline_rank": HEADLINE_RANK, "base_total_params": int(base_total), "break_mode": args.break_mode or "none", "ranks": [r["rank"] for r in viz_rows], "trainable_pct": [r["trainable_pct"] for r in viz_rows], "held_out_r2": [r["held_out_r2"] for r in viz_rows], "task_a_r2": [r["task_a_r2"] for r in viz_rows], "full": {"held_out_r2": full["held_out_r2"], "task_a_r2": full["task_a_r2"], "trainable_pct": full["trainable_pct"]}, "frozen": {"held_out_r2": frozen["held_out_r2"], "task_a_r2": frozen["task_a_r2"]},}(args.out / "demo").mkdir(parents=True, exist_ok=True)(args.out / "demo" / "vizdata.json").write_text(json.dumps(vizdata) + "\n") if args.rerun: rr.log("sweep/held_out_r2", rr.BarChart(np.array([r["held_out_r2"] for r in viz_rows]))) rr.log("sweep/trainable_pct", rr.BarChart(np.array([r["trainable_pct"] for r in viz_rows]))) rr.log("forgetting/task_a_r2", rr.BarChart(np.array([frozen["task_a_r2"], lora["task_a_r2"], full["task_a_r2"]]))) print(f"\nheadline: rank-{HEADLINE_RANK} LoRA trains {lora['trainable_pct']:.2f}% of the weights and recovers " f"{recovered * 100:.0f}% of full-FT's held-out gain ({frozen['held_out_r2']:.2f} -> " f"{lora['held_out_r2']:.2f} vs full {full['held_out_r2']:.2f}).")print(f"forgetting (the honest twist): task_A R^2 frozen {frozen['task_a_r2']:.2f} -> LoRA " f"{lora['task_a_r2']:.2f} vs full-FT {full['task_a_r2']:.2f}. Freezing W did NOT protect the old skill.")print(f"step 0: LoRA output is {lora['step0_frozen_gap']:.2e} from the frozen policy " f"[{'zero-init B: adapted == frozen' if not RAND_B else 'BREAK rand_init_B: prior already corrupted'}].")print(f"metrics: {args.out / 'metrics.json'} | vizdata: {args.out / 'demo' / 'vizdata.json'}")if args.rerun: print(f"recording: {args.out / 'lora.rrd'} — open it with: rerun {args.out / 'lora.rrd'}")task_A's fit, seed 0: frozen base +1.00 → LoRA −1.88 → full-FT −0.95. It collapses —
under LoRA as badly as under full fine-tuning, here worse. Freezing W did not protect it, and the
reason is right there in the update rule. The adapter (alpha/r)·B(A x) is added to every input,
including task_A's, and a single low-rank linear map cannot switch itself off for one skill and on
for another — that would take a multiplication it doesn't have. So the correction we trained for the
new skill bleeds straight onto the old one. LoRA's real win here is parameter efficiency, not free
memory. "Frozen weights" is not "frozen behavior." (Real systems do see LoRA forget less than full
fine-tuning — but that comes from scale and regularization, not from the frozen-W intuition, and in
this clean toy full fine-tuning, which can localize its update to the new skill's own weights, often
protects the old skill better.)
Break it: why B starts at zero
You wrote the LoRA update in the exercise; the single most natural way to get it wrong is to
initialize B the way you'd init any other linear layer instead of zeroing it. Predict what that does,
then run --break rand_init_B.
The adapted policy is no longer identical to the frozen one at step 0 — the step-0 gap jumps from
exactly 0.0 to ~0.17 (seed-robust). You have perturbed the model you paid to pretrain before a
single gradient step. Here adaptation still recovers the held-out fit, so the outcome doesn't
collapse — the lesson is subtler and more important: zero-init B is what makes LoRA a no-op at
step 0, so you begin adaptation from the pretrained model itself rather than from a randomly
jostled copy of it. That is the whole reason the real implementation zeroes it.
Read the real thing
Everything here has a production form, and it is remarkably close. The read-the-real-thing segment
pairs this chapter with huggingface/peft, src/peft/tuners/lora/layer.py, class Linear. You will
recognize every piece: update_layer builds lora_A = nn.Linear(in, r) and lora_B = nn.Linear(r, out)
— our A and B; reset_lora_parameters kaiming-inits lora_A and zeroes lora_B — our exact
init and the thing --break violates; scaling = lora_alpha / r — our alpha/r; and the forward is
result + lora_B(lora_A(x)) * scaling — our W x + (alpha/r)·B(A x). The same forty lines wrap the
attention and MLP projections of a billion-parameter VLA. The mechanism is size-invariant; that is the
point of building it small.
What's next
You built the single most-used fine-tuning primitive in the modern stack from scratch, and measured it
honestly: a low-rank bypass recovers most of full fine-tuning's fit for ~1% of the trainable weights
(the elbow) — and does not hand you free memory (freezing W did not stop the old skill from being
forgotten). That is the real practitioner's calculus: LoRA is how you adapt a frozen policy cheaply,
and knowing exactly what it does and does not protect is how you use it without fooling yourself.