LM6-T3 — Frozen and Trainable Parameters

Learning objectives

After this lesson, you should be able to:

• Describe head-only, partial and full fine-tuning and what each trains.
• Freeze and unfreeze parameters in PyTorch and count the trainable ones.
• Estimate the memory for weights, gradients and optimizer state under each strategy.
• Choose a strategy from data size, task difficulty and budget.

The idea: spend the update budget where it counts

A pretrained backbone computes a representation of the text. The classification head reads that representation. If the representation already separates your classes, only the head needs to learn, and training a few hundred parameters on a small dataset is fast and hard to overfit. If the classes depend on features the backbone does not expose well, the head alone is not enough and some of the backbone must change.

The options

• Head only (sometimes called linear probing): the backbone is frozen. Cheapest, safest against forgetting, and fastest because the frozen backbone needs no gradient. Quality depends entirely on the quality of the pretrained representation for your task. On hard tasks it can be far behind; on the toy task of the notebook it fails completely.
• Partial fine-tuning: train the head and the last one or few blocks and the final norm. The upper layers are the most task-specific, so adapting them costs little and often recovers most of the quality of full fine-tuning.
• Full fine-tuning: train everything. Most flexible and most expensive; needs the most care with the learning rate and data size.
• Parameter-efficient methods (adapters, low-rank updates): keep the weights frozen and train small added modules. Quality close to full fine-tuning at a fraction of the memory (Module 8).

How freezing works

Every parameter has a flag requires_grad. When it is False, autograd does not compute a gradient for it, the optimizer never sees it, and it holds no optimizer state. If the whole backbone is frozen, backpropagation stops at the head and the backbone only runs forward. Freezing is therefore a saving of memory and compute as well as a constraint.

Learning rates by group

A randomly initialised head needs a larger learning rate (for example 10⁻³) than pretrained layers (10⁻⁵ to 5·10⁻⁵). Optimizers accept parameter groups with their own rates. Some recipes decay the rate layer by layer from top to bottom, so that lower layers, which hold more general features, change least.

Caching features

If the backbone is frozen and dropout is off, its output for a given text never changes. You can run the backbone once over the dataset, store the pooled vectors, and train the head on those vectors for many epochs at almost no cost. This makes head-only experiments, hyperparameter sweeps and error analysis very fast.

The mathematics: what each strategy costs

Let P be the total number of parameters and P_train the number that are trainable. In float32 with Adam, the memory for weights, gradients and optimizer state is

4·P

the weights, 4 bytes each; all parameters, frozen or not, are stored.

4·P_train

gradients, only for trainable parameters.

8·P_train

Adam’s two moving averages (4 bytes each), only for trainable parameters.

activations

not included: they depend on batch size and sequence length. With a frozen backbone and no gradient through it they are not stored at all.

For GPT-2 small (P = 124,412,160, a tied head) with a 2-class head:

The last block has 7,085,568 parameters (Lesson LM4-T5); adding the final norm (1,536) and the head (1,538) gives 7,088,642. Head-only training needs little more than the model itself, which is why it fits on any machine that can run the model.

Compute. The forward pass costs about 2·P FLOPs per token whatever you train. The backward pass costs about 4·P per token if gradients flow through the whole network, about 4·P_below less if the lower part is frozen and cut out of the backward pass. With only the head trainable no backward pass through the backbone is needed, so a training step costs about a third of a full step.

Capacity versus data. The more parameters you train, the more examples you need to avoid overfitting. A rough guide: head only works with tens to hundreds of examples, partial fine-tuning with hundreds to thousands, full fine-tuning with thousands and more, but task difficulty changes all three.

A worked example: freezing and counting in PyTorch

The code builds a small stand-in backbone (two blocks as linear layers) and a head, then applies the three strategies and counts the trainable parameters and the optimizer memory. Parameter groups with different learning rates are shown at the end.

import torch
import torch.nn as nn

d, K = 64, 2
backbone = nn.ModuleDict({
    "embed": nn.Embedding(100, d),
    "block1": nn.Linear(d, d),
    "block2": nn.Linear(d, d),
    "norm": nn.LayerNorm(d),
})
head = nn.Linear(d, K)
model = nn.ModuleDict({"backbone": backbone, "head": head})

