Classification Loss, Accuracy & the Fine-Tuning Run

The surgically modified model outputs 2 logits at its last token. This unit implements the evaluation utilities (step 7 of the chapter’s 10-step plan), then runs the fine-tuning (steps 8–9) — five epochs to a 95%+ spam detector.

6.6 From logits to a verdict

Prediction works exactly like next-token selection — softmax then argmax — except over 2 outputs instead of 50,257. The book’s p. 102 figure, with its real numbers:

The book’s p. 102 pipeline: the same softmax→argmax mechanics as text generation, over 2 classes. Note the untrained model’s confident mistake on the lottery scam.

Explore the first example’s logits — with only two classes, softmax reduces to a tug-of-war:

last-token logits for "You won the lottery" (untrained head)sum = -0.39 (arbitrary)

not spam

3.5983

spam

-3.9902

softmax(x)sum = 1, all positive

not spam

0.9995

spam

0.0005

softmax([3.5983, −3.9902]) ≈ [0.99, 0.01] → argmax picks index 0, «not spam». The softmax is optional for the decision: argmax on the raw logits gives the same answer.

</> Prediction and accuracy over a data loader

logits = outputs[:, -1, :]  #A
label = torch.argmax(logits)
print("Class label:", label.item())

def calc_accuracy_loader(data_loader, model, device, num_batches=None):
  model.eval()
  correct_predictions, num_examples = 0, 0

  if num_batches is None:
      num_batches = len(data_loader)
  else:
      num_batches = min(num_batches, len(data_loader))
  for i, (input_batch, target_batch) in enumerate(data_loader):
      if i < num_batches:
          input_batch = input_batch.to(device)
          target_batch = target_batch.to(device)

          with torch.no_grad():
              logits = model(input_batch)[:, -1, :]  # Logits of last output token
          predicted_labels = torch.argmax(logits, dim=-1)  #B

          num_examples += predicted_labels.shape[0]
          correct_predictions += (predicted_labels == target_batch).sum().item()  #C
      else:
          break
  return correct_predictions / num_examples

train_accuracy = calc_accuracy_loader(train_loader, model, device, num_batches=10)
print(f"Training accuracy: {train_accuracy*100:.2f}%")

#A The classification read-out established in §6.5: last position, 2 logits.
#B Batched argmax: one predicted label per example.
#C Element-wise compare against the true labels, sum the hits — accuracy is correct/total.

Why optimize a loss you don’t report, and report a metric you don’t optimize?

Classification accuracy — the percentage of correct predictions — is what we ultimately want, but a count of argmax hits is not differentiable: it gives gradient descent nothing to climb. So training minimizes cross-entropy on the last-token logits as a smooth proxy (push probability toward the right class), and accuracy is reported alongside. The only change to ch5’s calc_loss_batch:

</> calc_loss_batch, classification edition — one line differs from ch5

def calc_loss_batch(input_batch, target_batch, model, device):
  input_batch = input_batch.to(device)
  target_batch = target_batch.to(device)
  logits = model(input_batch)[:, -1, :]  #A
  loss = torch.nn.functional.cross_entropy(logits, target_batch)  #B
  return loss

# calc_loss_loader: EXACTLY the same as chapter 5
with torch.no_grad():
  train_loss = calc_loss_loader(train_loader, model, device, num_batches=5)
  val_loss = calc_loss_loader(val_loader, model, device, num_batches=5)
  test_loss = calc_loss_loader(test_loader, model, device, num_batches=5)

#A THE difference: only the last token's logits — ch5 scored every position with model(input_batch).
#B No flattening needed: logits are already (batch, 2) and targets (batch,) — one classification per example.

6.7 Fine-tuning on supervised data

The training loop is chapter 5’s with two bookkeeping changes — it counts examples instead of tokens, and reports accuracy after each epoch:

</> train_classifier_simple — ch5's loop, classification edition

def train_classifier_simple(model, train_loader, val_loader, optimizer, device,
                          num_epochs, eval_freq, eval_iter):
  train_losses, val_losses, train_accs, val_accs = [], [], [], []
  examples_seen, global_step = 0, -1

  for epoch in range(num_epochs):
      model.train()

      for input_batch, target_batch in train_loader:
          optimizer.zero_grad()
          loss = calc_loss_batch(input_batch, target_batch, model, device)
          loss.backward()
          optimizer.step()
          examples_seen += input_batch.shape[0]  #A
          global_step += 1

          if global_step % eval_freq == 0:  #B
              train_loss, val_loss = evaluate_model(
                  model, train_loader, val_loader, device, eval_iter)
              train_losses.append(train_loss)
              val_losses.append(val_loss)
              print(f"Ep {epoch+1} (Step {global_step:06d}): "
                    f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")

      # Calculate accuracy after each epoch
      train_accuracy = calc_accuracy_loader(train_loader, model, device,
                                            num_batches=eval_iter)  #C
      val_accuracy = calc_accuracy_loader(val_loader, model, device,
                                          num_batches=eval_iter)
      print(f"Training accuracy: {train_accuracy*100:.2f}% | ", end="")
      print(f"Validation accuracy: {val_accuracy*100:.2f}%")
      train_accs.append(train_accuracy)
      val_accs.append(val_accuracy)

  return train_losses, val_losses, train_accs, val_accs, examples_seen

# evaluate_model: same as chapter 5

optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.1)  #D
num_epochs = 5
train_losses, val_losses, train_accs, val_accs, examples_seen = \
  train_classifier_simple(
      model, train_loader, val_loader, optimizer, device,
      num_epochs=num_epochs, eval_freq=50, eval_iter=5,
  )

#A New: track examples instead of tokens — the natural progress unit for classification.
#B The familiar periodic loss checkpoints (every 50 steps here).
#C New: per-epoch accuracy on eval_iter batches of each split — the metric that actually matters, measured cheaply.
#D AdamW again, gentler learning rate (5e-5) than pretraining — we're nudging a good model, not building one. Note model.parameters() includes frozen params; requires_grad=False simply means they get no gradients. ~5 minutes on an M3 MacBook Air, <30 s on a V100/A100.

The result, after five epochs:

Split⇅	Accuracy (%)⇅
Training	97.21	+
Validation	97.32	+
Test	95.67	+

Final accuracies (the notes print these three numbers — no intermediate curves). Click rows for interpretation.

Choosing the number of epochs

There’s no universal rule; five is usually a good starting point. If the loss plot shows overfitting after the first few epochs, reduce; if the trendline suggests validation loss would keep improving, increase.

Key idea — small changes, big reuse

Everything hard was already built: the model (ch4), the weights (ch5), the loop (ch5). Classification fine-tuning changed one line of the loss, one output layer, and what gets counted — and produced a 95.67%-on-unseen-data spam detector in five minutes on a laptop. The final unit puts it to work on new messages.

📝 Check yourself: training the classifier

0 / 5

1.Why train with cross-entropy loss when what we actually care about is accuracy?
2.How does the classification calc_loss_batch differ from chapter 5's?
3.Predict: an untrained-head model sees "You won the lottery" and outputs last-token logits [3.5983, −3.9902]. What does argmax predict?
4.Which two tracking changes distinguish train_classifier_simple from ch5's train_model_simple?
5.Final results: train 97.21%, validation 97.32%, test 95.67%. What's the healthiest reading of these numbers?