Rigging the Lottery: Making All Tickets Winners

In plain words

Making a large neural network small usually means training it in full and then cutting it down. Sparse training tries to keep the network small from the start: most connections are switched off, and only the rest are trained. The hard part is choosing which connections stay on.

RigL changes that choice during training. At regular intervals it switches off the weakest active connections and switches on the same number of new ones, picked where the error signal (the gradient) suggests they would help most. New connections start at zero, so the network's output does not jump when they appear. The authors report that this gives better accuracy than earlier sparse-training methods at the same training cost, and that in several settings it is competitive with pruning a dense network.

TL;DR

RigL keeps a fixed number of weights active and rewires them during training: every ΔT steps it drops the smallest-magnitude active weights and grows as many new ones where the dense gradient is largest, starting them at zero §3. In the authors' ImageNet ResNet-50 experiments it beats the other sparse-training methods in the same uniform-sparsity setting, and with 5x longer training it is above the magnitude-pruning baselines it cites fig. 2. On CIFAR-10 it matches pruning at fewer training FLOPs §4.3; on character-level WikiText-103 it does not match pruning §4.2. Read the notation and warm-up notes under Pitfalls and Paper vs code before implementing.

Problem

Sparse networks save memory and inference FLOPs, but the most accurate sparse models have usually been obtained from dense networks, so training still costs at least as much as a dense model §1. Static sparse training, where the mask is fixed from the start, performs worse than dense-to-sparse methods; the paper cites several works for this, and its own ImageNet runs show it §4.1.1. Earlier dynamic methods either grow connections at random §2 or need dense gradients and momentum for every parameter at every step, which is costly in memory and compute §2. The paper asks whether a sparse network with a fixed size and fixed cost can keep its quality while its connections change during training §1, §3.

Assumptions: fixed sparsity per run (no change of layer budgets during training, so FLOPs can be fixed in advance §3 (1)); FLOP counts assume the training uses sparsity fig. 2 caption.

Method

Figure and table numbers below are counted from the order of the LaTeX source, not checked against the rendered PDF. Equations in §3 are unnumbered in the source, so they are cited by section.

RigL starts from a random sparse network with a per-layer sparsity. Every ΔT iterations, until T_end, it updates each layer's mask in two parts. First it drops a fraction of the active connections with the smallest weight magnitude. Then it grows the same number of connections at positions that are currently inactive, choosing those with the largest gradient magnitude. Between updates it trains the active weights as an ordinary sparse network alg. 1. The number of connections replaced in layer l is k = f_decay(t) · (1 − s^l) · N^l, where the fraction f_decay(t) starts at α and decays with a cosine schedule §3 (2)–(3).

The paper's building blocks, in the order it describes them §3:

1. Sparsity distribution §3 (1). The paper compares three ways of setting per-layer sparsity for a given total sparsity S: Uniform (every layer gets S, and the first layer is kept dense), Erdős–Rényi (ER, introduced by SET §3 (1)), and Erdős–Rényi-Kernel (ERK), which adds the kernel width and height to the scaling. For ER and ERK, the sparsity scales with 1 minus a ratio of layer dimensions; the paper's wording is ambiguous, see Pitfalls. Bias and batch-norm parameters are kept dense in all methods §3 (1). The distribution is fixed for the whole run, so the FLOP budget is known in advance §3 (1).

2. Update schedule §3 (2). Four parameters: ΔT (iterations between updates), T_end (last iteration with an update), α (initial fraction of connections updated) and f_decay (a function that may shrink the fraction over time). The paper uses cosine annealing, which it found slightly better than the alternatives §3 (2):

   f_{decay}(t;~\alpha,~T_{end})=\frac{\alpha}{2}\left(1+\cos\left(\frac{t\pi}{T_{end}}\right)\right) §3 (2)

   Constant and inverse-power schedules are reported in app. G.

3. Drop criterion §3 (3). Every ΔT steps, drop the connections given by the top-k of −|θ^l|, where k is the number of connections to replace. In words: drop the k active weights with the smallest magnitude.

4. Grow criterion §3 (4). Grow the k inactive positions with the largest magnitude of the current minibatch gradient. Positions that are still active after the drop are excluded, while dropped positions may be grown again. Newly activated connections are initialised to zero, so they do not change the output at the moment they appear. The paper says random values and small values along the gradient direction gave no better results §3 (4).

