Fine-Tuning the Encoder With Your Own Loop

Context

Fine-tuning a pretrained text encoder trains every one of its weights, together with a new linear head, on your own labeled text, with the loop you already wrote: a forward pass, a loss, a backward pass and an optimizer's step, batch after batch. AdamW takes the steps, and a schedule raises the learning rate from almost nothing and lowers it again across the run.

Every intent router, ticket tagger and complaint classifier built on a pretrained encoder is trained this way, whether the loop is written out or hidden inside a library's trainer. Writing it yourself puts four settings in plain view, the rate, the schedule, the mode the network is in and the epoch you keep, and those four decide whether the encoder learns your task or forgets what it knew.

Lesson 5.4.1's frozen BERT-Mini read Banking77 through features nobody trained for it and stayed at 0.7955, well under Lesson 5.3.1's bag of embeddings at 0.8760. Fine-tuning passes the bag, and it fails in two quiet ways on the way there: at 1e-3 BERT-Mini ends under the frozen head with its first attention layer's query weight moved by 61 percent of itself, and a validation pass that forgets model.eval() reads a different score every time, with no error, because BERT's dropout is still switching values off.

Goal

• Fine-tune all 11,170,560 parameters of BERT-Mini and a new nn.Linear(256, 77) head with a loop written by hand over encoder.batches, and save the reusable pieces as finetune.py.
• Explain from a ten-step run why AdamW's weight decay is decoupled from the gradient, and what Adam's folded-in decay does to a parameter whose gradient is zero.
• Measure three learning rates on text, read how far each moved three weight tensors, and name the rate that undoes the pretraining.
• Build a linear warmup and decay with LambdaLR, stepped once per batch, measure it against the constant rate with every other setting held, and keep the best epoch as encoder_ft.pt.
• Show that validation read in training mode is lower and changes between two passes over the same queries, and fix it with model.eval().
• Score the chosen network on the test queries once, set it beside the bag and the frozen head, and read where it wins, the queries holding words outside Lesson 5.3.1's vocabulary among them.

Prerequisites

None of this is required, and the lesson stands on its own. Both are here if you want to watch a fine-tuning loop written out by the library's own team first, or read one line by line with its optimizer and its scheduler.

• YouTube - "Write your training loop in PyTorch" by Hugging Face. The course's own walk through fine-tuning a pretrained model without the library's trainer: the data loaders, the optimizer, the learning-rate scheduler and the loop.
• Hugging Face LLM Course - Chapter 3, "A full training loop". The same fine-tune as a PyTorch loop on the page: AdamW at 5e-5, a linear schedule from get_scheduler that decays the rate to 0, model.train() before the training loop and model.eval() before the evaluation loop.

Table of Contents

• Fine-Tuning Every Weight
• AdamW and Decoupled Weight Decay
• The Learning Rate on Text
• Warmup and Decay With LambdaLR
• Dropout at Validation
• Testing Once Against the Bag and the Frozen Head
• Where the Fine-Tuned Network Wins

Theory

Every block runs in ~/dev/bootcamp, beside text.py, Lesson 5.3.1's, image_loop.py, Lesson 4.5.2's, and encoder.py and bert_features.pt, Lesson 5.4.1's, with BERT-Mini and its tokenizer already in the cache Lesson 5.4.1's first load filled, so nothing new downloads. The first subtitle saves finetune.py, which every later script imports. Each fine-tuning run builds its network afresh from the pretrained weights, seeds its head and its dropout with 42 and shuffles the training queries with a generator of its own seeded 42, so two runs at different settings start from the same weights and see the same batches in the same order. An epoch of fine-tuning takes roughly a minute on a processor, and the trained figures are this machine's on torch's default threads. A Python block that opens with its imports is a script of its own, and a block that does not continues the script above it.

Fine-Tuning Every Weight

What. Fine-tuning trains all 11,170,560 of BERT-Mini's parameters and a new seeded nn.Linear(256, 77) head together, every batch running Lesson 4.2.1's five steps: the forward pass through Lesson 5.4.1's pooled and the head, the loss, backward(), the optimizer's step and zero_grad(). The batches come from encoder.batches, 251 of 32 an epoch, shuffled by a seeded generator, and the loss is image_loop.py's loss_fn, Lesson 4.3.2's nn.CrossEntropyLoss. Why. Lesson 5.4.1's frozen head read features BERT-Mini learned by self-supervision, filling in hidden words, and some of what separates one bank intent from another was never part of that task. Fine-tuning lets every layer move toward the 77 intents, the step Lesson 5.2.2 took for ResNet-18 on satellite tiles. When. When a frozen head leaves room under a baseline, as Lesson 5.4.1's did under the bag, and an epoch that trains every weight is affordable.

Save the pieces every run shares as finetune.py:

from huggingface_hub.utils import logging as hub_logging
from transformers import logging
logging.set_verbosity_error()
logging.disable_progress_bar()
hub_logging.set_verbosity_error()

import torch
from torch import nn
from encoder import tokenizer, load, labels, batches, pooled
from text import train, val, test, intents
from image_loop import loss_fn

def build():
    model = load()             # pretrained, in evaluation mode
    torch.manual_seed(42)      # the head and every dropout draw
    head = nn.Linear(256, 77)  # 256 features in, a logit per intent
    return model, head

