Why is the B matrix initialized to zero in LoRA?

So the update starts at exactly zero and the model begins training from the base model's exact behavior. With A random and B zero, the first step sees the original model.

Implementing LoRA from Scratch in PyTorch

Q: What does LoRA actually train?

Two small matrices per wrapped layer, A and B, whose product is a low-rank update added beside the frozen weight. The original weight never changes during training.

What LoRA is in practice

LoRA freezes a layer's weight matrix W and learns a low-rank update beside it. Instead of changing W, it adds B·A, where A and B are small matrices with an inner dimension r far smaller than W, scaled by α/r. Only A and B are trained. A is initialized at random and B at zero, so at the start the update is exactly zero and the model behaves like the frozen base; training moves it from there. That is the whole idea, and it fits in a small module.

Starting point: the base model

We need a trained model to adapt. We use the character-level JSON transformer we presented in this cookbook, condensed so the notebook runs on its own. Two dialects are defined here. Dialect A is the original. Dialect B is the same, plus one new field, code, whose value is two uppercase letters and a hyphen followed by digits. The uppercase letters and the hyphen are the characters that appear only in B. We build the vocabulary from both dialects up front, so the embedding and output head already have rows for those characters, even though the base model trains only on A and never produces them.

import json
import random
import torch
import numpy as np 

np.random.seed(3)

KEYS    = ["id", "user", "limit", "query", "active"]
METHODS = ["get_weather", "query_db", "send_email"]
LEAVES  = [lambda: random.randint(0, 999),
           lambda: random.choice(["alice", "bob", "ok"]),
           lambda: random.choice([True, False])]

def params_obj():
    keys = random.sample(KEYS, random.randint(1, 3))
    return {k: random.choice(LEAVES)() for k in keys}

def record_A():
    return {"method": random.choice(METHODS), "params": params_obj()}

UPPER = "ABCDEFGHJKMNPQRSTVWXYZ"
def record_B():
    r = record_A()
    # new field: uppercase letters and a hyphen
    r["code"] = "".join(random.choice(UPPER) for _ in range(2)) + "-" + str(random.randint(10, 99))
    return r                                   

textA = "\n".join(json.dumps(record_A()) for _ in range(8000))
textB = "\n".join(json.dumps(record_B()) for _ in range(8000))

# vocabulary from BOTH dialects, so the head has rows for the new characters
chars = sorted(set(textA + textB)); vocab_size = len(chars)
stoi  = {c: i for i, c in enumerate(chars)}
itos  = {i: c for i, c in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
# base trains on dialect A
dataA = torch.tensor(encode(textA), dtype=torch.long)
# LoRA adapts on dialect B   
dataB = torch.tensor(encode(textB), dtype=torch.long)   

def code_rate(text):
    # fraction of generated lines that include the new field                           
    lines = [l for l in text.split("\n") if l.strip()]
    ok = 0
    for l in lines:
        try:
            ok += "code" in json.loads(l)
        except json.JSONDecodeError:
            pass
    return ok / max(len(lines), 1)

The model and a small training loop, unchanged from the previous cookbook.

import torch.nn as nn
from torch.nn import functional as F

block_size, n_embd, n_head, n_layer = 128, 128, 4, 3
device = "cuda" if torch.cuda.is_available() else "cpu"

class Head(nn.Module):
    def __init__(self, head_size):
        super().__init__()
        self.key   = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        k, q, v = self.key(x), self.query(x), self.value(x)
        wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5
        wei = wei.masked_fill(self.tril[:x.size(1), :x.size(1)] == 0, float("-inf"))
        return F.softmax(wei, dim=-1) @ v

class MultiHead(nn.Module):
    def __init__(self, n_head, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(n_head)])
        self.proj  = nn.Linear(n_head * head_size, n_embd)
    def forward(self, x):
        return self.proj(torch.cat([h(x) for h in self.heads], dim=-1))

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.sa  = MultiHead(n_head, n_embd // n_head)
        self.ff  = nn.Sequential(nn.Linear(n_embd, 4 * n_embd), nn.ReLU(),
                                 nn.Linear(4 * n_embd, n_embd))
        self.ln1, self.ln2 = nn.LayerNorm(n_embd), nn.LayerNorm(n_embd)
    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        return x + self.ff(self.ln2(x))

class JSONTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb   = nn.Embedding(block_size, n_embd)
        self.blocks  = nn.Sequential(*[Block() for _ in range(n_layer)])
        self.ln_f    = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)
    def forward(self, idx, targets=None):
        T = idx.size(1)
        x = self.token_emb(idx) + self.pos_emb(torch.arange(T, device=idx.device))
        logits = self.lm_head(self.ln_f(self.blocks(x)))
        loss = None if targets is None else F.cross_entropy(
            logits.view(-1, vocab_size), targets.view(-1))
        return logits, loss