5. Cost §3, final paragraph; app. H. The dense gradient is needed only at update steps and can be computed layer by layer and discarded after the top-k selection. The extra cost is amortised when ΔT > 1/(1 − s), so the average work stays proportional to 1 − S §3. Per-sample training FLOPs are (3·f_S·ΔT + 2·f_S + f_D) / (ΔT + 1), where f_S and f_D are the sparse and dense forward FLOPs; the cost of the top-k selection is omitted app. H.

# RigL: drop and grow connections during training  [alg. 1]
# Input:  network, dataset D, per-layer sparsities s[l]
# Input:  dT (update interval), T_end (last update), alpha (initial drop)
# Output: weights and masks of the sparse network

# Cosine decay of the fraction of connections replaced  [§3 (2)]
def f_decay(step, alpha, T_end):
    # alpha / 2 * (1 + cos(step * pi / T_end))  [§3 (2)]
    return alpha / 2 * (1 + cos(step * pi / T_end))

def train_rigl(network, dataset, s, dT, T_end, alpha, lr):
    # Random sparse start with the chosen per-layer sparsities  [§3 (1)]
    masks = random_masks(network, sparsities=s)
    weights = network.initial_weights() * masks
    for step in range(num_steps):
        # Sample a minibatch B_t from D  [alg. 1]
        batch = sample_minibatch(dataset)
        # Loss summed over the samples of the batch  [alg. 1]
        loss = sum_loss(network(weights * masks, batch))
        if step % dT == 0 and step < T_end:
            # Fraction of the active weights to replace now  [§3 (2)]
            fraction = f_decay(step, alpha, T_end)
            for layer in network.layers:
                # Gradient for every position, zeros included  [§3 (4)]
                dense_grad = grad_of_effective_weights(loss, layer)
                # Number of connections to drop and to grow  [§3 (2)]
                k = fraction * (1 - s[layer]) * layer.size
                # Drop the k smallest-magnitude active weights  [§3 (3)]
                dropped = smallest_k(abs(weights[layer]), k, masks[layer])
                masks[layer][dropped] = 0
                # Grow the k largest-gradient inactive positions  [§3 (4)]
                grown = largest_k(abs(dense_grad), k, masks[layer] == 0)
                # New connections start at zero  [§3 (4)] [alg. 1]
                # Reference code differs for regrown positions, see Pitfalls
                weights[layer][grown] = 0
                masks[layer][grown] = 1
        else:
            # Ordinary sparse SGD step  [alg. 1]
            # Alg. 1 writes the step size as alpha; here it is lr
            weights = weights - lr * grad(loss, weights * masks)

Hyperparameters

Values as stated in the paper. "not stated" means the paper gives no value.

Results

Author-reported. Setup: ImageNet-2012 with ResNet-50 (top-1 accuracy), MobileNet-v1 and -v2 on ImageNet, character-level WikiText-103 (validation bits per step; values only in fig. 4 left), WideResNet-22-2 on CIFAR-10 (test accuracy; dense 94.1% §4.3), and an MLP on MNIST app. B. Baselines marked * are numbers reported in the original papers, not rerun by the authors fig. 2 caption. FLOPs are normalised to the dense model fig. 2.

ImageNet-2012, ResNet-50 fig. 2 left

Pruning* results come from Gale et al. (arXiv:1902.09574). The caption says the DNW* numbers come from Kusupati et al. (2020) fig. 2. Text: at 96.5% sparsity, ResNet-50 with ERK reaches 72.75% top-1, about 3.5 points above the extended magnitude-pruning result §4.1.1.

MobileNet on ImageNet fig. 3 right

Text: RigL trains 75% sparse MobileNets with no loss in performance §4.1.2. A MobileNet-v1 widened by a factor of 1.98 and kept 75% sparse matches the dense baseline's FLOPs and parameter count, and gains 4.3 points top-1 under RigL §4.1.2.

High sparsity, ResNet-50 app. L, table 4

The Static (ERK) row repeats the 80% and 90% values of fig. 2 under the 95% and 96.5% headers; see Pitfalls.

Character-level WikiText-103 §4.2; fig. 4 left

Final validation loss in bits per step is shown only in the figure; the text gives no numbers. RigL exceeds the other sparse-training methods and its results keep improving with more steps; it does not match magnitude pruning in this setting §4.2.

