Winning the Lottery by Preserving Network Training Dynamics with Concrete Ticket Search

AI summary

Overview

Research area: Neural network pruning / model compression, specifically the Lottery Ticket Hypothesis (LTH) and Pruning-at-Initialization (PaI).

Technical level: Advanced. The paper relies on combinatorial optimization, the Gumbel-softmax / Concrete relaxation, constrained optimization, and knowledge distillation.

Scope: The paper proposes Concrete Ticket Search (CTS), a single-shot differentiable search over binary pruning masks that draws "winning tickets" near initialization, and compares it against saliency-based PaI methods and Lottery Ticket Rewinding (LTR) on image classification tasks.

What This Paper Is About

The Lottery Ticket Hypothesis says dense, randomly initialized networks contain sparse trainable subnetworks ("winning tickets"), but the best method for finding them, Lottery Ticket Rewinding (LTR), needs repeated retraining and is computationally prohibitive. Cheaper Pruning-at-Initialization (PaI) methods score each weight independently with first-order saliency metrics, which the authors argue ignore inter-weight dependencies, produce a large accuracy-sparsity gap, and fail basic sanity checks. The goal of this work is an efficient algorithm that finds high-performing subnetworks near initialization, matching or beating LTR while requiring far less computation and no sensitive hyperparameter tuning.

Key Contributions

1. A diagnosis of first-order saliency pruning. The authors argue that synaptic saliency metrics of the form S(theta) = (dR/dtheta) ⊙ theta only measure the effect of removing weights independently, and that ignoring inter-weight coupling explains PaI's failure in the sparse regime. They show with VGG-16 on CIFAR-10 that saliency-based methods (SNIP, GraSP) often perform worse than LTR and can be surpassed by naive magnitude pruning even on the objectives those methods are designed to optimize.

2. Concrete Ticket Search (CTS). A method that reframes subnetwork discovery as a holistic combinatorial optimization problem, using a Concrete (Gumbel-softmax) relaxation of the binary mask space so gradients flow to per-weight retention probabilities.

3. GradBalance, a sparsity-control scheme. Instead of a tunable L1/L2 regularizer or a naive Lagrange multiplier, the authors directly balance the gradient norms of the objective and a normalized sparsity constraint with an exponential moving average, avoiding task-specific hyperparameter tuning. A traditional Lagrange multiplier variant is also presented.

4. Knowledge-distillation-inspired pruning objectives. Beyond task loss, relative change in loss, and negative gradient norm, the authors propose three new objectives: reverse KL divergence between the sparse and dense output distributions, feature-map matching with per-layer normalization, and gradient matching between sparse and dense layers. Reverse KL (CTS_KL) is reported as especially effective.

Main Findings

• CTS_KL beats LTR at high sparsity on ResNet-20 / CIFAR-10 in the abstract's example: at 99.3% sparsity, CTS_KL reaches 74.0% top-1 accuracy in 7.9 minutes, while LTR reaches 68.3% at the same sparsity in 95.2 minutes.

• At 1.44% density on ResNet-20 / CIFAR-10 (Table I), LTR attains 80.90% test accuracy with 3058 epochs, Quick CTS_KL attains 79.04% with 180 epochs, and CTS_KL attains 80.26% with 320 epochs. LTR, Quick CTS_KL, and CTS_KL all pass sanity checks and avoid hyperparameter tuning; LTR is marked as passing both.

• Saliency baselines at the same 1.44% density: SNIP 67.73% (160 epochs), GraSP 62.59% (160 epochs), SynFlow 70.18% (161 epochs) — all listed as failing sanity checks; Edge-popup 10.00% (320 epochs, fails sanity checks); Gem-Miner 77.89% (320 epochs, passes sanity checks but does not avoid hyperparameter tuning, and the authors note its authors tuned training hyperparameters extensively).

• Large computation reductions: at roughly 200 times compression (0.47% density), the proposed method draws tickets in 192 times fewer training iterations (Quick CTS) and 24 times fewer training iterations (CTS) than LTR on CIFAR-10 tasks, and the speedup increases as subnetworks become sparser.

• Advantage over saliency methods is consistent; advantage over LTR is regime-dependent: CTS outperforms saliency-based methods in the sparsity-accuracy tradeoff across all sparsities, but its advantages over LTR emerge most clearly only in the highly sparse regime.

• Rewinding detail: results in Table I use rewinding iteration k = 3000, i.e., epoch 7.5.

• Sanity checks: the drawn tickets pass sanity checks proposed in the cited works (random shuffling of the pruning mask within each layer and reinitialization are the examples discussed), unlike SNIP, GraSP, SynFlow, and Edge-popup.

• Objective behavior at initialization (Figure 1, VGG-16 on CIFAR-10): tickets drawn by saliency pruning rarely outperform baselines even on the objectives they are designed to optimize; LTR is not a pruning-at-initialization method, so its ticket need not be strong at initialization.

Methodology in Plain English

CTS treats the choice of which weights to keep as a single big optimization problem over the entire mask, rather than scoring each weight on its own.

• Probabilistic mask: each weight gets a trainable retention probability alpha_j, and the mask is a sample from independent Bernoulli distributions with those probabilities. The expected number of active weights is the sum of the alphas.

• Differentiable sampling: because sampling a discrete mask blocks gradients, the authors use the Concrete relaxation (Gumbel-softmax). Each soft mask element is a sigmoid of (logit(alpha_j) + logistic noise) divided by a temperature tau. This produces low-variance, nearly unbiased gradients and avoids the weight stagnation the authors attribute to straight-through estimators.

• Hard mask at the end: after search, the final mask is the top-kappa fraction of retention probabilities.