def set_trainable(model, mode):
    for p in model.parameters():
        p.requires_grad = False
    groups = [model["head"]]
    if mode in ("last", "all"):
        groups += [model["backbone"]["block2"], model["backbone"]["norm"]]
    if mode == "all":
        groups += [model["backbone"]["embed"], model["backbone"]["block1"]]
    for g in groups:
        for p in g.parameters():
            p.requires_grad = True
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

total = sum(p.numel() for p in model.parameters())
for mode in ("head", "last", "all"):
    n = set_trainable(model, mode)
    mb = (4 * total + 12 * n) / 1e6
    print(f"{mode:5s} trainable {n:6d} of {total}  ->  {mb:.2f} MB (weights + gradients + Adam)")
# head  trainable    130 of 14978 ...
# last  trainable   4418 of 14978 ...
# all   trainable  14978 of 14978 ...

# frozen parameters get no gradient
set_trainable(model, "head")
x = torch.randint(0, 100, (3,))
h = model["backbone"]["norm"](model["backbone"]["block2"](model["backbone"]["block1"](model["backbone"]["embed"](x))))
loss = model["head"](h).sum()
loss.backward()
print(model["backbone"]["block1"].weight.grad is None, model["head"].weight.grad is not None)   # True True

# parameter groups: a larger learning rate for the new head than for pretrained layers
set_trainable(model, "last")
opt = torch.optim.AdamW([
    {"params": model["head"].parameters(), "lr": 1e-3},
    {"params": list(model["backbone"]["block2"].parameters()) + list(model["backbone"]["norm"].parameters()), "lr": 3e-5},
], weight_decay=0.01)
print([g["lr"] for g in opt.param_groups])           # [0.001, 3e-05]

Only parameters with requires_grad = True are given to the optimizer here, so frozen weights hold no Adam state. If you pass all parameters, PyTorch skips those without gradients but you should still filter, so that nothing changes silently if a flag is flipped later.

Choosing

Engineering implications and limitations

Freeze before building the optimizer. Set requires_grad first, then create the optimizer from the parameters that require gradients. Changing the flags after the optimizer exists does not change its parameter list.

Keep frozen parts in eval mode. A frozen backbone with dropout still on gives noisy features and is not what you evaluate. Call backbone.eval() if it is fully frozen; layer norm and dropout behave differently in the two modes.

Use torch.no_grad() for the frozen part. Wrapping the backbone forward pass in no_grad avoids storing its activations at all, which is a large memory saving.

Verify what trains. Print the names of parameters with requires_grad and their count at the start of every run. Many silent failures come from layers that were meant to train and did not (or the reverse).

Weight tying and freezing. If the output head is tied to the token embedding, freezing one freezes the other. A classification head is new and separate, so this does not matter for classification; it matters in instruction tuning (Module 7).

Gradual unfreezing. Train the head first, then unfreeze one block at a time from the top while lowering the learning rate. It can stabilise training on small data at the cost of more phases.

Compare against the baseline. A logistic regression on cached pooled features is a head-only classifier; if full fine-tuning does not beat it by a meaningful margin on held-out data, the extra cost is not justified.

Common misconceptions

Engineering takeaways

1. Head only, partial and full fine-tuning train 0.001%, 5.7% and 100% of GPT-2 small’s parameters (1,538; 7,088,642; 124,412,160).
2. Frozen parameters are stored but have no gradient and no optimizer state: M ≈ 4P + 12·P_train bytes in float32 with Adam.
3. Set requires_grad before creating the optimizer, use parameter groups for different learning rates, and print what is trainable.
4. Cache the pooled features of a frozen backbone to experiment on the head almost for free.
5. More trainable parameters need more labelled data; parameter-efficient methods (Module 8) sit in between.

References and further reading

See Chapter 6 of the book (initializing a model with pretrained weights, and finetuning the model on supervised data, including the choice of layers to train).

• PyTorch documentation: torch.Tensor.requires_grad, torch.no_grad and the optimizer parameter groups (torch.optim).
• Build a Large Language Model (From Scratch), Chapter 6: Finetuning for Classification (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: preparing the labelled data and measuring the classifier with the right metrics.

Related content

Nothing below is required to finish this lesson.

• Before
  LM6-T2: Classification Heads and Sequence Representations

  Replace the vocabulary head with a small classification head and pick the token whose state represents the sequence.

• Lab
  LM6-L2: Fine-Tune a Language Model for Classification

  Add a classification head to a small GPT, compare frozen and fine-tuned variants and analyse the errors.