@torch.no_grad()
def sample(model, n=4000):
    model.eval()
    idx = torch.tensor([[stoi["\n"]]], device=device)
    for _ in range(n):
        logits, _ = model(idx[:, -block_size:])
        nxt = torch.multinomial(F.softmax(logits[:, -1], dim=-1), 1)
        idx = torch.cat([idx, nxt], dim=1)
    model.train()
    return "".join(itos[i] for i in idx[0].tolist())

def train(model, data, steps, lr=3e-4, batch_size=64):
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr)
    for _ in range(steps):
        ix = torch.randint(len(data) - block_size, (batch_size,))
        xb = torch.stack([data[i:i + block_size]         for i in ix]).to(device)
        yb = torch.stack([data[i + 1:i + block_size + 1] for i in ix]).to(device)
        _, loss = model(xb, yb)
        opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
    return model

Train the base on dialect A. It learns valid JSON but never the code field, since it never saw one.

base = JSONTransformer().to(device)
train(base, dataA, steps=3000)

print("base code-rate:", round(code_rate(sample(base, 4000)), 3))   # expect near 0

The base code-rate in this example is 0 which makes sense given the training data.

LoRA from scratch

Here is the whole adapter. It wraps a linear layer, keeps the original frozen, and adds the low-rank path. delta returns the effective weight change, B·A scaled by α/r, which we use later to see what moved.

import math

class LoRALinear(nn.Module):
    def __init__(self, base, r=8, alpha=16):
        super().__init__()
        # the frozen pretrained linear
        self.base = base                                  
        out_features, in_features = base.weight.shape
        self.A = nn.Parameter(torch.empty(r, in_features))
        self.B = nn.Parameter(torch.zeros(out_features, r))   # B = 0 -> adapter starts as a no-op
        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))      # A starts small and random
        self.scale = alpha / r

    def forward(self, x):
        return self.base(x) + (x @ self.A.T @ self.B.T) * self.scale

    def delta(self):
        return (self.B @ self.A) * self.scale             # effective weight change, shape of base.weight

Wrap the model and freeze it

We wrap the attention query and value projections and the output head, freeze every parameter, then re-enable only the adapter matrices. After this, the only trainable parameters are the LoRA ones. Which layers you target is a choice; the query and value projections are a common one.

def add_lora(module, targets=("query", "value", "lm_head"), r=8, alpha=16):
    for name, child in list(module.named_children()):
        if isinstance(child, nn.Linear) and name in targets:
            setattr(module, name, LoRALinear(child, r, alpha))
        else:
            add_lora(child, targets, r, alpha)

add_lora(base)                                  # wrap attention query/value and the output head

for p in base.parameters():
    p.requires_grad_(False)                     # freeze everything
for m in base.modules():
    if isinstance(m, LoRALinear):
        m.A.requires_grad_(True); m.B.requires_grad_(True)   # train only the adapters

trainable = sum(p.numel() for p in base.parameters() if p.requires_grad)
total     = sum(p.numel() for p in base.parameters())
print(f"trainable {trainable} of {total}  ({trainable / total:.2%})")

That comes out to 32256 (4.89%) of the model's parameters.

Adapt to the new dialect