• Sparsity enforcement: the constraint is written as L_sparsity = E[||m||_0] / (kappa*d) - 1, which self-scales across target densities and approaches zero when the target is met. GradBalance sets the balancing coefficient to the ratio of the objective gradient norm to the sparsity gradient norm, zeroes it out once the constraint is satisfied, and smooths it with an exponential moving average. The effective density is set slightly above the requested one (kappa_eff = 1.1 kappa). Alphas are initialized at kappa so the constraint starts at zero.

• Frozen weights, single shot: ticket search happens once, briefly after a short initial training period of k steps, with network weights frozen — so the search is not disturbed by stochastic gradient noise. Training then continues to convergence. The Concrete temperature tau is deliberately not annealed, and data augmentation is skipped during search.

• Objectives: candidates are plain task loss, relative change in loss (related to SNIP), negative gradient norm (related to GraSP), reverse KL between softmax outputs of ticket and dense parent, per-layer-normalized feature MSE, and per-layer-normalized gradient MSE. Teacher–student style objectives need two forward passes per step.

• Truncation note: the paper content supplied ends partway into Section V (Experiments), at the mention of VGG-16 and ResNet-20 on CIFAR-10 and the start of another dataset name ("Re..."). Results beyond the abstract, Figure 1, and Table I are therefore not described here.

Why This Matters

Impact on research: The paper challenges a widely used assumption in Pruning-at-Initialization — that per-weight, first-order saliency scores are sufficient. It argues the failure of diverse PaI methods stems from ignoring weight interdependence, and offers a holistic search formulation plus a hyperparameter-light sparsity controller as an alternative. It also connects pruning objectives to training-dynamics preservation and knowledge distillation.

Real-world applications:

• Deploying very sparse models on edge devices and embedded hardware where memory and compute are tightly constrained.
• Reducing the cost of model compression pipelines, since ticket search avoids the massive repeated retraining of LTR.
• Compressing large pretrained image classifiers for inference-serving systems where retraining budgets are limited.
• Enabling pruning research on models where iterative retraining is impractical (for example, very deep or very large backbones), subject to the sparse-training hardware support the authors mention.

Industry relevance: The headline practical claim is cost. CTS_KL reportedly matches or exceeds LTR accuracy at high sparsity in a small fraction of the time and, in the Table I comparison, Quick CTS_KL gets within roughly two accuracy points of LTR using 180 epochs instead of 3058. That ratio is directly relevant to teams deciding whether sparse-model training is economically viable.

Future Directions

• Closing the dense-regime gap with LTR. The authors state GradBalance applies the sparsity constraint one-sidedly: if the network becomes sparser than the target, retention probabilities are driven only by the objective gradient, which can produce overly sparse tickets at high density (large kappa). They note CTS is designed primarily for the high-sparsity regime.
• Efficient hardware-level sparse training. The authors motivate the frozen-weight, single-shot search partly by a hypothesized future with hardware support for efficient sparse training, which would turn these search savings into real training-computation reductions.
• Extending objective design. The three proposed objectives (reverse KL, feature matching, gradient matching) open the question of which training-dynamics signal best characterizes a genuine lottery ticket, and whether others exist.
• Broader architecture and dataset coverage. The supplied content only shows VGG-16 and ResNet-20 on CIFAR-10 plus a reference to LTR comparisons; the full experimental scope beyond this excerpt is not available here, and the paper's own summary states experiments on varying image classification tasks, leaving the generalizability of CTS to other domains open.

Target Audience

Researchers and graduate students working on neural network pruning, sparsity, the Lottery Ticket Hypothesis, or training dynamics; practitioners who need to compress image classifiers but cannot afford iterative retraining; and readers comfortable with constrained optimization, reparameterized gradient estimators, and knowledge distillation. Beginners will struggle with the methodology section without prior exposure to Concrete/Gumbel-softmax relaxations.

Authors’ abstract

The Lottery Ticket Hypothesis asserts the existence of highly sparse, trainable subnetworks ('winning tickets') within dense, randomly initialized neural networks. However, state-of-the-art methods of drawing these tickets, like Lottery Ticket Rewinding (LTR), are computationally prohibitive, while more efficient saliency-based Pruning-at-Initialization (PaI) techniques suffer from a significant accuracy-sparsity trade-off and fail basic sanity checks. In this work, we argue that PaI's reliance on first-order saliency metrics, which ignore inter-weight dependencies, contributes substantially to this performance gap, especially in the sparse regime. To address this, we introduce Concrete Ticket Search (CTS), an algorithm that frames subnetwork discovery as a holistic combinatorial optimization problem. By leveraging a Concrete relaxation of the discrete search space and a novel gradient balancing scheme (GRADBALANCE) to control sparsity, CTS efficiently identifies high-performing subnetworks near initialization without requiring sensitive hyperparameter tuning. Motivated by recent works on lottery ticket training dynamics, we further propose a knowledge distillation-inspired family of pruning objectives, finding that minimizing the reverse Kullback-Leibler divergence between sparse and dense network outputs (CTS-KL) is particularly effective. Experiments on varying image classification tasks show that CTS produces subnetworks that robustly pass sanity checks and achieve accuracy comparable to or exceeding LTR, while requiring only a small fraction of the computation. For example, on ResNet-20 on CIFAR10, it reaches 99.3% sparsity with 74.0% accuracy in 7.9 minutes, while LTR attains the same sparsity with 68.3% accuracy in 95.2 minutes. CTS's subnetworks outperform saliency-based methods across all sparsities, but its advantage over LTR is most pronounced in the highly sparse regime.