CIFAR-10, WideResNet-22-2 §4.3; fig. 4 right

Dense baseline 94.1% test accuracy. Some 50%-sparse networks generalise better than the dense baseline. RigL matches pruning with a fraction of the training resources, while static sparse networks fall behind pruning as sparsity grows §4.3. Per-sparsity values are shown only in the figure.

MNIST MLP, structured-pruning comparison app. B, table 2

The training costs of SBP, L0 and VIB are lower bounds computed by the authors from the reported accuracies, not measured app. B.

Restarting from the RigL network, with lottery-ticket initialisation app. E, table 3

The table does not state the network or sparsity used.

Paper vs code

Searched: the paper names github.com/google-research/rigl in the abstract, footnote and §4. The authors are at Google Brain and DeepMind, and the repository sits in the google-research organisation. Code cloned at commit d39fc7d46505cb3196cb1edeb32ed0b6dd44c0f9. The TensorFlow 1 path (rigl/sparse_optimizers*.py, rigl/sparse_utils.py, rigl/imagenet_resnet/, rigl/cifar_resnet/) was read for the settings above. rigl/rigl_tf2/ was not compared in detail. The character-level GRU experiment has no code in the repository.

• Warm-up. The paper says the ImageNet learning rate warms up linearly to 1.6 at epoch 5 §4.1. In the reference script the ResNet schedule starts at epoch 0, so the warm-up term divides by zero and is overridden by the constant entry for every epoch imagenet_train_eval.py imagenet_train_eval.py. The MobileNet schedule starts at epoch 8, so its warm-up does apply imagenet_train_eval.py. For ResNet the learning rate is therefore 1.6 from the first step. The peak value agrees: 0.1 × 4096 / 256 = 1.6 imagenet_train_eval.py.
• Momentum. The paper says SGD with momentum §4.1. The code uses Nesterov momentum imagenet_train_eval.py.
• Drop schedule. The paper uses cosine annealing §3 (2). The ImageNet and CIFAR-10 scripts default to a constant drop fraction imagenet_train_eval.py resnet_train_eval.py. Cosine is the default only in the MNIST script mnist_train_eval.py. The repository does not show which setting produced the ImageNet and CIFAR-10 results. Launch flags are not given for them.
• First layer. The paper keeps the first layer dense only for Uniform §3 (1). The ImageNet default is first_layer_sparsity=0, so the first layer stays dense imagenet_train_eval.py. The README's ERK evaluation command passes --first_layer_sparsity=-1, which prunes the first layer at the default sparsity README.md. The paper gives no first-layer rule for ERK.
• Architecture defaults. The CIFAR-10 script defaults to depth 16 and width 4 resnet_train_eval.py. The paper's WideResNet-22-2 needs resnet_depth=22 and resnet_width=2 passed explicitly; depth 22 satisfies the script's 6n+4 check resnet_train_eval.py.
• Dropped-and-regrown connections. The paper says newly activated connections are initialised to zero §3 (4). The code initialises only positions that were inactive before the update. A position dropped and regrown in the same update keeps its old value, because RigL calls the update with reinit_when_same=False sparse_optimizers_base.py sparse_optimizers_base.py.
• Momentum of new connections. The paper says nothing about optimizer state. In the code, the momentum slot of a newly grown connection is reset to the value rigl_acc_scale × masked gradient, which is zero by default sparse_optimizers_base.py imagenet_train_eval.py.
• ERK sparsities. The paper says ER and ERK scale sparsity with 1 minus a ratio of dimensions §3 (1). The code solves for a global constant that scales the ratio so that total parameters match the budget; a layer whose density would exceed 1 is made dense and the solve repeats sparse_utils.py. The sparsity of each layer is 1 − constant × ratio sparse_utils.py. The paper does not describe the dense-layer rule.
• Grow and drop scores. The drop score is |θ| plus Gaussian noise with standard deviation 1e-5, used only for tie-breaking sparse_optimizers_base.py. Growth uses the absolute gradient of the masked weights from the current minibatch sparse_optimizers_base.py sparse_optimizers_base.py.
• Checkpoint numbers. The README's ResNet-50 checkpoint table matches the paper to rounding at 1x (for example uniform 80%: 74.60 in the README, 74.6 in fig. 2) README.md. At 5x the differences are larger: uniform 80% is 76.71 in the README and 76.6 in fig. 2; ERK 80% is 77.17 against 77.1 README.md. The README also lists 99% sparse ERK checkpoints from 1x to 100x training, which the paper does not report README.md.
• Mask update stop. The code's first update happens at step 0 because the last-update counter starts at minus the interval sparse_optimizers_base.py. This matches Algorithm 1 at t = 0 alg. 1.

