Training a Classifier That Actually Generalizes

The cross-entropy loss, properly

In Guide 1 we met cross-entropy in a single line: it measures how surprised the model is by the true answer. Now we earn the right to lean on it for the rest of this guide by understanding why it behaves so well. Recall the setup of image classification: the network outputs a vector of raw, unbounded scores called logits; softmax squashes them into probabilities that are all positive and sum to 1; and the label is a one-hot label encoding — a vector that is 1 for the true class and 0 for every other class.

Cross-entropy between the one-hot target y and the predicted distribution p.

Read it symbol by symbol. The sum runs over all classes i. The term y_i is the one-hot target (1 for the true class, 0 otherwise), and p_i is the softmax probability the model assigned to class i. Because y_i is zero for every class except the correct one, the entire sum collapses to a single term: L = −log(p_true). So cross-entropy really only cares about one number — the probability you gave the correct answer. If p_true = 1 the loss is −log 1 = 0 (perfect). If p_true = 0.5 the loss is −log 0.5 ≈ 0.69. If p_true = 0.01 the loss is −log 0.01 ≈ 4.6 (badly wrong). The loss shoots toward infinity as the true probability approaches zero — that is exactly the 'never be confidently wrong' pressure we want.

The gradient of cross-entropy with respect to the logits is breathtakingly simple.

This little equation is the engine of the whole training process, so let us unpack every symbol. z_i is the logit (the raw score) for class i. p_i is the softmax probability for class i. y_i is the one-hot target (1 for the true class, else 0). The gradient is simply predicted-minus-true: p_i − y_i. For the correct class, p_true − 1 is negative, so gradient descent pushes that logit up. For every wrong class, p_i − 0 is positive, so descent pushes those logits down. And the size of each nudge is proportional to how wrong you were. Worked example with 3 classes and the true class = cat (index 0): suppose the logits are z = [2.0, 1.0, 0.1]. Softmax gives p ≈ [0.659, 0.242, 0.099]. The gradient is p − y = [0.659 − 1, 0.242 − 0, 0.099 − 0] = [−0.341, +0.242, +0.099]. The cat logit gets a negative gradient (so it rises), and both wrong logits get positive gradients (so they fall) — exactly 'push the right class up, the wrong classes down'.

Why hard labels hurt: label smoothing

Here is a subtle pathology hiding in plain sight. With a pure one-hot label encoding, cross-entropy L = −log(p_true) is only truly minimised when p_true reaches exactly 1 — and for softmax, p_true = 1 requires the true logit to run off to +infinity relative to the others. The optimiser can never quite get there, so it just keeps pushing the gap wider and wider. The model becomes pathologically over-confident: it learns to output 0.9999 even on ambiguous images. That over-confidence hurts generalisation (the model memorises sharp boundaries instead of robust ones) and wrecks calibration (its stated confidence stops matching its real accuracy).

label smoothing is the cure, and its slogan is 'be confident, but leave a little doubt.' Instead of telling the model the answer is 100% cat and 0% everything else, we tell it the answer is, say, 92% cat with a sliver of probability spread across the other classes. The model no longer has any incentive to chase infinite logits, because the target it is aiming for is finite and reachable.

The smoothed target replaces the hard one-hot value.

Symbol by symbol: y_c is the original one-hot entry for class c (1 or 0), ε (epsilon) is the small smoothing strength (a common choice is 0.1), and K is the number of classes. Plug in the two cases. For the true class, y_c = 1, so y_c^LS = (1−ε)·1 + ε/K = 1 − ε + ε/K. For every other class, y_c = 0, so y_c^LS = 0 + ε/K = ε/K. Worked example with K = 5 classes and ε = 0.1: the true class gets 1 − 0.1 + 0.1/5 = 0.9 + 0.02 = 0.92, and each of the four other classes gets 0.1/5 = 0.02. The full softened target for a 'class-0 is true' example is [0.92, 0.02, 0.02, 0.02, 0.02]. Notice it still sums to 0.92 + 4×0.02 = 1.00 — it is a valid probability distribution, just no longer all-or-nothing.

