Fedor Pakhurov
RUGitHub ↗
Menu
Course contents 24 / 37

24 / 37 · Adaptation and Efficient Models

Fine-tuning

Continue training a pretrained model on a labelled task: what to train (everything, or only a new head), at what learning rate, and what it costs in memory. A GPT-2 model on sentiment gives 0.890 with everything trained, 0.737 with the head alone, and 0.900 when the head is trained first and everything second.

Prerequisites

Transfer learning in NLP · GPT and autoregressive language modeling · Gradient descent

On this page 12

The transfer learning note found that training all of a pretrained encoder on the task scored higher than reading off its frozen features on AG News: 0.949 on a held-out tenth of the training set against 0.902 on the test set. That raises the practical questions this note answers: which parameters to train, at what learning rate, in what order, and what the choice costs. The experiments move to a decoder-only model and a new task, so that the conclusions are not specific to BERT.

Notation. Pretrained parameters θ\theta (the body), a task head with parameters ϕ\phi initialized at random; learning rate η\eta; the loss L\mathcal L is the task cross-entropy.

The setting#

Task: binary sentiment on the Stanford Sentiment Treebank (SST-2), movie-review phrases labelled positive or negative; 50,000 of the 67,349 training phrases are used, and the 872-phrase validation split is the test set (the official test labels are hidden). The classes are 44 % negative and 56 % positive in training, close to balanced, so accuracy is reported.

Model: GPT-2 small (124 M parameters, from the GPT note) with a linear head ϕ∈R768×2\phi \in \mathbb R^{768 \times 2} on the hidden state of the last non-padding token. In a causal model only the last position has attended to the whole input, so it is the one that can summarize the phrase; the [CLS] convention of BERT reads the first position instead because every position there sees everything. A [PAD] token is added to the tokenizer and the embedding table grows by one row.

Training: AdamW, batch size 128, 1,000 steps (about 2.6 passes over the 50,000 phrases), mixed precision, evaluation every 250 steps.

Baseline: train everything#

Full fine-tuning updates θ\theta and ϕ\phi together with the same optimizer:

(θ,ϕ)←(θ,ϕ)−η ∇θ,ϕ L.(\theta, \phi) \leftarrow (\theta, \phi) - \eta\, \nabla_{\theta,\phi}\, \mathcal L .

The learning rate is 2×10−62 \times 10^{-6}, three orders of magnitude below the 10−310^{-3} used to train the small Transformer from scratch in the architecture note. Pretrained weights already encode what the model knows; a large step destroys that before the head has learned to use it, a failure called catastrophic forgetting. Small steps move the body only as far as the task requires.

Result: 0.890 validation accuracy, loss 0.277.

Cost: the optimizer keeps two moment estimates per parameter and the backward pass keeps a gradient per parameter. In 32-bit floats the 124 M weights take 0.5 GB, the gradients another 0.5 GB and the AdamW states 1.0 GB, before activations: about four times the memory of the model itself. For a 7 B-parameter model the same arithmetic gives 112 GB, more than one GPU holds, and that is the number the PEFT note sets out to reduce.

Linear probing: train only the head#

Freeze θ\theta, train ϕ\phi alone at a larger learning rate (10−410^{-4}, since the head starts from random values and has nothing to forget). The body’s output is a fixed feature and the head is logistic regression on it, as in the frozen-features experiment of the previous note, except that here the features are computed on the fly.

Result: 0.737, loss 0.541.

The gap to full fine-tuning is 15 points. The last-token hidden state of a language model was trained to predict the next token, not to separate sentiments, and a linear read-out finds only part of the signal. The BERT [CLS] vector on AG News gave 0.902 under a similar regime, but that comparison changes the model, the task, the data and the training of the head at once. It cannot tell whether topic is more linearly present than sentiment, or whether a bidirectional summary token is a better feature than a last-position causal state. Both are plausible hypotheses; testing either needs a comparison that changes one factor, such as the same model probed on both tasks, or an encoder and a decoder of similar size probed on SST-2. Linear probing is a measurement of what the representation contains, and a cheap but weak classifier.

Wall-clock time was about the same as full fine-tuning (6.5 against 7.5 minutes): every step still runs the full forward pass, and the backward pass through the frozen body is skipped but was not the dominant cost at this size. Memory, not time, is what probing saves.

LP-FT: head first, then everything#

