Training and Fine-Tuning Deep Models
By the end: You choose a transfer-learning strategy for a given data and compute budget, estimate approximate training memory by component, state what mixed precision does and does not save, and judge a comparison against seed-to-seed spread.
At a glance
Pick a rung on the freeze-to-full ladder from data size and domain, budget memory by component, and compare strategies over several seeds.
The ladder
Frozen plus head, top blocks, full fine-tune, from scratch: cheapest first.
Bytes per parameter
Adam training costs about 16 bytes per trainable parameter, fp32 or mixed.
Mixed precision
Halves activations and speeds compute; does not shrink optimizer states.
Freeze to fit
A frozen weight keeps only itself in memory: no gradient, no moments.
Seeds
Three or more seeds per strategy; read gaps against the spread.
Recall
Checkpoint warm-up, no notes, one line each: (1) the output size of a 5 x 5 convolution with stride 2 and padding 1 on 28 x 28; (2) the shape of the attention score matrix for 3 tokens; (3) why a hidden state decays over many steps; (4) how you would tell a leak from overfitting in a loss curve; (5) what the Adam optimizer stores per parameter besides the weight and gradient.
Concept
A pretrained model is a feature extractor paid for by someone else's compute. Your decisions: what to keep fixed, what to retrain, how gently, and with what evidence. Like an experienced hire, keep most of what it knows and retrain the specifics. The analogy breaks on measurement: for a model you can see what was lost, with a retention set and a baseline.
Predict first
Full fine-tuning with Adam: going from fp32 to mixed precision (fp16 copies plus an fp32 master copy), what happens to the weights, gradients and optimizer states per parameter?
Text description of the figure
A stacked bar of approximate training memory split into weights, gradients, optimizer states and activations, with controls for the parameter count, the precision, the optimizer and the batch size. The values are labelled as approximate.
- 1. Frozen backbone + new headA linear probe on pretrained features. Fewest trainable parameters; least risk of forgetting.
- 2. Unfreeze the top blocksSmaller learning rates for lower layers, or unfreeze gradually.
- 3. Full fine-tuneEvery parameter trains. Most flexible, most memory, most seed variance on small data.
- 4. Train from scratchOnly with plenty of data or a domain unlike any pretraining corpus.
Starting from a trained probe, then fine-tuning, avoids a random head distorting good features (Kumar et al., 2022).
Frozen backbone + head
Features already fit; few labels cannot support more trainable parameters.
Fine-tune top blocks or all
Enough data to adapt; similar domain keeps it stable.
Careful
Top blocks with small rates, strong baseline, retention set; transfer may not help.
Full fine-tune or scratch
Enough data to relearn features that do not transfer.
- weights
- gradients
- fp32 master
- Adam moments
- activations
Planning estimate, not a measurement: assumed activation factor 20, no framework overhead, no attention-score storage, no KV cache. Mixed precision moved only the activation bar (4.0 GB instead of 8.1 GB).
4 of 7 strategies fit a 16 GB card. LoRA cuts the states to almost nothing but keeps 4.0 GB of activations, because backprop must still pass through all 24 layers to reach the lowest adapter; only freezing whole blocks or checkpointing shrinks that bar.
Each bar is the application's estimator split into its parts. Freezing removes states, checkpointing removes activations, and LoRA removes states but not activations.
| Component | fp32 | Mixed precision | Frozen parameter |
|---|---|---|---|
| Working weights | 4 | 2 | 4 (fp32) or 2 |
| Gradients | 4 | 2 | 0 |
| fp32 master copy | 0 | 4 | 0 |
| Adam moments (two) | 8 | 8 | 0 |
| Total | 16 | 16 | 4 or 2 |
Activations come on top and scale with batch x sequence x hidden x layers; mixed precision roughly halves them.
Freeze the split
Versioned train, validation and test; identical for every strategy
Log everything
Config, seed, split id, library versions, hardware
Three or more seeds
Per strategy; small-data fine-tuning varies a lot across seeds
Read against spread
A gap inside the seed spread is not a ranking
Open the test set once
For the single chosen strategy
Bit-identical results are not guaranteed across releases, platforms or CPU versus GPU, even with fixed seeds (PyTorch reproducibility notes).
- The ladder, cheapest first: frozen backbone plus a new head (a Linear probe); top blocks unfrozen; full fine-tune; from scratch. Early layers tend to be general, later ones task-specific (Transfer learning).
- Memory is weights, gradients (trainable parameters only), Optimizer state and activations. With Adam that is about 16 bytes per trainable parameter in fp32 and also with mixed precision; mixed precision mainly halves activations and speeds compute.
- What shrinks states is training fewer parameters, sharding, 8-bit optimizers or parameter-efficient methods (Day 44). These are planning estimates; measure the peak of a real step.
- Reproducibility: fix seeds and log config, split and library versions, but bit-identical results are not guaranteed across releases or hardware. Run at least three seeds per strategy.
Assumptions
The target task is related enough to the pretraining distribution for the early features to transfer, and your preprocessing reproduces the pretrained model's exactly.
The validation and test sets are versioned and untouched, and the same for every strategy.
The memory estimate is a planning tool; the decision to buy hardware or change strategy rests on a measured peak.
Application and workplace decision
The 4,000 complaints go to a 12-way classifier. Two candidates, three strategies each (frozen backbone plus head, last two blocks plus head, full fine-tune):
- a 110M-parameter encoder: 12 layers, hidden 768, batch 32, 256 tokens;
- a 1.3B decoder: 24 layers, hidden 2048, batch 8, 256 tokens.
Before running, predict which of the six rows fit in 16 GB with mixed precision.
- The code estimates weights, gradients, master copy, Adam moments and activations.
- Activations use an assumed 20 values per hidden unit per layer, an order-of-magnitude figure.
- It is a planning estimate; the real number is a measured peak.
Python
GB = 1e9
def memory_gb(params, trainable, trainable_layers, hidden, batch, seq, mixed, act_factor=20):
"""Approximate training memory in GB: weights + gradients + optimizer states + activations."""
w_bytes = 2 if mixed else 4
weights = params * w_bytes
grads = trainable * w_bytes
master = trainable * 4 if mixed else 0 # fp32 master copy used for the update
adam = trainable * 8 # two fp32 moments per trainable weight
# activations are kept only for layers that backprop passes through (assumed factor)
acts = batch * seq * hidden * trainable_layers * act_factor * w_bytes
return (weights + grads + master + adam + acts) / GB
models = {"110M encoder": dict(params=110e6, layers=12, hidden=768, batch=32, seq=256),
"1.3B decoder": dict(params=1.3e9, layers=24, hidden=2048, batch=8, seq=256)}
head = 4e6 # new classification head
print("model strategy trainable fp32 GB mixed GB fits 16 GB (mixed)")
for name, m in models.items():
L = m["layers"]
for strat, frac in [("frozen backbone + head", 0), ("last 2 blocks + head", 2 / L), ("full fine-tune", 1)]:
trainable = head + frac * (m["params"] - head)
args = dict(params=m["params"], trainable=trainable, trainable_layers=round(frac * L),
hidden=m["hidden"], batch=m["batch"], seq=m["seq"])
f32, mp = memory_gb(mixed=False, **args), memory_gb(mixed=True, **args)
print(f"{name:14s} {strat:24s} {trainable / 1e6:8.0f}M {f32:9.1f} {mp:10.1f} {'yes' if mp <= 16 else 'no'}")
Mini-assessment
Technical depth
The memory model, defined, and bytes per parameter derivedThe planning equation symbol by symbol, what one Adam step keeps per trainable weight, why mixed precision still costs 16 bytes, and the rounding behind the master copy.
- all parameters and trainable parameters
- bytes per working value: 4 in fp32, 2 in mixed precision
- the fp32 master copy, only in mixed precision
- Adam's two fp32 moments
- batch, sequence, hidden size, trainable layers; 20 is an assumed factor
- Training memory (planning sum)all parts as if resident
- countsPer-parameter termssame for any batch
- every stored weightWorking weight4 or 2 B, frozen too
- trainable onlyGradient, master copy, moments12 B fp32, 14 B mixed
- a modelActivation termbatch x sequence x width x depth
- until backwardInputs kept for the backward passfreed as grads appear
- countsPer-parameter termssame for any batch
The sum is a planning number, not a measured peak: a real step frees activations during the backward pass. Adam's m and v are running averages of g and g^2, so they must persist and cannot be recomputed like activations. Freezing moves a parameter from 16 bytes to 2 or 4.
Forward
The working weight w: 4 bytes in fp32, 2 in fp16 or bf16
Backward
One gradient per trainable weight, in the working precision: 4 or 2 bytes
Adam state
Two running averages per weight, m (mean of g) and v (mean of g^2), kept in fp32: 4 + 4 = 8 bytes
Master copy (mixed only)
An fp32 copy of w receives the update, because tiny updates vanish in fp16 (next figure): 4 bytes
Sum, fp32
4 + 4 + 8 = 16 bytes per trainable parameter
Sum, mixed
2 + 2 + 4 + 8 = 16 bytes: the halved copies are paid back by the master copy
Frozen parameter
Only the working weight: 4 or 2 bytes. No gradient, no state, no master copy
The same accounting appears in Rajbhandari et al. (2020), Section 3.1. Activations, temporary buffers and framework overhead come on top.
| Optimizer | Weight | Gradient | State | Total |
|---|---|---|---|---|
| SGD, no momentum | 4 | 4 | 0 | 8 |
| SGD with momentum | 4 | 4 | 4 (velocity) | 12 |
| Adam or AdamW | 4 | 4 | 8 (m and v) | 16 |
| Frozen (any optimizer) | 4 | 0 | 0 | 4 |
Swapping Adam for SGD saves a quarter to a half of the states but changes training dynamics (Day 33). Training fewer parameters moves every trainable row to the frozen row.
The weight
w = 1.0. Representable fp16 numbers near 1.0 are 2^-10 = 0.000977 apart
The update
A typical fine-tuning step: learning rate x gradient = 0.0001
In fp16
1.0 + 0.0001 is closer to 1.0 than to 1.000977, so the sum rounds back to exactly 1.0: the step is lost
In fp32
Spacing near 1.0 is 2^-23 = 1.2e-7, so 1.0 + 0.0001 = 1.0001 is kept
Over 1,000 steps
fp16 weights never move; the fp32 master copy has moved by 0.1
Checked in numpy (np.float16 and np.spacing), including the 1,000-step loop. This is a precision problem in the update, distinct from the range problem of tiny gradients in the mixed-precision chapter.
Counting a 110M encoderadvancedThe parameter count of a 12-layer, 768-wide encoder derived from its matrices, so a model card's headline number can be checked and split into what a strategy trains.
Attention, one layer
W_Q, W_K, W_V, W_O, each H x H with a bias: 4H^2 + 4H
Feed-forward, one layer
H to 4H and back, with biases: 8H^2 + 5H
Two LayerNorms
A scale and a shift per unit each: 4H
One layer
12H^2 + 13H = 12 x 589,824 + 9,984 = 7,087,872
Twelve layers
12 x 7,087,872 = 85.1 million
Embeddings
30,000 tokens x 768 + 512 positions x 768 + small extras: about 23.4 million
Pooler and total
768 x 768 + 768 = 0.6 million; total about 109 million, the ~110M quoted for BERT-base
Vocabulary and layout follow Devlin et al. (2019) for BERT-base; counted in Python. Per layer the count grows with H^2, never with sequence or batch. A fifth sits in the embedding table; "last two blocks" trains 2 x 7.1M plus the head, a frozen backbone only the head.
Activation memory, derived and computedadvancedWhich tensors the backward pass needs, the linear and the S-squared term, a worked table for the 1.3B decoder, and where the estimate breaks.
- batch, sequence length, hidden size: tensors with one value per token per unit
- values stored per token and unit for backward; the code assumes 20, an order-of-magnitude figure
- attention scores, one (S, S) matrix per head; gone with fused kernels such as FlashAttention
- layers the backward pass runs through: frozen layers below the lowest trainable one store nothing
- bytes per stored value: 4 in fp32, 2 in mixed precision, which is why mixed precision halves this term
The code uses only the first term. Setting the two terms equal gives the crossover S = c H / a: past it the S-squared term dominates. Measure the real peak with the framework's allocator statistics.
| Sequence length | Linear term per layer | Score term per layer | Larger term |
|---|---|---|---|
| 256 | 8 x 256 x 2048 x 20 x 2 B = 168 MB | 16 x 8 x 256^2 x 2 B = 17 MB | Linear (10x) |
| 2,048 | 1.34 GB | 1.07 GB | About equal |
| Crossover | S = 20 x 2048 / 16 = 2,560 | Equal to the linear term | Above it, scores dominate |
| Over 24 layers, S = 256 | 4.0 GB (the application's activation bar) | 0.4 GB extra if scores are stored | Linear |
Computed in Python under the stated assumptions (factor 20, 16 heads, naive attention). Long-sequence fine-tunes hit the S-squared term first; checkpointing or a fused kernel is then the lever, not precision.
- the layer's input from the forward pass, kept until the backward pass reaches this layer
- arrives only during the backward pass, from the layers above
- so every input on the backward path is resident at the same time; frozen layers below the lowest trainable one keep nothing
Batch and sequence move only this term; the per-parameter states are the same for any batch. Where it breaks: c depends on architecture and kernels, fused attention removes the S-squared term, and allocator overhead adds to the peak; measure torch.cuda.max_memory_allocated on one real step.
The 110M example, step by stepThe memory equation term by term on one row of the application, then the breakdown of the row that does not fit and all six strategies.
Trainable parameters
Head 4M + 2/12 of the other 106M = 4M + 17.7M = 21.7M
Working weights, all
110M x 2 bytes = 0.22 GB: frozen weights still sit in memory
Gradients
21.7M x 2 bytes = 0.04 GB
fp32 master copy
21.7M x 4 bytes = 0.09 GB
Adam moments
21.7M x 8 bytes = 0.17 GB
Activations
Only 2 layers store them: 32 x 256 x 768 x 2 x 20 x 2 bytes = 0.50 GB
Total
About 1.03 GB, the application's 1.0 GB row; half of it is activations
Each line is one term of the memory equation. Redo it on paper, then check it against the printed row. The measured peak will differ: freed activations lower it, while the CUDA context, temporary buffers and fragmentation raise it.
- Working weights (2 bytes)2.6 GB
- Activations (assumed factor)8.1 GB in fp32; the only part mixed precision halves+4 GB
- Gradients (2 bytes)+2.6 GB
- fp32 master copy (4 bytes)+5.2 GB
- Adam moments (8 bytes)+10.4 GB
- Total, against a 16 GB card24.8 GB
Computed from the code's formula, in the order a step needs them: forward (weights, activations), backward (gradients), update (master copy, moments). The card is full before the update terms are added.
The application's printout. Planning estimates; measure the peak of one real training step before committing.
Properties of the estimate, and when it breaksadvancedWhich parts of the estimate are exact counts and which are assumptions, how each assumption fails, and the one-line check for each.
| Assumption | Holds when | Breaks when | Diagnostic |
|---|---|---|---|
| 16 bytes per trainable parameter | Adam or AdamW with fp32 moments | 8-bit optimizer states, sharded states (ZeRO) or plain SGD | Sum the sizes of the optimizer's state tensors after one step |
| A frozen weight costs only itself | requires_grad is False, so no gradient and no Adam state is ever created | A layer frozen in the config still has requires_grad True | Count parameters with requires_grad True and compare with the plan |
| Blocks of equal size | Transformer blocks dominate the count | Embeddings hold a fifth of 110M: the code's 'last 2 blocks' trains 17.7M, two real blocks have 14.2M | Count per module from named_parameters instead of a fraction |
| 20 stored values per token, unit and layer | A standard block, as an order of magnitude | Other activations, dropout masks, kernels or the S-squared scores change the constant | Peak at two batch sizes: the slope is memory per example, the intercept the fixed part |
| No activations below the lowest trainable layer | Only the head or top blocks train | An adapter or unfrozen layer near the input sends backprop through the whole stack | Find the lowest module with requires_grad True |
| No overhead | Never exactly | CUDA context, allocator cache, temporary buffers and fragmentation add to the peak | Compare max_memory_allocated with memory_reserved and the card's free memory |
The per-parameter terms are counts and exact for a given optimizer and precision; the activation term is a model with an assumed constant. One measured step (reset_peak_memory_stats, then max_memory_allocated) replaces every assumed number. Verified against the application code and the 110M count above.
Conservative part
The sum ignores that activations are freed while new gradients appear.
Optimistic part
It leaves out overhead, buffers and fragmentation. A plan that just fills the card on paper is no plan.
Two peaks beat any constant
peak(B) is close to fixed + B x per-example activations: peaks at two batch sizes give both terms.
Which term dominates
Defect-image CNN: few weights, big feature maps, activations lead. 1.3B model on short text: states lead.
Neither error has a fixed size, which is why one measured step replaces every assumed number.
What shrinks memory, and what does notFreezing, sharding, checkpointing and parameter-efficient methods compared, and which bar each one moves.
Train fewer parameters
Freezing removes gradients, master copy and moments for those weights.
Sharding
Splits states across devices (ZeRO, Rajbhandari et al., 2020); helps only with more than one device.
Activation checkpointing
Stores a subset, recomputes the rest: memory ~ sqrt(depth) for about one extra forward (Chen et al., 2016).
Parameter-efficient methods
Small added matrices (Day 44): states go, activations stay above the lowest adapter.
Smaller batch or sequence
Shrinks activations only; may need gradient accumulation.
Mixed precision, step by stepadvancedWhy fp16 needs loss scaling and bf16 usually does not, shown on a gradient histogram.
| Criterion | fp32 | fp16 | bf16 |
|---|---|---|---|
| Bytes | 4 | 2 | 2 |
| Range | Wide | Narrow: small gradients underflow to zero | Same exponent range as fp32, fewer mantissa bits |
| Needs loss scaling | No | Yes | Usually not |
| Typical role | Master weights, optimizer moments, sensitive operations | Working copies and activations | Working copies and activations on hardware that supports it |
Micikevicius et al. (2018) for fp16 mixed precision with loss scaling; PyTorch's AMP documentation for autocast and gradient scaling.
Unscaled, 22.5% of these gradients become exactly 0 in fp16 and 73% more lose precision as subnormals. In bf16, 0 of 6,000 would.
- rounded to zero
- overflow (inf)
Scaling the loss by S shifts the whole histogram right by log10 S. Here S = 2^16 loses the least (nothing to zero, nothing to overflow); too large an S overflows, which dynamic scaling detects as inf and answers by skipping the step and lowering S.
Synthetic gradient magnitudes (6,000 seeded log-normal draws centred near 3e-7); real networks differ by layer and over training. The fp16 and bf16 limits are the formats' exact thresholds.
fp16 has a narrow exponent range; multiplying the loss shifts every gradient by the same factor, and the update divides it out again on the fp32 master copy.
- Forward in fp16autocast; some ops kept in fp32
- Loss times scale Slifts small gradients
- Backward in fp16gradients times S
- Unscale; inf or nan?
- no: to Update fp32 master
yes - Skip the steplower S, next batch
- Update fp32 mastercopy back to fp16
Dynamic scaling also raises S again after a run of clean steps. bf16 has fp32's exponent range, so it usually needs no loss scaling.
Fine-tuning gently, the knobs and the four strategiesProbe first, unfreeze gradually, lower rates for lower layers, check retention; the knobs and their defaults; probe, partial, full and LoRA compared.
Linear probe first
Train the new head on frozen features until it converges
Then unfreeze
Top blocks first; lower layers later or not at all (gradual unfreezing)
Discriminative rates
Smaller learning rates for lower layers (Howard and Ruder, 2018)
Warmup and early stopping
Small first steps; stop at the validation minimum (Day 34)
Retention check
Score a held-out set from the original task to see what was forgotten
| Knob | Default or common start | Trade-off |
|---|---|---|
| AdamW learning rate | 1e-3 in PyTorch; fine-tuning usually far lower | Speed of adaptation vs distorting pretrained features |
| AdamW weight decay | 0.01 (decoupled) | Pull toward zero vs pull toward the pretrained weights |
| AdamW betas, eps | (0.9, 0.999), 1e-8 | Rarely the first knob; eps matters in low precision |
| Warmup steps | a few percent of training (heuristic) | Protects the new head's first updates |
| Batch and accumulation | what fits; accumulate to the target | Activations vs gradient noise and the sensible learning rate |
| Layer-wise rates | lower for lower layers | Keeps general features vs slower adaptation |
| Epochs and early stopping | stop at the validation minimum | Underfitting vs memorising a small set |
AdamW defaults read from the PyTorch 2.14 page. Rows marked heuristic are common practice, not library facts.
| Criterion | Probe (frozen + head) | Top blocks + head | Full fine-tune | LoRA (Day 44) |
|---|---|---|---|---|
| Trainable parameters | Head only | A few blocks | All | Small low-rank matrices |
| Memory saved | States and activations of the backbone | States and activations of frozen blocks | Nothing | States only; activations remain |
| Forgetting risk | None: the backbone is unchanged | Low to moderate | Highest, especially at high rates | Lower; also learns less (Biderman et al., 2024) |
| Serving | One shared backbone | A partial copy per task | A full copy per task | Adapters can be merged into the weights |
Seed spread and the decisionHow often a few seeds pick the truly better strategy, and how to read the three strategies' results.
- true gap 0.005 (below the spread): 136 seeds for a 95% chance
- true gap 0.015 (below the spread): 16 seeds for a 95% chance
- true gap 0.03 (above the spread): 4 seeds for a 95% chance
With 3 seeds and seed spread 0.025, the truly better strategy comes out on top 60%, 77%, 93% of the time for gaps of 0.005, 0.015, 0.03. A gap well below the spread is close to a coin flip.
Lines: normal-theory probability; open dots: 1,000 seeded simulated comparisons per point. Assumes independent runs with equal spread for both strategies; pairing runs that share a seed and data order usually needs fewer runs.
Computed from the normal model and checked by simulation. The gap has to be large relative to the seed spread before a few seeds rank strategies reliably.
| Strategy | Macro-F1, mean +/- sd over 3 seeds | Reading |
|---|---|---|
| Frozen + head | 0.71 +/- 0.01 | Stable, lowest |
| Last two blocks | 0.74 +/- 0.02 | Clear gain over frozen: 0.03 against spreads of 0.01-0.02 |
| Full fine-tune | 0.75 +/- 0.04 | 0.01 above last two blocks, well inside its own spread |
Fine-tuning on small datasets is known to vary a lot across seeds (Dodge et al., 2020).
What a staff-level answer addsadvancedNuances and misreadings that separate an experienced answer from a textbook one.
| Criterion | Common reading | Staff-level reading |
|---|---|---|
| Budgeting | 'Will the model fit?' | By component: states follow trainable params; activations batch, sequence, depth |
| Mixed precision | 'It halves memory' | With Adam it halves activations; per-parameter states stay about the same |
| LoRA | 'It makes training fit anywhere' | It removes states; long sequences can still be dominated by activations |
| The decision rule | Chosen after seeing the runs | Stated before the runs: which gap, against which spread, at which cost |
| After the win | Ship the target-task winner | Check a retention set: the win can cost what made the base model useful |
Memory anchor
Freeze to shrink the states, not the weights; mixed precision shrinks activations; three seeds before a ranking.
Useful for
Choosing which layers to train under a data and GPU budget, and defending the choice with a memory estimate and a seed-spread comparison.
Reading a fine-tuning proposal for what it leaves out - preprocessing parity, a baseline, seeds, a retention set.
Planning hardware or choosing between full and parameter-efficient methods (Day 44) for a defect-image model on a factory line, a clinical-note classifier or a call-centre transcript tagger.
Not suitable for
Producing a measured memory number - use a profiler or peak allocation on a real batch.
Deciding between model families without your own labelled evaluation on the same split.
Replacing the study of a specific framework's memory and mixed-precision behaviour.
Traps
Reporting one run
Fine-tuning small datasets varies a lot across seeds, and a single run can show a gain that disappears on another seed, or a divergence that is not the strategy's fault.
How to avoid it
Run three or more seeds per strategy on a fixed split, report mean and spread, and require differences to exceed the spread before ranking strategies.
Mismatched preprocessing and silent forgetting
Different tokenizer, image normalisation or resolution from the pretraining setup degrades features without an error, and full fine-tuning at a high learning rate can erase what made the model useful.
How to avoid it
Reproduce the pretrained model's preprocessing exactly, start with a probe, use smaller learning rates for lower layers, and keep a retention set and the classical baseline in the report.
Senior reasoning
A senior sets the experiment before the run: a ladder of strategies from probe to full fine-tune, one fixed split, a baseline, three or more seeds and a memory estimate that rules out what cannot fit. They read gaps against spread, weigh compute and maintenance, and keep the simplest strategy that meets the need. Their checkpoint question is always the same: which earlier day tells me what this result can and cannot show?
Interview question
"You have 5,000 labelled examples and a pretrained model. How do you decide what to fine-tune, and how do you know it worked?"
Recall cards scheduled after completion
- Estimate training memory per trainable parameter with Adam, and say what mixed precision changes.
- How do you decide between probe, partial and full fine-tuning?
Prerequisites and related concepts
References
- How transferable are features in deep neural networks?core
Early layers are general and later layers task-specific, and transfer degrades with distance between tasks; the basis for the freeze-train ladder.
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscore
Mixed-precision Adam needs about 2 + 2 bytes per parameter for fp16 weights and gradients plus 12 bytes for fp32 weights and two moments, i.e. 16 bytes per parameter before activations (Section 3.1). (Section 3.1)
- Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distributionoptional
Full fine-tuning can distort pretrained features, supporting a linear probe first and then fine-tuning.
- Universal Language Model Fine-tuning for Text Classificationoptional
Gradual unfreezing and layer-wise learning rates as fine-tuning techniques for small datasets.
- Mixed Precision Trainingoptional
Mixed-precision training with an fp32 master copy of weights and loss scaling for half-precision range limits.
- Automatic Mixed Precision package - torch.amp (PyTorch documentation)optional
Gradient scaling multiplies the loss because small float16 gradients can underflow to zero.
- Reproducibility (PyTorch notes)core
Completely reproducible results are not guaranteed across releases, platforms or CPU versus GPU, and deterministic modes trade speed for determinism.
- Fine-Tuning Pretrained Language Models: Weight Initializations, Data Orders, and Early Stoppingoptional
Distinct seeds give substantially different BERT fine-tuning results, and many runs diverge on small datasets.
- AdamW (PyTorch documentation)optional
AdamW defaults (learning rate, betas, eps, decoupled weight decay 0.01) behind the knobs table.
- LoRA Learns Less and Forgets Lessoptional
LoRA underperformed full fine-tuning on the target domain but preserved more of the base model's behaviour outside it.
- On the difficulty of training recurrent neural networksoptional
Exploding and vanishing gradients in recurrent networks and gradient-norm clipping (used in probe q-040-p2).
- Are Transformers Effective for Time Series Forecasting?optional
Simple linear models outperformed several transformer forecasters on standard benchmarks (used in probe q-040-p4).
- How to Use t-SNE Effectivelyoptional
Cluster sizes and between-cluster distances in a t-SNE plot can mislead (used in probe q-040-p8).
- Accounting for Variance in Machine Learning Benchmarksoptional
Variance from data sampling, initialisation and hyperparameter choice markedly affects benchmark conclusions.
- Training Deep Nets with Sublinear Memory Costoptional
Activation checkpointing trades roughly one extra forward pass for memory that grows with the square root of depth.
- CUDA semantics (PyTorch documentation)optional
Memory management and the functions that report allocated and peak memory, for measuring instead of estimating.
Optional sections
Deep dive
Why 16 bytes is the planning number. fp32 Adam: 4 (weights) + 4 (gradients) + 4 + 4 (moments). Mixed precision: 2 + 2 + 4 (fp32 master) + 4 + 4. Both are 16 per trainable parameter, which is why frameworks and papers quote the same figure and why states, not the dtype of the working copy, dominate for Adam. Frozen parameters add their working copy only (2 or 4 bytes).
Activations. Per transformer layer, activation memory grows with batch x sequence x hidden and, for the attention scores, with batch x heads x sequence squared unless an efficient kernel avoids storing them. Activation checkpointing stores a subset and recomputes the rest, at about one extra forward pass. Gradient accumulation reduces per-step activations at the cost of more steps.
Learning rates. Fine-tuning learning rates are usually one to two orders of magnitude below from-scratch values; with AdamW use decoupled weight decay and keep biases and normalisation scales out of it (Days 33-34). Warm-up protects the head's first steps from large updates.
Retention and drift. Fine-tuning can reduce performance on what the model did before. Keep a small retention set from the original task or behaviour, evaluate it before and after, and set an acceptance threshold in advance. Day 44 revisits this for parameter-efficient methods.
Reproducibility checklist. Seed every library, record the split identifier and file hash, library and driver versions, hardware, and the exact config. Enable deterministic algorithms only when needed; expect a speed cost and that results can still differ across platforms. The comparison that matters is between strategies under the same seeds, not a bit-exact rerun.
Hands-on lab: Probe, partial and full fine-tuning on a related task, with seed spread (scikit-learn digits)
Pretrain a small MLP on a 10-class digit task, then reuse it for a related target task (odd versus even digit) with 20, 50 or 150 labelled examples. Compare logistic regression on pixels, an MLP from scratch, a frozen pretrained layer with a new head, fine-tuning with a random head, and probe-then-fine-tune. Five seeds each. Read every difference against the spread. Then state what this tiny experiment cannot tell you about a 110M-parameter model on text.
import warnings
import numpy as np
from sklearn.datasets import load_digits
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPClassifier
warnings.filterwarnings("ignore")
X, y = load_digits(return_X_y=True)
X = X / 16.0
Xs, Xr, ys, yr = train_test_split(X, y, test_size=0.45, random_state=0, stratify=y) # source part / rest
Xpool, Xtest, ypool, ytest = train_test_split(Xr, yr % 2, test_size=0.5, random_state=0) # target: odd digit
def hidden(m, A):
return np.maximum(0, A @ m.coefs_[0] + m.intercepts_[0])
def run(seed, n):
idx = np.random.default_rng(seed).choice(len(Xpool), n, replace=False)
Xl, yl = Xpool[idx], ypool[idx]
pre = MLPClassifier(hidden_layer_sizes=(64,), max_iter=300, random_state=seed).fit(Xs, ys) # source task
out = {"logreg on pixels": LogisticRegression(max_iter=2000).fit(Xl, yl).score(Xtest, ytest)}
scratch = MLPClassifier(hidden_layer_sizes=(64,), max_iter=300, random_state=seed).fit(Xl, yl)
out["MLP from scratch"] = scratch.score(Xtest, ytest)
head = LogisticRegression(max_iter=2000).fit(hidden(pre, Xl), yl)
out["frozen features + new head"] = head.score(hidden(pre, Xtest), ytest)
for name, probe_first in (("fine-tune, random head", False), ("probe, then fine-tune", True)):
ft = MLPClassifier(hidden_layer_sizes=(64,), random_state=seed, learning_rate_init=1e-3)
ft.partial_fit(Xl, yl, classes=[0, 1])
ft.coefs_[0], ft.intercepts_[0] = pre.coefs_[0].copy(), pre.intercepts_[0].copy() # pretrained layer
if probe_first: # start from the already-trained head
ft.coefs_[1], ft.intercepts_[1] = head.coef_.T.copy(), head.intercept_.copy()
for _ in range(200):
ft.partial_fit(Xl, yl)
out[name] = ft.score(Xtest, ytest)
return out
for n in (20, 50, 150):
runs = [run(s, n) for s in range(5)]
print("labelled target examples:", n)
for k in runs[0]:
v = np.array([r[k] for r in runs])
print(f" {k:28s} {v.mean():.3f} +/- {v.std():.3f}")
Interview practice
In a hospital imaging group, a teammate reports that fine-tuning a pretrained backbone on chest X-ray crops "beats the linear probe by 2 points" from one run. What do you ask, and how do you decide?