Augmentation: free data from old data

Loss tweaks shape how the model learns; data augmentation changes what it learns from. The idea is to apply label-preserving transformations to your training images — random crop, horizontal flip, colour jitter, small rotations — so the model sees many different views of the same picture. Each transformation is chosen so it does not change the correct answer in image classification: a shifted, flipped, or slightly recoloured cat is still labelled 'cat'. By being shown all these variants, the model is forced to learn the right invariances — to recognise a cat regardless of where it sits, which way it faces, or how warm the lighting is.

Why does this help so much? Recall the overfitting picture from Guide 2: a model with lots of capacity can memorise the exact training pixels, scoring perfectly on the training set while flopping on new images. Augmentation fights this by enlarging the effective dataset — if every epoch each image is cropped and recoloured a fresh random way, the model almost never sees the exact same pixels twice, so there is nothing fixed to memorise. It must instead learn the underlying, generalisable features. Augmentation is, pound for pound, one of the cheapest and most powerful regularisers you have.

# A typical training-time augmentation pipeline (all label-preserving).
# Each epoch, every image is transformed a fresh random way,
# so the model almost never sees the exact same pixels twice.
train_transform = Compose([
    RandomResizedCrop(224),          # random zoom + crop -> position/scale invariance
    RandomHorizontalFlip(p=0.5),     # OK for cats; NOT ok for text or '6'/'9'
    ColorJitter(brightness=0.4,      # small lighting / colour changes
                contrast=0.4,
                saturation=0.4),
    ToTensor(),
    Normalize(mean=IMAGENET_MEAN,    # match the stats the backbone expects
              std=IMAGENET_STD),
])

# IMPORTANT: the validation/test set gets NO random augmentation,
# only a deterministic resize + center crop, so the scorecard stays honest.
val_transform = Compose([
    Resize(256),
    CenterCrop(224),
    ToTensor(),
    Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])

Mixup: blending images and labels

Per-image augmentation transforms one picture at a time. mixup takes a leap: it combines two training examples at once, blending both their pixels and their labels by the same proportion. Take a cat image and a dog image, overlay them at 70% / 30% opacity to get a ghostly hybrid, and — crucially — set the target to 70% cat and 30% dog instead of a single hard label. You are teaching the model something one-hot data never could: how to behave in the space between classes.

Blend the pixels and the labels with the same weight λ.

Symbol by symbol: x_a and x_b are two input images and y_a, y_b are their one-hot labels (the same one-hot vectors from Guide 1). λ (lambda) is the mixing weight, a number in [0, 1]. So x̃ is the pixel-wise blend (multiply every pixel of image a by λ, every pixel of image b by 1−λ, and add), and ỹ is the matching soft label, blended with the identical weight. The model trains on this single synthetic pair (x̃, ỹ). Worked example with λ = 0.7, where image a is a cat (one-hot [1, 0]) and image b is a dog (one-hot [0, 1]): the blended image is x̃ = 0.7·(cat pixels) + 0.3·(dog pixels), a translucent cat-dog overlay, and the blended target is ỹ = 0.7·[1, 0] + 0.3·[0, 1] = [0.7, 0.3]. The label split exactly mirrors the pixel split — that consistency is the whole point.

import numpy as np
import torch
import torch.nn.functional as F

def mixup_batch(images, labels, num_classes, alpha=0.2):
    # Turn integer class ids into one-hot vectors (see Guide 1).
    y = F.one_hot(labels, num_classes).float()

    lam = np.random.beta(alpha, alpha)        # mixing weight in [0, 1], fresh each step
    perm = torch.randperm(images.size(0))     # pick a random partner for each image

    mixed_x = lam * images + (1 - lam) * images[perm]
    mixed_y = lam * y      + (1 - lam) * y[perm]
    return mixed_x, mixed_y                    # train on blended pixels AND blended labels

