Lottery tickets

The idea

When a pruned network recovers well, is that because its trained weights were good, or because its mask (the set of surviving connections) picked a sub-network that was trainable from the start? The lottery ticket experiment separates the two. Find a mask by pruning a trained network, reset the surviving weights to their values at initialisation (or very early in training), and retrain with the mask fixed. If this "ticket" trains well while the same mask with fresh random weights does not, the initial values mattered. Frankle and Carbin called it a winning ticket.

How it works

1. Initialise a network, save its weights W0, train it, and also save an early snapshot Wk (the rewind point).
2. Prune the trained net by magnitude to get a 0/1 mask m.
3. Reset: W = W0 * m (rewind to init) or W = Wk * m (late rewinding).
4. Retrain with the mask fixed.
5. Repeat steps 2-4, pruning more each round (iterative magnitude pruning).
6. Controls at the same sparsity: the ticket mask with a new random init, and a random mask on the original init.

The code, line by line

lottery_ticket in lab/experiments/shrink.py uses a 2x128 net with 20,480 weights. init = net(hid, seed=s) is W0. scratch(init.copy(), s, callback=...) trains a copy for 3,000 steps, and snap.setdefault("early", n.copy()) if ep == 1 keeps the weights after the second epoch: step 126, about 4% of training.

branches maps each arm to its starting weights: init, early, or None (keep the trained weights and fine-tune). Each round sets sp = 1 - 0.5 ** rnd and calls mask_src = magnitude_prune(state[name], sp) on that arm's net from the previous round. magnitude_prune counts sp over all weights, so each round halves the survivors and each mask is a subset of the last.

m = mask_src.copy() if start is None else start.copy() picks the weights, m.masks = [...] attaches the mask, and m.apply_masks() zeroes the pruned entries. scratch(m, s, lr=FT_LR if start is None else 3e-3) retrains for 3,000 steps. Both learning rates are 3e-3, so the arms differ only in their starting weights. state[name] = m carries each arm's lineage into the next round.

The controls follow. "ticket mask, fresh init" puts the rewind-to-init mask on net(hid, seed=999 + s). "random mask, original init" draws rr.random(mk.shape) < mk.mean() per layer, keeping each layer's density but not its pattern.

run_lottery in lab/methods.py is a short version for this page: 2x64, 1,000 steps per round, one rewind point.

What the simulations show

From lottery_ticket (2x128, mean of 2 seeds; dense accuracy 0.904):

This is the lottery-ticket pattern. Rewound tickets hold up best at every sparsity, and the gap grows as the net gets sparser. At 96.9% (640 weights), with the original init, the ticket mask scores 0.6118 and a random mask 0.3219. With the ticket mask, the original init scores 0.6118 and a fresh one 0.4801. So both the mask and its initial values matter. Rewinding to init and to the early snapshot are indistinguishable here. Fine-tuning the trained weights, what most pruning pipelines do, lands with the controls once sparsity passes 90%.

Page runs (2x64, one seed, prune rounds 6; dense 0.8514):

In the smaller, shorter-trained net, rewinding beats a random reinit by 0.055 to 0.063 at 50% and by 0.017 to 0.035 at 75%. From 87.5% on the three arms sit within 0.035 of each other and swap order from round to round: one seed cannot separate them.

The toy cannot test large vision models, where Frankle et al. found that rewinding to init stops working and an early rewind point is needed.

When it breaks

• Compounding schedules. If a pruning function ranks only survivors, a schedule of total sparsities overshoots. An earlier version of this experiment did that: its "75%" round was really 87.5%. Check the nonzero count.
• Tickets are expensive to find. You need a full training run per round. It is a probe, not a shortcut.
• Small or briefly trained nets show little ticket effect where it should matter most, at high sparsity: in the page runs, rewound and randomly reinitialized tickets sit within 0.035 of each other from 87.5% on.
• Masks are unstructured, so there is no speedup without sparse kernels.

Try this

• Set prune rounds to 1 and switch rewind to between early and random reinit.
• Set prune rounds to 3 (87.5%) and compare all three rewind points. The params column reads 916, of which 148 are biases.
• Run prune rounds 6 with each rewind and watch accuracy fall towards a 244-parameter net.

Going further

Frankle et al. (2020) showed that rewinding to an early iteration makes tickets work at scale, and linked it to linear mode connectivity: past that point, retraining with a different data order ends in the same basin. The YOLOv8n cockpit never rewinds. Its loop fine-tunes the trained weights (the "fine-tune trained" arm), the practical choice when the original training run cannot be repeated, as with a pretrained YOLOv8n. A ticket experiment on it would need the COCO training recipe and many full trainings.

Source code

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

@experiment("lottery_ticket", G, "Lottery tickets: is the mask or the init the secret?",
            "Iterative magnitude pruning, then retrain the sparse net from (a) its original "
            "init, (b) an early-training rewind point, (c) a fresh random init, (d) its "
            "trained weights. Plus a random mask control.")
def lottery_ticket():
    d = data()
    hid = [128, 128]
    rows = []
    for s in SEEDS[:2]:
        init = net(hid, seed=s)
        # train once, snapshotting the weights after epoch 2 (~4% of training):
        # the "late rewinding" point of Frankle et al.
        snap: dict = {}
        trained = scratch(init.copy(), s, callback=lambda ep, n: (
            snap.setdefault("early", n.copy()) if ep == 1 else None))
        early = snap["early"]
        branches = {"rewind to init": init, "rewind to early": early, "fine-tune trained": None}
        state = {k: trained for k in branches}
        for rnd in range(1, 6):
            sp = 1 - 0.5 ** rnd
            for name, start in branches.items():
                mask_src = magnitude_prune(state[name], sp)
                m = (mask_src.copy() if start is None else start.copy())
                m.masks = [None if mk is None else mk.copy() for mk in mask_src.masks]
                m.apply_masks()
                m = scratch(m, s, lr=FT_LR if start is None else 3e-3)
                state[name] = m
                rows.append(dict(sparsity=sp, method=name, seed=s, **ev(m)))
            # fresh init with the winning-ticket mask
            m = net(hid, seed=999 + s)
            m.masks = [None if mk is None else mk.copy() for mk in state["rewind to init"].masks]
            m.apply_masks()
            rows.append(dict(sparsity=sp, method="ticket mask, fresh init", seed=s, **ev(scratch(m, s))))
            # random mask, same per-layer sparsity
            m = net(hid, seed=s)
            rr = np.random.default_rng(rnd + 10 * s)
            m.masks = [(rr.random(mk.shape) < mk.mean()).astype(np.float32)
                       for mk in state["rewind to init"].masks]
            m.apply_masks()
            rows.append(dict(sparsity=sp, method="random mask, original init", seed=s, **ev(scratch(m, s))))
        rows.append(dict(sparsity=0.0, method="dense", seed=s, **ev(trained)))
    return dict(rows=agg(rows, ["sparsity", "method"]), refs=refs(), plots=[
        dict(x="sparsity", y="acc", group="method"),
    ])

Simulation figures

Lottery tickets: is the mask or the init the secret?

Iterative magnitude pruning, then retrain the sparse net from (a) its original init, (b) an early-training rewind point, (c) a fresh random init, (d) its trained weights. Plus a random mask control.

Full table: Lottery tickets: is the mask or the init the secret?

Papers

• Frankle & Carbin. The Lottery Ticket Hypothesis (2019)
• Frankle et al. Linear Mode Connectivity and the Lottery Ticket Hypothesis (2020)