def train_epoch(model, head, optimizer, generator, scheduler=None, df=train):
    model.train()              # dropout on
    total = 0.0
    for batch, y in batches(df, 32, generator):
        loss = loss_fn(head(pooled(model, batch)), y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        if scheduler is not None:
            scheduler.step()   # once per batch
        total += loss.item() * len(y)
    return total / len(df)     # mean training loss

def evaluate(model, head, df):
    model.eval()               # dropout off
    correct = 0
    with torch.no_grad():
        for batch, y in batches(df, 250):
            correct += (head(pooled(model, batch)).argmax(dim=1) == y).sum().item()
    return correct / len(df)

The file opens with Lesson 5.4.1's five lines that set both libraries' messages to errors only, before encoder.py loads its tokenizer, so no script that imports it prints the load report or the Hub's message. image_loop.py's own train_epoch calls model(x) on a tensor, where here the tokenizer's batch is a dictionary that pooled unpacks into the encoder and the network is two modules, so the loop is written again with the same five steps. df=train is the default frame, and a run on other queries passes its own. The network before any step:

from finetune import build, evaluate
from text import val

model, head = build()
print(model.training, sum(p.numel() for p in model.parameters()), sum(p.numel() for p in head.parameters()))
print(f"{evaluate(model, head, val):.4f}")

False 11170560 19789
0.0215

build() returns BERT-Mini with training at False, because from_pretrained switches a loaded model to evaluation mode, which is why train_epoch opens with model.train(): without it every dropout layer would stay off for the whole fine-tune. The network trains 11,170,560 parameters in the encoder and 19,789 in the head, 256 times 77 weights and 77 biases. Before any step the seeded head calls 0.0215 of the validation queries right, 43 of 2,000, near a guess's one in 77.

DON'T:

• DON'T assume a loaded model is in training mode. build() returned BERT-Mini with training at False, and a loop without model.train() fine-tunes with every dropout layer off.

AdamW and Decoupled Weight Decay

What. torch.optim.AdamW is Lesson 4.3.4's Adam with the weight decay taken out of the gradient: each step first multiplies every parameter by 1 - lr * weight_decay, Lesson 4.4.1's shrink, then takes Adam's step on the loss's gradient alone. optim.Adam(..., weight_decay=wd) instead adds wd * p to the gradient, Lesson 4.4.1's SGD form, before Adam's two running averages see it. AdamW's default weight_decay is 0.01. Why. Adam divides each parameter's gradient by the root of its running squared average, so a decay folded into the gradient is divided too, and the shrink a parameter receives then depends on the rest of its gradient rather than on the parameter. Decoupled, every parameter loses the same share at every step, which is what weight decay meant for SGD. When. Whenever Adam's kind of step and a weight decay are both wanted, and in every fine-tune of a transformer, whose embedding table holds rows no batch looks up.

import torch
from torch import optim
from finetune import tokenizer, train

for opt_class, decay in ((optim.Adam, 0.0), (optim.Adam, 0.01), (optim.AdamW, 0.01)):
    w = torch.tensor([2.0, 2.0], requires_grad=True)
    optimizer = opt_class([w], lr=0.1, weight_decay=decay)
    for step in range(10):
        w.grad = torch.tensor([0.0, 100.0])  # no gradient, then a large one
        optimizer.step()
    print(opt_class.__name__, decay, w.detach())
print(optim.AdamW([w]).defaults["weight_decay"])
seen = set()
for ids in tokenizer(list(train["text"]))["input_ids"]:
    seen.update(ids)
print(len(seen), tokenizer.vocab_size)

Adam 0.0 tensor([2.0000, 1.0000])
Adam 0.01 tensor([1.0246, 1.0000])
AdamW 0.01 tensor([1.9801, 0.9846])
0.01
2258 30522

Without decay the first entry, whose gradient is 0, stays at 2.0, and the second moves by 0.1 a step to 1.0, Adam's step of size lr whatever the gradient, Lesson 4.3.4's. Adam's folded decay hands the first entry a gradient of 0.01 * 2.0, the only gradient it has, and Adam scales that up to a step of nearly lr, so ten steps take it to 1.0246, almost as far as the entry pushed by a gradient of 100. AdamW multiplies both entries by 1 - 0.1 * 0.01, 0.999, at every step, so the first ends at 2.0 times 0.999 to the tenth, 1.9801, and the second takes its gradient's steps on top of the same shrink. BERT-Mini's case is the first entry's: the 8,003 training queries look up 2,258 distinct token ids, [CLS] and [SEP] among them, of the 30,522 rows of its embedding table, so the other 28,264 rows receive a gradient of exactly zero at every step. Adam's folded decay would drag each of them toward zero by nearly the full rate at every step, and AdamW shrinks them by lr times 0.01 of themselves a step.

DON'T:

• DON'T pass weight_decay to optim.Adam and expect Lesson 4.4.1's shrink. Ten steps took an entry with no gradient from 2.0 to 1.0246, where AdamW shrank it to 1.9801.
• DON'T read an AdamW built without weight_decay= as an optimizer without decay. Its default is 0.01, and every AdamW in this lesson applies it.

The Learning Rate on Text

What. The same fine-tune run three times, three epochs each, at 1e-4, the rate Lesson 5.2.2 fine-tuned ResNet-18 at, at 3e-4, and at 1e-3, the rate Lesson 5.2.3 found undoing ImageNet's pretraining under Adam. Each run is read by its validation accuracy per epoch and by how far three weight tensors moved from their pretrained values, Lesson 5.2.2's measure: the word embedding table, the first layer's query projection and the last layer's output projection. Why. A rate belongs to a network, an optimizer and a dataset together, so the rate that served a vision fine-tune is where a text one starts looking, not its answer. Too small and three epochs end before the network has adapted, too large and the steps overwrite the weights the pretraining paid for. When. Before any fine-tune whose rate was copied from another task, and again whenever the model, the optimizer or the data changes.

import torch
from finetune import build, train_epoch, evaluate, val

WATCHED = ("embeddings.word_embeddings.weight", "encoder.layer.0.attention.self.query.weight", "encoder.layer.3.output.dense.weight")
for lr in (1e-4, 3e-4, 1e-3):
    model, head = build()
    # clone: a state dict shares the parameters' memory, Lesson 4.2.3's
    start = {n: model.state_dict()[n].clone() for n in WATCHED}
    optimizer = torch.optim.AdamW(list(model.parameters()) + list(head.parameters()), lr=lr)
    generator = torch.Generator().manual_seed(42)
    for epoch in range(1, 4):
        loss = train_epoch(model, head, optimizer, generator)
        print(lr, epoch, f"{loss:.4f}", f"{evaluate(model, head, val):.4f}")
    now = model.state_dict()
    print([round(((now[n] - start[n]).norm() / start[n].norm()).item(), 4) for n in WATCHED])

0.0001 1 2.6712 0.7130
0.0001 2 1.0019 0.8365
0.0001 3 0.5389 0.8665
[0.0211, 0.0761, 0.1015]
0.0003 1 1.8670 0.8135
0.0003 2 0.5369 0.8575
0.0003 3 0.3108 0.8795
[0.0453, 0.1614, 0.1562]
0.001 1 1.8047 0.7060
0.001 2 0.8379 0.7420
0.001 3 0.8136 0.7065
[0.1426, 0.6134, 0.3395]

The three runs share the pretrained weights, the seeded head and the shuffles, so the rate is their one difference. At 1e-4 validation climbs to 0.8665 and is still rising after three epochs. At 3e-4 it leads at every epoch and reaches 0.8795, with the first layer's query projection moved by 16.1 percent of itself. At 1e-3 it rises to 0.7420 and falls back to 0.7065, under Lesson 5.4.1's frozen head at 0.7955, its training loss stuck at 0.8379 and 0.8136 over the last two epochs where 3e-4's fell to 0.3108, and it moved the query projection by 61.3 percent of itself and the whole word table by 14.3 percent: the steps overwrote what pretraining had put there, the failure Lesson 5.2.3 found in vision. On these queries, with AdamW and three epochs, the best of the three rates is 3e-4, and Lesson 5.2.2's 1e-4 was a place to start rather than the answer.

DON'T:

• DON'T copy a fine-tuning rate from a vision run. Lesson 5.2.2's 1e-4 trailed 3e-4 at every epoch here.
• DON'T keep training at a rate whose training loss has stopped falling. At 1e-3 it sat at 0.8379 and 0.8136 while validation fell back to 0.7065.
• DON'T read the smallest rate as the safe one. 1e-4 moved every watched tensor least and ended 1.3 points under 3e-4, its three epochs spent before it had adapted.

Warmup and Decay With LambdaLR

What. LambdaLR(optimizer, f) is a scheduler, Lesson 4.3.4's Assignment's object that rewrites the optimizer's lr on a timetable, there StepLR. It sets each group's rate to the lr the optimizer was built with times f(step), where step counts its own step() calls from 0. Here f rises linearly from 1/75 to 1 over the first 75 steps, a warmup, then falls linearly to 0 at step 753, the end of three epochs of 251 batches, a decay, and the scheduler is stepped once per batch. Why. At the first steps the head is a seeded draw, so its gradients are large and point nowhere useful, and a full rate lets them push every layer below it, the reason Lesson 5.2.3 tried a warm head. A warmup keeps those first steps small without a head trained beforehand, and the decay makes the last steps small, so the run ends settled instead of still moving. When. In any fine-tune of a pretrained transformer, with the warmup a small share of the steps, checked against the constant rate on the task's own validation queries.

import copy
import torch
from torch import optim
from torch.optim.lr_scheduler import LambdaLR
from finetune import build, train_epoch, evaluate, val

STEPS, WARMUP = 3 * 251, 75

def warmup_decay(step):
    return min((step + 1) / WARMUP, (STEPS - step) / (STEPS - WARMUP))

w = torch.zeros(1, requires_grad=True)
trial = optim.AdamW([w], lr=3e-4)
scheduler = LambdaLR(trial, warmup_decay)
rates = []
for step in range(STEPS):
    rates.append(trial.param_groups[0]["lr"])
    trial.step()  # no gradient, so nothing moves
    scheduler.step()
print(len(rates), [f"{rates[i]:.2e}" for i in (0, 1, 37, 74, 75, 251, 502, 752)])
print(f"{sum(rates) / len(rates):.2e}")

753 ['4.00e-06', '8.00e-06', '1.52e-04', '3.00e-04', '3.00e-04', '2.22e-04', '1.11e-04', '4.42e-07']
1.50e-04

trial is an optimizer over one tensor that never gets a gradient, so its 753 steps change nothing and the scheduler prints its timetable without training anything. The rate starts at 3e-4 over 75, 4.00e-06, reaches the peak at step 74, where the warmup term reaches 1, and holds it at 75, where the decay term starts from 1, is 2.22e-04 when the second epoch starts at step 251 and 1.11e-04 when the third starts at 502, and ends at 4.42e-07. Averaged over the 753 steps it is 1.50e-04, half the peak. The fine-tune at a peak of 3e-4, the sweep's best constant rate, keeping the best epoch by validation with Lesson 4.4.2's copy.deepcopy, of both halves of the network:

model, head = build()
optimizer = optim.AdamW(list(model.parameters()) + list(head.parameters()), lr=3e-4)
scheduler = LambdaLR(optimizer, warmup_decay)
generator = torch.Generator().manual_seed(42)
best_acc, best_state = 0.0, None
for epoch in range(1, 4):
    loss = train_epoch(model, head, optimizer, generator, scheduler)
    val_acc = evaluate(model, head, val)
    print(epoch, f"{loss:.4f}", f"{val_acc:.4f}", f"{scheduler.get_last_lr()[0]:.2e}")
    if val_acc > best_acc:
        best_acc, best_state = val_acc, copy.deepcopy({"encoder": model.state_dict(), "head": head.state_dict()})
torch.save(best_state, "encoder_ft.pt")
print(f"{best_acc:.4f}", scheduler.last_epoch)

1 2.2182 0.8125 2.22e-04
2 0.5518 0.8715 1.11e-04
3 0.2804 0.8955 0.00e+00
0.8955 753

The sweep's 3e-4 run read 0.8135, 0.8575 and 0.8795 from the same pretrained weights, the same seeded head and the same shuffles at the same peak rate, so the schedule is the one setting the two runs do not share. It cost a tenth of a point after the first epoch, where the warmup held the rate low and the training loss ended at 2.2182 against 1.8670, and bought 1.4 points after the second and 1.6 after the third, 0.8955 against 0.8795. The schedule itself changes three things at once, the warmup, the decay and a mean rate of half the constant one, and these two runs cannot say which of the three the 1.6 points belong to. scheduler.last_epoch counts the scheduler's steps, 753 here and not epochs, since this one is stepped per batch. Every epoch set a new best, so best_state holds the third, and encoder_ft.pt holds a dictionary of two state dicts, "encoder" and "head": the chosen network for the rest of the lesson.

DON'T:

• DON'T step a per-batch schedule once per epoch, the way Lesson 4.3.4's StepLR was stepped. This warmup reaches its peak after 74 steps, which would be 74 epochs.
• DON'T credit the 1.6 points to the warmup. The scheduled run changed the warmup, the decay and the mean rate together, and only a run that changes one of them can name it.
• DON'T save the encoder's state dict alone. The head trained with it, and encoder_ft.pt keeps both under their own keys so a later script rebuilds the network the run chose.

Dropout at Validation

What. BERT-Mini holds 13 nn.Dropout layers, Lesson 4.4.2's, each at 0.1: one after the embeddings, and three in each of its four layers, on the attention weights, after the attention's output projection and after the feed-forward network. Two validation passes over the same 2,000 queries with the fine-tuned network in training mode, then two in evaluation mode, count the answers that change between the passes of each pair. Why. In training mode every dropout layer zeroes a fresh random tenth of its inputs on every call, so a pass in that mode scores a different damaged copy of the network each time. Nothing raises, because the arithmetic is valid in both modes, and train_epoch leaves the network in training mode, so the pass that follows a training epoch is exactly where the mode is wrong. When. Before every validation, test or served prediction from a network that holds dropout or BatchNorm, and above all straight after a training epoch.

import torch
from torch import nn
from finetune import build, evaluate, batches, pooled, labels, val

model, head = build()
state = torch.load("encoder_ft.pt", weights_only=True)
print(model.load_state_dict(state["encoder"]), head.load_state_dict(state["head"]))
print(sum(isinstance(m, nn.Dropout) for m in model.modules()), model.config.hidden_dropout_prob, model.config.attention_probs_dropout_prob)

def predictions(df):
    with torch.no_grad():
        return torch.cat([head(pooled(model, batch)).argmax(dim=1) for batch, y in batches(df, 250)])

Z = labels(val)
model.train()
first, second = predictions(val), predictions(val)
print(f"{(first == Z).float().mean().item():.4f}", f"{(second == Z).float().mean().item():.4f}", (first != second).sum().item())
model.eval()
first, second = predictions(val), predictions(val)
print(f"{(first == Z).float().mean().item():.4f}", f"{(second == Z).float().mean().item():.4f}", (first != second).sum().item())
print(f"{evaluate(model, head, val):.4f}")

<All keys matched successfully> <All keys matched successfully>
13 0.1 0.1
0.8890 0.8895 100
0.8955 0.8955 0
0.8955

Both halves of encoder_ft.pt load with every key matched, Lesson 4.4.3's report, and config holds the two dropout rates, 0.1 on the hidden values and 0.1 on the attention weights. In training mode the same network over the same queries reads 0.8890 and then 0.8895, and 100 of the 2,000 answers differ between the two passes. In evaluation mode both passes read 0.8955 and no answer changes, the score evaluate reads and the training run kept. Left in training mode, a validation pass here reads 0.65 and 0.60 of a point low and moves on its own from pass to pass, so two settings compared through such passes are compared through dropout's draws as much as through the settings, the flag Lesson 4.4.3 met on one listing priced three ways.

DON'T:

• DON'T read validation straight after train_epoch without model.eval(). The same network read 0.8890 and 0.8895 and changed 100 answers between two passes over the same queries.
• DON'T compare two settings by scores read in training mode. A reading that moves by itself between passes cannot separate a setting's effect from the dropout draws.
• DON'T trust the mode a network arrives in. from_pretrained returned BERT-Mini in evaluation mode, train_epoch leaves it in training mode, and evaluate leaves it in evaluation mode for whatever runs next.

Testing Once Against the Bag and the Frozen Head

What. One script retrains the two cheaper models by their own lessons' recipes: Lesson 5.3.1's bag of embeddings, dimension 64, Adam at 0.01 and ten epochs of image_loop.py's train_epoch, with order and the weights both seeded 42 right before it is built, and Lesson 5.4.1's frozen head, a seeded nn.Linear(256, 77) trained on the features in bert_features.pt for 30 epochs of Adam at 0.001 by that lesson's loop. It then loads encoder_ft.pt, reads its validation score beside theirs, and scores the 3,080 test queries once with it. Why. A fine-tune is worth its cost only against what cheaper models reach on the same validation queries, and these two are the cheaper ones. The validation queries chose the rate, the schedule and the epoch, so only the test queries can say what the chosen network does on queries no decision touched, Lesson 3.1.2's rule, scored once as Lesson 5.2.3 scored its own choice. When. Once, after every choice has been made on the validation queries.

import torch
from torch import nn, optim
from text import vocab, order, train_loader, val_loader, val_set
from image_loop import train_epoch, accuracy, loss_fn
from finetune import build, evaluate, batches, pooled, tokenizer, intents, val, test

class Bag(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.embed = nn.Embedding(len(vocab), dim, padding_idx=0)
        self.head = nn.Linear(dim, len(intents))

    def forward(self, ids):
        real = (ids != 0).sum(dim=1, keepdim=True)
        return self.head(self.embed(ids).sum(dim=1) / real)

order.manual_seed(42)
torch.manual_seed(42)
bag = Bag(64)
optimizer = optim.Adam(bag.parameters(), lr=0.01)
for epoch in range(10):
    train_epoch(bag, optimizer, train_loader)
print(f"{accuracy(bag, val_loader):.4f}")

cached = torch.load("bert_features.pt", weights_only=True)
F, Y, G, Z = cached["F"], cached["Y"], cached["G"], cached["Z"]
torch.manual_seed(42)
frozen = nn.Linear(256, 77)
optimizer = optim.Adam(frozen.parameters(), lr=0.001)
shuffle = torch.Generator().manual_seed(42)
for epoch in range(30):
    perm = torch.randperm(len(F), generator=shuffle)
    for i in range(0, len(F), 32):
        loss = loss_fn(frozen(F[perm[i:i + 32]]), Y[perm[i:i + 32]])
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
with torch.no_grad():
    print(f"{(frozen(G).argmax(dim=1) == Z).float().mean().item():.4f}")

model, head = build()
state = torch.load("encoder_ft.pt", weights_only=True)
model.load_state_dict(state["encoder"])
head.load_state_dict(state["head"])
print(f"{evaluate(model, head, val):.4f}")
print(f"{evaluate(model, head, test):.4f}")

0.8760
0.7955
0.8955
0.8990

Bag is Lesson 5.3.1's class unchanged, and the bag's 0.8760 and the frozen head's 0.7955 are their lessons' figures to the digit, so both recipes are the same training. On the validation queries, with every model of this section beside them:

The fine-tune with the schedule is the choice, encoder_ft.pt, and the 3,080 test queries score it once at 0.8990, a little above its validation 0.8955, as Lesson 5.3.3's test score of the bag, 0.8808, sat a little above the bag's 0.8760. Against the bag, on the queries each lesson's choice was scored on once, the fine-tune leads by 1.8 points on the test queries and by 1.95 on validation, with 11,190,349 trained parameters against 145,933, about 77 times as many. Against Lesson 5.3.3's encoder, built from the same parts and trained from scratch, the lead is 6.6 points, and those two differ in their pretraining, their tokenizer, their size, their optimizer, their rate and their schedule, so the table cannot credit the gap to any one of them.

DON'T:

• DON'T score the bag, the frozen head and the fine-tune on the test queries to choose among them. The validation queries chose, and the test queries scored the choice once, at 0.8990.
• DON'T credit the gap over Lesson 5.3.3's encoder to pretraining alone. The two runs differ in six settings, and no run here changes one of them alone.

Where the Fine-Tuned Network Wins

What. The bag's and the fine-tuned network's answers on the same 2,000 validation queries, read side by side: on the 145 that hold a word outside Lesson 5.3.1's vocabulary and on the other 1,855, query by query, intent by intent, and on Lesson 5.3.1's own probe queries, revoke, rescind, fiver dosh quid and the two money-moving queries in opposite orders. Why. One accuracy each says the fine-tune leads and nothing about where. Lesson 5.3.1 traced the bag's failures to <unk> and to word order, and BERT-Mini is built against both, with Lesson 5.4.1's subword tokenizer and Lesson 5.3.3's position embeddings, so those are the places to look first, and the counts say whether the lead is where the design says it should be. When. Before claiming why one model beats another, and before choosing which errors the next hour goes to.

X = val_set.tensors[0]
has_unk = (X == 1).any(dim=1)
bag.eval()
model.eval()
with torch.no_grad():
    bag_right = bag(X).argmax(dim=1) == Z
    tuned_right = torch.cat([head(pooled(model, batch)).argmax(dim=1) for batch, y in batches(val, 250)]) == Z
for name, right in (("bag", bag_right), ("tuned", tuned_right)):
    print(name, has_unk.sum().item(), f"{right[has_unk].float().mean().item():.4f}", f"{right[~has_unk].float().mean().item():.4f}")
print((bag_right & ~tuned_right).sum().item(), (~bag_right & tuned_right).sum().item(), (~bag_right & ~tuned_right).sum().item())

def predict(query):
    batch = tokenizer([query], return_tensors="pt")
    with torch.no_grad():
        probs = head(pooled(model, batch)).softmax(dim=1)[0]
    return intents[probs.argmax()], round(probs.max().item(), 3)

for query in ("how do i cancel a transfer", "how do i revoke a transfer", "how do i rescind a transfer", "fiver dosh quid", "can i move money from my card to my account", "can i move money from my account to my card"):
    print(tokenizer.tokenize(query), predict(query))
per = []
for k in range(77):
    mine = Z == k
    per.append((tuned_right[mine].float().mean().item() - bag_right[mine].float().mean().item(), intents[k], mine.sum().item()))
per.sort()
print([(n, c, round(d, 3)) for d, n, c in per[:3]], [(n, c, round(d, 3)) for d, n, c in per[-3:]])

bag 145 0.7655 0.8846
tuned 145 0.7931 0.9035
81 120 128
['how', 'do', 'i', 'cancel', 'a', 'transfer'] ('cancel_transfer', 0.984)
['how', 'do', 'i', 'rev', '##oke', 'a', 'transfer'] ('cancel_transfer', 0.9)
['how', 'do', 'i', 'res', '##cin', '##d', 'a', 'transfer'] ('failed_transfer', 0.359)
['five', '##r', 'dos', '##h', 'qui', '##d'] ('direct_debit_payment_not_recognised', 0.108)
['can', 'i', 'move', 'money', 'from', 'my', 'card', 'to', 'my', 'account'] ('topping_up_by_card', 0.68)
['can', 'i', 'move', 'money', 'from', 'my', 'account', 'to', 'my', 'card'] ('topping_up_by_card', 0.756)
[('virtual_card_not_working', 9, -0.444), ('verify_my_identity', 22, -0.182), ('disposable_card_limits', 30, -0.133)] [('declined_card_payment', 26, 0.154), ('get_physical_card', 26, 0.154), ('unable_to_verify_identity', 21, 0.238)]

X is the validation queries encoded by Lesson 5.3.1's vocabulary, in the same order encoder.batches reads val in, so the two models' answers line up query by query, and model.eval() stands before the hand-read predictions even though evaluate left the network in that mode. The bag's two figures are Lesson 5.3.1's own, 0.7655 on the 145 queries with an unknown word and 0.8846 on the rest. The fine-tune scores 0.7931 and 0.9035: it leads on both groups, by 2.76 points where the bag's vocabulary runs out and by 1.89 where it does not, and its own unknown-word queries still sit 11 points under its others. The lead is a net figure: 120 queries the fine-tune gets right and the bag misses, 81 the other way round, and 128 both miss, so 39 queries, 1.95 points of 2,000, and the bag still beats it on 81.

The probes show what the pieces buy and what they do not. revoke is two pieces, rev and ##oke, and the query comes back cancel_transfer at 0.9, where the bag, which read revoke as <unk>, answered transfer_into_account at 0.578. rescind is three pieces, and comes back failed_transfer at 0.359: a word cut into stored pieces is not a word understood, and the answer is wrong with less than half the probability on it. fiver dosh quid is six pieces and comes back at 0.108, a hesitant answer where the bag's three copies of the untrained <unk> row gave card_acceptance at 0.62. The two queries that move money in opposite directions now get different probabilities, 0.68 and 0.756, so the network separates the two inputs the bag could not, and it still files both under topping_up_by_card. Intent by intent, the largest gain is unable_to_verify_identity, up 23.8 points on 21 queries, and the largest fall virtual_card_not_working, down 44.4 points on 9, and each of the six differences printed is 4 or 5 queries: per intent the two models differ by a handful.

DON'T:

• DON'T credit the whole gain to the subword tokenizer. The fine-tune led by 1.89 points on queries the bag's vocabulary covers fully, and by 2.76 on those it does not.
• DON'T read a word cut into known pieces as a word understood. rescind became res, ##cin and ##d and came back failed_transfer at 0.359.
• DON'T read a lead of 1.95 points as the fine-tune being right wherever the bag is. The bag was right on 81 validation queries the fine-tune missed.

Lab Exercises

Lab Exercise 1: A Schedule Stepped per Epoch

In per_epoch.py, build the network, an AdamW at 3e-4 over both halves, and the fourth subtitle's LambdaLR with the same warmup_decay, STEPS and WARMUP. Train one epoch with train_epoch given no scheduler, then call scheduler.step() once, the way Lesson 4.3.4's StepLR was stepped once per epoch. Print the rate before the epoch, then the training loss, the validation accuracy, the rate after the step and scheduler.last_epoch. Say in a comment how this epoch compares with the scheduled run's first epoch, how many epochs this warmup would take to reach the peak, and what one change fixes it.

Lab Exercise 2: Serving the Chosen Network

In serve.py, build the network, keep a copy of build()'s head weight, and load only the "encoder" half of encoder_ft.pt. Print the validation accuracy with build()'s own head, then with a second nn.Linear(256, 77) built without a seed of its own. Load the "head" half, and print the validation accuracy again beside how far the trained head's weight sits from build()'s draw, the size of the change over the draw's size. Then write predict(queries), which returns each query's top intent and its probability in evaluation mode, and run it on four misspelled queries, once as one padded batch and once a query at a time. For each query print the tokens Lesson 5.3.1's vocabulary turns into <unk>, BERT-Mini's tokens, the batched answer, and whether it equals the lone one. Say in a comment why the first two accuracies differ and what a serving script has to load.

import torch
from torch import optim
from torch.optim.lr_scheduler import LambdaLR
from finetune import build, train_epoch, evaluate, val

STEPS, WARMUP = 3 * 251, 75

def warmup_decay(step):
    return min((step + 1) / WARMUP, (STEPS - step) / (STEPS - WARMUP))

model, head = build()
optimizer = optim.AdamW(list(model.parameters()) + list(head.parameters()), lr=3e-4)
scheduler = LambdaLR(optimizer, warmup_decay)
generator = torch.Generator().manual_seed(42)
print(f"{scheduler.get_last_lr()[0]:.2e}")
loss = train_epoch(model, head, optimizer, generator)
scheduler.step()
print(f"{loss:.4f}", f"{evaluate(model, head, val):.4f}", f"{scheduler.get_last_lr()[0]:.2e}", scheduler.last_epoch)
# One epoch at 4e-6 reads 0.0305 against the scheduled run's 0.8125:
# the warmup reaches the peak after 74 steps, 74 epochs at this pace.
# The fix is the one line inside train_epoch, scheduler.step() per batch.

4.00e-06
4.3275 0.0305 8.00e-06 1

import torch
from torch import nn
from text import tokenize, encode
from finetune import build, evaluate, pooled, tokenizer, intents, val

model, head = build()
start = head.weight.detach().clone()
state = torch.load("encoder_ft.pt", weights_only=True)
model.load_state_dict(state["encoder"])
print(f"{evaluate(model, head, val):.4f}")
print(f"{evaluate(model, nn.Linear(256, 77), val):.4f}")
head.load_state_dict(state["head"])
print(f"{evaluate(model, head, val):.4f}", f"{((head.weight - start).norm() / start.norm()).item():.4f}")

def predict(queries):
    model.eval()
    batch = tokenizer(queries, padding=True, return_tensors="pt")
    with torch.no_grad():
        probs = head(pooled(model, batch)).softmax(dim=1)
    return [(intents[p.argmax()], round(p.max().item(), 3)) for p in probs]

queries = ["my crad was declined at the shop", "i lost my wallett", "how long does a transfr take", "why is there an extra fee on my statment"]
together = predict(queries)
alone = [predict([q])[0] for q in queries]
for q, a, b in zip(queries, together, alone):
    print([t for t, i in zip(tokenize(q), encode(q)) if i == 1], tokenizer.tokenize(q), a, a == b)
# build()'s head is the training run's own starting draw, seed 42, which
# the encoder learned to write for, so it reads 0.8975. An unseeded draw
# read 0.0020. A serving script loads both halves of encoder_ft.pt.

0.8975
0.0020
0.8955 0.3961
['crad'] ['my', 'cr', '##ad', 'was', 'declined', 'at', 'the', 'shop'] ('declined_transfer', 0.618) True
['wallett'] ['i', 'lost', 'my', 'wallet', '##t'] ('lost_or_stolen_card', 0.574) True
['transfr'] ['how', 'long', 'does', 'a', 'trans', '##fr', 'take'] ('transfer_timing', 0.287) True
['statment'] ['why', 'is', 'there', 'an', 'extra', 'fee', 'on', 'my', 'stat', '##ment'] ('extra_charge_on_statement', 0.412) True

Theory Assessment

Q1: The Decay Adam Divides

Ten steps of optim.Adam(..., weight_decay=0.01) took an entry whose gradient is zero from 2.0 to 1.0246, and optim.AdamW at the same setting took it to 1.9801. Say what each optimizer does with the decay in its step, why Adam's moved that entry almost as far as a gradient of 100 moved the other, and which of BERT-Mini's parameters sit in that entry's place during this lesson's fine-tune.

Hint: Adam's step divides by the root of the running squared gradient. Ask what that gradient is for an entry whose only gradient is the decay.

Q2: Three Rates, One Difference

At 1e-4, 3e-4 and 1e-3 the fine-tune read 0.8665, 0.8795 and 0.7065 after three epochs and moved the first layer's query projection by 7.6, 16.1 and 61.3 percent of itself. Say what the three runs share, what 1e-3's training loss and moved figures show, why 1e-4 lost to 3e-4 although it moved the weights least, and what the sweep says and does not say about Lesson 5.2.2's 1e-4.

Hint: Hold each run's last two training losses side by side, and set 1e-3's end beside the frozen head's 0.7955.

Q3: Two Passes, 100 Answers

In training mode the fine-tuned network read 0.8890 and 0.8895 on the same 2,000 validation queries and changed 100 answers between the two passes, and in evaluation mode it read 0.8955 twice. Say what differs between the two passes in training mode, why neither raises, what the moving reading does to a comparison of two settings, and where in this lesson's code the mode would be wrong if evaluate did not set it.

Hint: Count BERT-Mini's dropout layers and what each does on every call in training mode, then ask which mode train_epoch leaves the network in.

Q4: The Schedule's 1.6 Points

A teammate sets the scheduled run's 0.8955 beside the constant 3e-4 run's 0.8795, concludes that warmup is worth 1.6 points, puts a 75-step warmup with no decay into the next project's constant-rate script, and, to confirm, scores both of this lesson's networks on the test queries. Say what the two runs share and what separates them, why the 1.6 points cannot be credited to the warmup, what run would measure the warmup alone, and what scoring both networks on the test queries does.

Hint: Read the scheduled run's rate when the second epoch starts and its mean over the 753 steps, against the constant run's 3e-4 at every step.

Assignment

Assignment: Ten Queries an Intent

Transfer learning earns its keep where labels are few, and Banking77's 8,003 training queries are a lot of labels for 77 intents. Write few.py, which keeps 10 training queries per intent, 770 in all, and measures what that does to the bag and to the fine-tune.

Two things this lesson did not use. train.groupby("category").sample(n=10, random_state=42) draws 10 rows from every group of the frame, each group's draw seeded, so every intent keeps exactly 10 queries. And the bag's vocabulary must be rebuilt from those 770 queries alone, with Lesson 5.3.1's tokenize, <pad> at 0, <unk> at 1 and the rest by count, most frequent first, and every query encoded by it and padded or cut to MAX_LEN.

Print the number of queries, intents and vocabulary entries, and how many of the 2,000 validation queries hold a token that vocabulary lacks. Train the bag on the 770 at dimension 64 with Adam at 0.01, in batches of 32 shuffled by a generator seeded 42, the weights seeded 42 right before it is built, for 30 epochs, printing validation every 5 epochs and the best epoch with its score. Then fine-tune BERT-Mini on the same 770 for 10 epochs with AdamW at a peak of 3e-4 and the warmup and decay of this lesson laid over the run's own steps, a tenth of them warmup, printing each epoch's training loss and validation accuracy and then the best. Say in a comment how each model's best compares with its run on all 8,003 queries, and why the gap between them changed.

Hints:

• 770 queries in batches of 32 make 25 batches an epoch, so 10 epochs are 250 scheduler steps and a tenth of them is 25, and train_epoch takes the frame to train on through df=.
• image_loop.py's train_epoch and finetune.py's share a name, and from image_loop import train_epoch as bag_epoch imports the first under a name of its own.

from collections import Counter
import torch
from torch import nn, optim
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader, TensorDataset
from text import tokenize, MAX_LEN
from image_loop import accuracy, train_epoch as bag_epoch
from finetune import build, train_epoch, evaluate, train, val, intents

few = train.groupby("category").sample(n=10, random_state=42)
counts = Counter(t for text in few["text"] for t in tokenize(text))
small = {"<pad>": 0, "<unk>": 1}
for token, n in counts.most_common():
    small[token] = len(small)

def tensors(df):
    rows = []
    for text in df["text"]:
        ids = [small.get(t, 1) for t in tokenize(text)][:MAX_LEN]
        rows.append(ids + [0] * (MAX_LEN - len(ids)))
    return torch.tensor(rows), torch.tensor([intents.index(c) for c in df["category"]])

X, Y = tensors(few)
V, Z = tensors(val)
print(len(few), few["category"].nunique(), len(small), (V == 1).any(dim=1).sum().item())

class Bag(nn.Module):
    def __init__(self, size, dim):
        super().__init__()
        self.embed = nn.Embedding(size, dim, padding_idx=0)
        self.head = nn.Linear(dim, len(intents))

    def forward(self, ids):
        real = (ids != 0).sum(dim=1, keepdim=True)
        return self.head(self.embed(ids).sum(dim=1) / real)

loader = DataLoader(TensorDataset(X, Y), batch_size=32, shuffle=True, generator=torch.Generator().manual_seed(42))
val_loader = DataLoader(TensorDataset(V, Z), batch_size=250)
torch.manual_seed(42)
bag = Bag(len(small), 64)
optimizer = optim.Adam(bag.parameters(), lr=0.01)
best = (0.0, 0)
for epoch in range(1, 31):
    bag_epoch(bag, optimizer, loader)
    acc = accuracy(bag, val_loader)
    if acc > best[0]:
        best = (acc, epoch)
    if epoch % 5 == 0:
        print("bag", epoch, f"{acc:.4f}")
print("bag best", f"{best[0]:.4f}", best[1])

770 77 915 643
bag 5 0.5390
bag 10 0.6290
bag 15 0.6410
bag 20 0.6470
bag 25 0.6500
bag 30 0.6505
bag best 0.6555 29

STEPS, WARMUP = 10 * 25, 25
model, head = build()
optimizer = optim.AdamW(list(model.parameters()) + list(head.parameters()), lr=3e-4)
scheduler = LambdaLR(optimizer, lambda s: min((s + 1) / WARMUP, (STEPS - s) / (STEPS - WARMUP)))
generator = torch.Generator().manual_seed(42)
best = (0.0, 0)
for epoch in range(1, 11):
    loss = train_epoch(model, head, optimizer, generator, scheduler, few)
    acc = evaluate(model, head, val)
    print("bert", epoch, f"{loss:.4f}", f"{acc:.4f}")
    if acc > best[0]:
        best = (acc, epoch)
print("bert best", f"{best[0]:.4f}", best[1])
# On 770 queries the bag falls from 0.8760 to 0.6555 and the fine-tune
# from 0.8955 to 0.7120, a gap of 5.65 points where all 8,003 gave 1.95:
# the bag's words and their vectors come from the 770 alone, while
# BERT-Mini's 30,522 pieces and their meanings came from pretraining.

bert 1 4.3587 0.0910
bert 2 3.6137 0.3825
bert 3 2.4857 0.5650
bert 4 1.7503 0.6150
bert 5 1.2929 0.6645
bert 6 0.9742 0.6815
bert 7 0.7703 0.6960
bert 8 0.6163 0.7060
bert 9 0.5274 0.7050
bert 10 0.4880 0.7120
bert best 0.7120 10

Additional Resources

Optional, and nothing later depends on either. Worth opening for the paper behind AdamW, or for the one behind BERT-Mini itself.

• Paper - Loshchilov and Hutter, "Decoupled Weight Decay Regularization". Why L2 regularization and weight decay are one thing for SGD and two things for Adam, and the decoupled update that became AdamW.
• Paper - Turc, Chang, Lee and Toutanova, "Well-Read Students Learn Better: On the Importance of Pre-training Compact Models". Pre-training and fine-tuning compact BERT models, and the 24 miniature BERTs released with the paper, BERT-Mini among them.

Summary

Nothing in the loop was new. The five steps of Module 4 ran over a pretrained encoder exactly as they ran over a single neuron, and what made fine-tuning work or fail lived in the settings around them: an optimizer that shrinks a weight no batch touches by a thousandth instead of nearly the full rate, a rate that adapted BERT-Mini at 3e-4 and overwrote it at 1e-3, a timetable for that rate, and a mode switch. Each was measured against a run that differed from it in that one setting, and the one comparison whose difference held several changes, the schedule, was read as the timetable it was rather than as a warmup.

The failures were quiet. A rate that undid the pretraining ended under the frozen head without an error, a validation pass in training mode moved by itself between two readings of the same queries, a schedule stepped per epoch spent a whole epoch at a 75th of its rate, and a serving script that forgot the head read 0.8975 only because it shared a seed with training. In each case the network ran, printed a number and looked like a model, and only a reading set beside another one showed what had happened.

And the gain was smaller, and more honest, than the architecture suggests. On 8,003 queries the fine-tune led the bag by under two points, with the bag still right on 81 queries the fine-tune missed, and its subword pieces turned revoke into a correct answer and rescind into a wrong one. With 770 queries the lead nearly tripled, because the pretrained weights carry what a small label set cannot teach. That is what a pretrained encoder is for: less a better reader of plentiful data than a way to need less of it.

Cheat Sheet

Part A - Reference Table

Part B - Compact Code Block

STEPS, WARMUP = 3 * 251, 75

def warmup_decay(step):
    return min((step + 1) / WARMUP, (STEPS - step) / (STEPS - WARMUP))

optimizer = optim.AdamW(params, lr=3e-4)  # decay 0.01 by default
scheduler = LambdaLR(optimizer, warmup_decay)
for batch, y in batches(train, 32, generator):
    loss = loss_fn(head(pooled(model, batch)), y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    scheduler.step()                      # every batch, not every epoch

model.eval()  # dropout off before any reading
torch.save({"encoder": model.state_dict(), "head": head.state_dict()}, "encoder_ft.pt")