CutMix: cut, paste, and share the label

cutmix keeps the label-blending insight of mixup but throws away the translucency. Instead of overlaying whole images like a double exposure, CutMix cuts a rectangular patch out of one image and pastes it solidly onto another, then mixes the two labels in proportion to the patch area. The result looks like a real, opaque photo — a dog with a square of cat stamped over its corner — with no ghosting. The model now has to make sense of two genuine, localised pieces of evidence sitting side by side.

The label split exactly matches the pixel-area split.

Symbol by symbol: y_a is the host image's one-hot label (the picture we paste onto), y_b is the source image's label (the picture the patch came from), and λ (lambda) is the fraction of the image that still belongs to the host. We define λ = 1 − (area of pasted patch) / (area of image), so 1 − λ is exactly the patch's area fraction. Worked example: paste a patch that covers 30% of the image. Then 1 − λ = 0.30, so λ = 0.70. The blended label is ỹ = 0.70·y_a + 0.30·y_b. If the host is a dog and the patch is a cat, the target becomes 70% dog + 30% cat — exactly matching the fact that 70% of the visible pixels are dog and 30% are cat. The accounting is literal: pixel area in, label proportion out.

Why prefer this to mixup? Because the cat evidence is confined to one corner, the model cannot lean on a single global cue (overall colour, average texture) to win — it must use localised evidence from different image regions, which improves its sense of where objects are and makes it more robust to occlusion (objects partly hidden behind other things). Rough guidance: CutMix tends to shine on large natural-image datasets and object-heavy tasks because its images stay realistic; mixup is often gentler and a good default when objects are small or fill the whole frame, or when translucent blends are acceptable. Many strong recipes simply randomly alternate between the two per batch and get the best of both.

Putting it together: a modern training recipe

None of these tricks is a silver bullet, but together they form the backbone of every champion classifier. The unifying theme is fighting overconfidence and overfitting from several angles at once: strong label-preserving augmentation enlarges the effective dataset; mixup and cutmix smooth the decision boundaries between classes; and label smoothing stops the model from chasing infinite logits. They attack the same disease — a model too eager to memorise and too sure of itself — through different mechanisms, which is why stacking them compounds the benefit rather than being redundant.

1. Start from a clean baseline: standard resize/crop/flip augmentation plus plain cross-entropy. Confirm it trains and even overfits a little — that proves the model has enough capacity to learn the task.
2. Strengthen label-preserving augmentation (RandomResizedCrop, flip if safe, mild colour jitter) to push back on overfitting.
3. Turn on label smoothing with ε ≈ 0.1 so the model stops chasing infinite logits and becomes better calibrated.
4. Add mixup OR cutmix, sampling λ from a Beta distribution each batch; many strong recipes randomly alternate between the two.
5. Train longer: these regularisers slow down memorisation, so the model needs more epochs to converge to its best validation score.
6. Select and report on validation top-1/top-5 accuracy, never on the (now inflated) training loss.

Two important cautions. First, do not over-regularise small datasets: aggressive mixup/cutmix plus heavy augmentation can starve a tiny dataset of clean signal, and the model underfits — accuracy drops. Dial the strength down (or off) when data is scarce, and watch the validation curve to decide. Second, beware double-counting the softness: mixup and cutmix already produce soft labels, so layering strong label smoothing on top can over-soften the targets and waste capacity. If you combine them, use a smaller ε, or let the mix-augmentation provide the softness on its own.

There is a bonus hiding in all of this. Every trick here nudges the model away from brittle 0.999 certainties toward softer, better-behaved probabilities — and well-behaved probabilities are the raw material of calibration, the art of making a model's stated confidence actually mean something. That is exactly where Guide 5 begins. The next guide, though, tackles the cases where even a well-trained classifier struggles: fine-grained, imbalanced, and structured labels. The recipe you just built is the launch pad for both.