Regularizers & training tricks

The idea

Most training tricks were invented for big networks that overfit: they memorise the training set and do worse on new data. A small network usually underfits, because it lacks the capacity to fit even the training data. This page tests four of them on a 2-layer MLP whose bottleneck is capacity, not data. Label smoothing trains on targets that are 90% the true class and 10% spread over all classes. Mixup trains on blends of two examples and their labels. EMA of weights keeps an exponential moving average of the weights and ships the average. Longer training uses four times the steps.

How it works

1. Label smoothing with ε = 0.1: y_smooth = (1 - ε) * onehot(y) + ε / C, so the model is never pushed to put 100% on one class.
2. Mixup: pair each example with a random partner, draw λ ~ Beta(1, 1) (uniform on 0-1), and train on x = λ x_a + (1 - λ) x_b with target λ y_a + (1 - λ) y_b.
3. EMA: after every optimizer step, shadow = d * shadow + (1 - d) * W with decay d close to 1. At the end, replace W by shadow.
4. Longer: the same cosine learning-rate schedule stretched over 4x the steps.

The code, line by line

CE in lab/train.py is cross-entropy against b.Y, a matrix of target probabilities (one-hot, or soft for mixup). With smoothing, Y = Y * (1 - self.ls) + self.ls / Y.shape[1] mixes in the uniform distribution. m = b.lab.astype(F)[:, None] masks out unlabeled rows and n counts the labeled ones. The loss is -(m * Y * log_softmax(z)).sum() / n, with logit gradient (softmax(z) - Y) * m / n.

Mixup lives in OnlineSource.batches. k = rng.permutation(len(j)) pairs examples within the batch, lam = rng.beta(1.0, 1.0, size=(len(j), 1)) draws one λ per example, and X and Y are blended with it. lab = lab & lab[k] keeps a row labeled only if both parents were.

fit in lab/train.py, the lab's training loop, uses AdamW (wd=1e-4), fixes the budget with total = max_steps or epochs * steps_per, and on each step sets opt.lr = lr * 0.5 * (1 + np.cos(np.pi * min(step / total, 1.0))), a cosine decay to 0. Stretching total 4x is all "4x longer" does.

For EMA, shadow = [p.copy() for p in net.params()] if ema else None copies the initial weights. Each step does s_ *= ema; s_ += (1 - ema) * p, and at the end p[...] = s_ replaces the trained weights. There is no warm-up or bias correction, so after T steps the random init still makes up ema ** T of the average. The small_tricks experiment uses ema=0.999 over 3,000 steps, leaving 0.0497 of the init in the shipped weights.

run_tricks in lab/methods.py trains a baseline and the chosen trick for 1,000 steps each (4,000 for "4x longer"), seed 0. Its EMA uses ema=1 - 10 / STEPS, which is 0.99 at this budget, so the init's share is about 0.99 ** 1000, effectively zero.

What the simulations show

From small_tricks, 3,000 steps, mean of 3 seeds, accuracy on 20 classes. The 3x256 teacher reaches 0.9139.

Training longer is the only one of the four tricks with a clear gain. Label smoothing is within noise (2x16 standard deviations are about 0.008). Mixup and EMA hurt. That fits underfitting: regularizers fight memorisation, and this model has none to fight. Distilling from the teacher (T=1, on the labeled inputs plus 20,000 unlabeled ones) adds only 0.0077 for 2x16 and 0.0148 for 2x32.

Page runs (1,000 steps, width 16): baseline 0.5803, label smoothing 0.5846, mixup 0.5397, EMA 0.5796, 4x longer 0.6909. At width 64, EMA gives 0.8508 against 0.8514. With a decay matched to the run, EMA is neutral.

With 8,000 clean labels the toy cannot show where regularizers usually help: big models on small or noisy data.

When it breaks