Kumar et al. (2022) showed this two-stage schedule (linear probing, then fine-tuning) preserves pretrained features better than fine-tuning from a random head, with the largest gains on data that differs from the fine-tuning distribution. Here full fine-tuning starts from the head trained by linear probing above and proceeds at 2×10−62 \times 10^{-6}.

Result: 0.900, loss 0.258, one point above fine-tuning from scratch at the same cost plus one probing run.

RegimeTrainableη\etaVal. accuracyVal. loss
Full fine-tuning124 M2×10−62\times10^{-6}0.8900.277
Linear probing1,53610−410^{-4}0.7370.541
LP-FT1,536, then 124 M10−410^{-4}, then 2×10−62\times10^{-6}0.9000.258
Figure

Three fine-tuning regimes: what is trained, and when

Fine-tuning regimesfull fine-tuning: body and head trained, η = 2·10⁻⁶; 0.890. linear probing: only the head trained, η = 10⁻⁴; 0.737. LP-FT: only the head trained, η = 10⁻⁴ → body and head trained, η = 2·10⁻⁶; 0.900.full fine-tuning0.890head φ, 1,536body θ, 124 Mη = 2·10⁻⁶linear probing0.737head φ, 1,536body θ, 124 Mη = 10⁻⁴LP-FT0.900head φ, 1,536body θ, 124 Mη = 10⁻⁴head φ, 1,536body θ, 124 Mη = 2·10⁻⁶trainedfrozenaccuracyFine-tuning regimesfull fine-tuning: body and head trained, η = 2·10⁻⁶; 0.890. linear probing: only the head trained, η = 10⁻⁴; 0.737. LP-FT: only the head trained, η = 10⁻⁴ → body and head trained, η = 2·10⁻⁶; 0.900.full fine-tuning0.890head φbody θη = 2·10⁻⁶linear probing0.737head φbody θη = 10⁻⁴LP-FT0.900head φbody θη = 10⁻⁴head φbody θη = 2·10⁻⁶trainedfrozen
GPT-2 small on SST-2, 1,000 steps per phase. Full fine-tuning trains the body and a randomly initialized head from the first step at a small learning rate. Linear probing freezes the body and trains only the head (1,536 weights) at a larger rate. LP-FT probes first and then fine-tunes everything starting from the trained head, so the first steps do not distort the pretrained features with gradients from a random layer. On the right, accuracy on the SST-2 validation set (872 phrases; one phrase is 0.11 points).

The one-point difference between the first and third rows is at the edge of what 872 examples can resolve: it is 9 phrases (one example is 0.11 %), and the standard error of a single accuracy near 0.9 on 872 examples is 0.9⋅0.1/872≈0.010\sqrt{0.9 \cdot 0.1 / 872} \approx 0.010, one point. So this run observed 0.900 against 0.890, which is consistent with the mechanism above but does not rank the two regimes, and does not show that they are equivalent either.

Hyperparameters that matter#

  • Learning rate is the main one: 10−510^{-5} to 5×10−55 \times 10^{-5} for encoders, lower for larger decoders. Too high forgets; too low does not move.
  • Schedule. Linear warm-up over the first few hundred steps followed by decay is the default for post-norm models and harmless for pre-norm ones.
  • Epochs. Two to four over the labelled set; more overfits small datasets.
  • Weight decay of 0.01 and dropout at the pretrained rate; both regularize a model that has far more parameters than the task has examples.
  • Seeds. Runs on small datasets vary by a point or more with the seed and data order; a decision needs several runs (Dodge et al., 2020).

Padding each batch, not the dataset#

A batch is a tensor, so its sequences must have one length: the shorter ones are padded and the attention mask hides the padding (tokenization note). The question is to which length. Padding the whole dataset to one maximum is simple and wastes compute on positions that are masked anyway. Dynamic padding tokenizes every example once, without padding, and pads only when a batch is formed, to the longest sequence of that batch. Batches are reshuffled every epoch, so their lengths are not known in advance; the padding belongs in the function that collates examples into a batch.

Worked example (toy lengths). A batch of eight phrases of 5, 9, 12, 7, 23, 4, 10 and 6 tokens holds 76 real tokens. Padded to 128, it is 1,024 positions, 7.4 % of them real; padded to its own longest phrase, 23, it is 184 positions, 41.3 % real. The feed-forward layers cost in proportion to the positions, 5.6 times less here. Self-attention compares every pair of positions within a sequence, 1282=16,384128^2 = 16{,}384 against 232=52923^2 = 529 pairs, 31 times less.