Pitfalls

• Overloaded notation. α is the initial drop fraction §3 (2), but Algorithm 1 uses α as the learning rate in its update step alg. 1. Use your optimizer's learning rate there.
• The set I_active. Algorithm 1 defines I_active = ArgTopK(−|θ|, k), which by the drop rule §3 (3) is the set of connections to drop. The "Changes since publishing" list says I_active means the connections kept app. N. The code drops the smallest magnitudes and keeps the largest sparse_optimizers_base.py. Implement the drop rule, not the changelog's reading.
• ER/ERK wording. The paper says sparsity "scales with 1 minus a ratio" §3 (1). The reference code makes density proportional to the ratio, so sparsity is 1 minus a scaled ratio. Use the code's behaviour: solve a single constant so that the total number of parameters meets the budget, and make oversized layers dense. See Paper vs code.
• Growth budget. Each update grows as many connections as it drops, so the parameter count and per-layer sparsity stay fixed §3 (2)–(4). Keep the count equal per layer; a global count would move parameters between layers.
• Gradient cost. The dense gradient is needed at update steps only. Its cost is amortised only if ΔT > 1/(1 − s) §3. At 99% sparsity this is ΔT > 100; at lower sparsity the condition holds easily. Check this when you pick ΔT.
• Batch-size sensitivity. Growth uses the current minibatch gradient. The paper uses batch size 4096 on ImageNet and 128 on CIFAR-10 and does not report the effect of a different batch size §4.1 §4.3.
• MobileNet sparsity labels. Sparsities in the MobileNet experiments are computed over pruned layers only. Real sparsities, counting the first layer and the depthwise convolutions, are 74.2, 84.1, 89 and 94 for the labels 75, 85, 90 and 95 §4.1.2. Report the real value when you compare.
• FLOP accounting. The FLOP counts assume the top-k selection and the random growth in SET are free app. H. Pruning is charged for its average density over training. The counts are estimates and do not measure wall-clock time.
• Replica consistency. Random drop and grow decisions must agree across data-parallel replicas, and the masked gradients must be all-reduced across replicas. The authors' own early bugs were fixed in these ways; the fixes raised SET by 0.1–0.3 points and RigL and SNFS by about 0.5–1 point at default lengths app. M.
• Table 4 duplicate row. The Static (ERK) row in the 95% and 96.5% table matches the 80% and 90% values of fig. 2 exactly, including its FLOPs (0.42x and 0.24x). Those FLOPs cannot be right for two different sparsities. Check against a run before using that row app. L.
• Test-set selection in the MNIST table. The authors note that the baselines' code reports the lowest test error seen during training as the final error. They report that lowest value in the table and the end-of-training error in parentheses app. B. Compare the parenthesised values if you want end-of-training numbers.
• Char-level LM is not reproducible from this repository. The GRU and WikiText-103 code is not in the repository, so the settings in app. I can only be checked against the paper.
• Figure and table numbers are counted from the LaTeX source order (see Method).

Lineage

• Builds on: Dynamic sparse training with random growth (SET, Mocanu et al. 2018) §2; §3 (1), which RigL replaces with gradient-based growth; SET's Erdős–Rényi sparsity distribution is used §3 (1). Also magnitude pruning and gradual pruning, as in "To Prune, or Not to Prune" [arXiv:1710.01878] §2; §4.1.
• Compared against: Static sparse training, SNIP §4.1.1, SET §4.1.1, DSR §2, SNFS (arXiv:1907.04840) §2, magnitude pruning (Gale et al., arXiv:1902.09574) §4.1, DNW fig. 2 caption, and structured pruning methods SBP, L0 and VIB on MNIST app. B.
• Related: The lottery ticket hypothesis (Frankle et al.); this paper finds that rewiring beats restarting from the original initialisation app. E (see also arXiv:1903.01611 on lottery tickets at scale, cited in §2).
• Superseded by: not stated.