• EMA decay too slow for the run. The average spans roughly 1 / (1 - ema) steps. With 0.999 over 1,000 steps, 37% of the shipped weights would still be the random init. Warm the decay up, or bias-correct it the way Adam corrects its moments.
• Mixup on an underfitting model makes a task it can't fit even harder.
• Label smoothing and distillation. A smoothed model's probabilities carry less class-similarity information. In teacher_smoothing, a 2x16 student distilled at T=4 reached 0.6398 from the plain teacher and 0.5738 from one trained with smoothing 0.1, though that teacher was more accurate (0.9216 vs 0.9139). At T=1 the gap nearly vanishes (0.6714 vs 0.6696).
• Unequal budgets. "4x longer" costs 4x the compute.

Try this

• Switch trick between label smoothing and mixup at width 64, a model closer to fitting the task, and see whether the sign changes.
• Run trick = 4x longer at widths 4, 16 and 64. The gain should shrink as the model gets big enough to converge early.

Going further

Big models overfit, so these tricks are standard there. Ultralytics' trainer, which the YOLOv8n cockpit's recover_job.py drives, keeps an EMA of the weights by default and ramps its decay up over early updates, avoiding the init problem above. STATUS.md lists "more train data, more epochs" as the cockpit's next lever. On this toy, more steps was the one trick that reliably helped a capacity-limited model.

Source code

The exact code that runs, read from the repository when this page was built.

class CE:
    """Cross-entropy on ground-truth labels (rows without a label are ignored)."""

    def __init__(self, weight: float = 1.0, smoothing: float = 0.0):
        self.w, self.ls = weight, smoothing

    def __call__(self, z, hid, b: Batch):
        N = len(z)
        if b.Y is None or not b.lab.any():
            return 0.0, np.zeros_like(z), {}
        Y = b.Y
        if self.ls:
            Y = Y * (1 - self.ls) + self.ls / Y.shape[1]
        m = b.lab.astype(F)[:, None]
        n = max(m.sum(), 1.0)
        loss = -(m * Y * log_softmax(z)).sum() / n
        dz = (softmax(z) - Y) * m / n
        return self.w * loss, self.w * dz, {}

    params = staticmethod(lambda: [])

def fit(net: MLP, source, terms: list, epochs: int = 40, lr: float = 3e-3, wd: float = 1e-4,
        bs: int = 128, seed: int = 0, ema: float | None = None, callback=None,
        max_steps: int | None = None) -> MLP:
    """Train ``net`` in place (cosine LR, AdamW). Returns the net (EMA weights if
    ``ema`` is set). ``callback(epoch, net)`` runs after every epoch."""
    rng = np.random.default_rng(seed)
    extra = [p for t in terms for p in t.params()]
    opt = Adam(net.params() + extra, lr=lr, wd=wd)
    steps_per = int(np.ceil(len(source) / bs))
    total = max_steps or epochs * steps_per
    shadow = [p.copy() for p in net.params()] if ema else None
    step = 0
    for ep in range(epochs):
        for b in source.batches(rng, bs):
            opt.lr = lr * 0.5 * (1 + np.cos(np.pi * min(step / total, 1.0)))
            z, hid, cache = net.forward(b.X, cache=True)
            dz = np.zeros_like(z)
            dhid: dict[int, np.ndarray] = {}
            for t in terms:
                _, g, gh = t(z, hid, b)
                dz += g
                for k, v in gh.items():
                    dhid[k] = dhid.get(k, 0) + v
            grads, _ = net.backward(cache, dz, dhid or None)
            grads += [t.dP for t in terms if isinstance(t, Hint)]
            opt.step(grads)
            net.apply_masks()
            if shadow is not None:
                for s_, p in zip(shadow, net.params()):
                    s_ *= ema; s_ += (1 - ema) * p
            step += 1
            if step >= total:
                break
        if callback:
            callback(ep, net)
        if step >= total:
            break
    if shadow is not None:
        for s_, p in zip(shadow, net.params()):
            p[...] = s_
    return net

Simulation figures

Bag of tricks for a small model

Fixed small architecture. What actually moves its accuracy: regularizers, longer training, weight averaging, over-parameterize-then-collapse, self-distillation, a teacher, or train-big-then-shrink?

Full table: Bag of tricks for a small model

Papers

• Szegedy et al. Rethinking the Inception Architecture (label smoothing) (2016)
• Zhang et al. mixup (2018)
• Izmailov et al. Averaging Weights Leads to Wider Optima (2018)