The predictions do not change, up to rounding: padded positions are masked out of attention, and the classifier reads a real token. Here that is the last non-padding token, found from the attention mask, which works for any padded length as long as the padding is on the right. Grouping examples of similar length into the same batch removes most of the remaining padding, at the price of batches that are less random.

What changed#

Frozen features + linear modelFine-tuning
Parameters updatedHeadHead and body
Learning rateLarge (random head)Small (pretrained body)
MemoryHead only≈ 4× model size in fp32
Feature qualityFixed by pretrainingAdapted to the task
SST-2, GPT-2 small0.7370.890 / 0.900 (LP-FT)

Limitations#

  • One full copy per task. The fine-tuned θ\theta is a new 124 M-parameter (or 7 B) artefact; ten tasks are ten copies.
  • Memory scales with the model: weights, gradients and optimizer states.
  • Forgetting. The fine-tuned model is specialized; its general abilities degrade.
  • Noise. Small validation sets cannot resolve differences of a point.
  • Last-token features are weak. A causal model’s representation is not designed for classification; an encoder of the same size is often better and cheaper for this task, although this course does not compare the two on SST-2.

What you should remember#

  • Fine-tuning: train body and head together at a small learning rate (10−610^{-6} to 10−510^{-5} range here); large rates erase pretraining.
  • Linear probing: freeze the body, train the head at a larger rate; cheap in memory, limited by what the representation already encodes (0.737 here).
  • LP-FT: probe first, then fine-tune from the trained head; 0.900 against 0.890 here, a difference within noise, so the case for it rests on the mechanism and on Kumar et al., and it should be checked on each task.
  • Memory for full fine-tuning is about four times the model’s weights in fp32.
  • Every task fine-tuned this way produces a full copy of the model.
  • Tokenize once without padding and pad each batch to its own longest sequence; with the attention mask and right padding the predictions are unchanged and the compute falls with the padding.

Notebook#

week10_PEFT.ipynb. Cell numbers count every cell, markdown included, from 1. Quoted numbers come from the run saved with the course; cells changed since then for current library versions have no saved output.

  • Cells 4–7 load SST-2 (train[:50000] and validation); cell 19 prints the class proportions (0.442 / 0.558 train, 0.491 / 0.509 validation).
  • Cells 10–14 load GPT-2 with a sequence-classification head, add [PAD], resize the embeddings and print 124,442,112 parameters. The Model wrapper’s forward runs the transformer, selects the hidden state at index attention_mask.sum(-1) - 1 (the last real token) and applies the score layer.
  • Cell 23 defines the training loop (AdamW, cross-entropy, torch.autocast in float16, evaluation every 250 of 1,000 steps with loss and accuracy curves) and evaluate.
  • Cells 25–27: full fine-tuning at 2×10−62\times10^{-6}, validation loss 0.2772 and accuracy 0.8899.
  • Cells 31–34: linear probing (AdamW(model.model.score.parameters(), lr=1e-4)), 0.5412 and 0.7374; the head is saved to head.pt.
  • Cells 56–60: LP-FT, loading head.pt and fine-tuning everything at 2×10−62\times10^{-6}: 0.2580 and 0.9002.
  • Cells 62–64 print predictions with probabilities for ten validation phrases.

Cells 35–54 (LoRA and prompt tuning on the same setup) belong to the next note.

The LLM notebook of the HSE-LAMBDA Deep Generative Models course, Transformers-LangChain.ipynb, gives the Trainer recipe for the same kind of task in markdown cells 32–35, which are not run (same cell-number convention). It fine-tunes distilbert-base-uncased on the Rotten Tomatoes movie-review sentences with learning rate 2⋅10−52 \cdot 10^{-5}, batch size 8, 2 epochs and evaluation every 500 steps (the default once the strategy is "steps"). The dataset is tokenized once with map(batched=True) and no padding, and DataCollatorWithPadding(tokenizer=tokenizer) pads each batch to its longest sequence, the dynamic padding described above. Two details need changing before it runs: current versions of Transformers accept only eval_strategy and reject evaluation_strategy, and the trainer evaluates on the test split although the dataset has a validation split, which is the one to monitor during training.

Next#

Full fine-tuning gives the best accuracy and costs a full copy of the model plus four times its memory during training. The next note keeps the body frozen and adds a small number of new parameters in the right places, so that the same SST-2 task is learned to the same accuracy while training 0.24 % of the weights.

References#