Train on dialect B. Because everything except the adapters is frozen, only A and B move and the base weights stay exactly as they were. The new behavior appears.

train(base, dataB, steps=2000)                  # only the LoRA parameters move

print("adapted code-rate:", round(code_rate(sample(base, 4000)), 3))   # expect high

The code-rate goes from 0% to 72.7%. The model now produces the new field, and it learned to do so by moving a small set of adapter parameters while the original weights stayed frozen.

Which characters the head learned

This is the part that connects the weights to the task. The output head maps a hidden vector to one logit per character, so its weight has one row per character, and so does the adapter's update. The size of each row tells us how much the adapter can move that character's logit. Rank the characters by it.

head = base.lm_head   # LoRALinear wrapping the output head
dW   = head.delta()   # (vocab, n_embd): one row per character
row  = dW.norm(dim=1) # how much the update can move each character's logit

print("characters the head update changed most:")
for i in row.argsort(descending=True)[:12].tolist():
    print(f"  {itos[i]!r:6}  {row[i]:.3f}")

Notice the hyphen is the character with the biggest change, which is expected as it was introduced in dialect B, we also note some of the capital letters being changed signifcantly.

characters the head update changed most:
'-'     1.808
's'     0.972
','     0.966
'c'     0.951
'g'     0.910
'e'     0.899
'd'     0.854
'_'     0.852
'R'     0.805
'N'     0.804
'w'     0.802
't'     0.797

How many directions the update used

We set the rank to 8, but the update does not have to use all of it. Take the singular values of one layer's update and see how fast they fall off.

dW = base.blocks[1].sa.heads[0].value.delta()   # a wrapped value projection (rank <= r)
s  = torch.linalg.svdvals(dW)
print("singular values:", [round(float(x), 3) for x in s[:8]])

Notice the energy concentrated in the first few values, with the rest small: the change the task needed was lower-rank than the budget we gave it. On a model this small the spectrum is noisy, so read the trend rather than the exact numbers. singular values: [1.14, 0.487, 0.363, 0.091, 0.074, 0.065, 0.05, 0.028]

Where the update landed

The relative size of each layer's update, |ΔW| / |W|, shows where the adaptation concentrated.

for name, m in base.named_modules():
    if isinstance(m, LoRALinear):
        d = m.delta()
        print(f"{name:34}  |dW|/|W| = {(d.norm() / m.base.weight.norm()):.4f}")

Compare the values to see which projections and which layers did the most work. At this scale, read the ranking as a rough guide rather than a precise measurement.

Merge the adapter

When fine tuning is done, you can fold the update back into the weight. Since B·A scaled is just a matrix the same shape as W, W + ΔW is an ordinary weight, and the wrapped layer produces the same outputs with the adapter zeroed. Merging removes the inference overhead and lets you ship a plain model.

@torch.no_grad()
def merge(model):
    for m in model.modules():
        if isinstance(m, LoRALinear):
            m.base.weight += m.delta()      # fold the update into the frozen weight
            m.B.zero_()                     # adapter is now a no-op; outputs are unchanged

xb = torch.tensor([encode(textB[:block_size])], device=device)
before, _ = base(xb)
merge(base)
after, _ = base(xb)
print("max logit difference after merge:", (before - after).abs().max().item())   # ~0

The maximum change in the logits after merging is about 1.38e-05, which is extremely small. The merged model is very close to the original model.

Common questions

What does LoRA actually train?

Two small matrices per wrapped layer, A and B, whose product is a low-rank update added beside the frozen weight. The original weight never changes during training; when you are finished you can fold the update back into it.

Why is B initialized to zero?

So the update starts at exactly zero and the model begins training from the base model's behavior rather than from a randomly perturbed version of it. With A random and B zero, the first step sees the original model; the adapter moves away from there.

Does a smaller rank lose information?

Not necessarily. The update often uses fewer directions than the rank you set, which is why small ranks are frequently enough. The right rank depends on how complex the change is: a focused behavior shift needs little, a broad change needs more.

Implementing LoRA from Scratch on a Toy Model

Quick Answer / TL;DR