PaperlayerView exact revision on arXiv

Source-linked AI summary

Top-KAST: Top-K Always Sparse Training

Siddhant M. Jayakumar, Razvan Pascanu, Jack W. Rae, Simon Osindero, Erich Elsen

arXiv:2106.03517v1cs.LGstat.ML

TL;DR

Finding effective sparse networks has often required dense training, even as large models increase resource demands. Top-KAST keeps training sparse by using top-magnitude weights in the forward pass and a larger update set in the backward pass. It reports comparable ImageNet performance to dense-to-sparse methods while remaining sparse, and strong results on small language models.

  • Problem

    Finding performant sparse subsets has often required dense-model training, despite the resource demands of increasingly large models.

    Source: §1 · p. 1

  • Method

    Top-KAST uses the top-magnitude parameter subset in the forward pass and updates a larger subset in the backward pass, keeping both passes sparse.

    Source: §1 · p. 2

  • Results

    Top-KAST remains sparse while matching dense-to-sparse methods on ImageNet and reports 1.00 bpc at 55M parameters versus a 0.99 baseline at 277M parameters for small language models.

    Source: §6 · p. 9

  • Takeaways & Limitations

    Practitioners can choose forward and backward sparsities according to their available resource budget.

    Source: §6 · p. 9

    The method is not a strictly valid coordinate descent method on either α or θ.

    Source: §2.3 · p. 4

Abstract

from arXiv · show

Sparse neural networks are becoming increasingly important as the field seeks to improve the performance of existing models by scaling them up, while simultaneously trying to reduce power consumption and computational footprint. Unfortunately, most existing methods for inducing performant sparse models still entail the instantiation of dense parameters, or dense gradients in the backward-pass, during training. For very large models this requirement can be prohibitive. In this work we propose Top-KAST, a method that preserves constant sparsity throughout training (in both the forward and backward-passes). We demonstrate the efficacy of our approach by showing that it performs comparably to or better than previous works when training models on the established ImageNet benchmark, whilst fully maintaining sparsity. In addition to our ImageNet results, we also demonstrate our approach in the domain of language modeling where the current best performing architectures tend to have tens of billions of parameters and scaling up does not yet seem to have saturated performance. Sparse versions of these architectures can be run with significantly fewer resources, making them more widely accessible and applicable. Furthermore, in addition to being effective, our approach is straightforward and can easily be implemented in a wide range of existing machine learning frameworks with only a few additional lines of code. We therefore hope that our contribution will help enable the broader community to explore the potential held by massive models, without incurring massive computational cost.

1 Introduction

Sparse networks can represent learned functions with few weights, but finding effective sparse subsets has often required dense training. Top-KAST keeps training sparse while adapting its active weights, and the paper reports strong results on ImageNet and language modeling.

1 Introduction

  • Sparse subsets can represent a converged model’s learned function at high accuracy, and large sparse models have been reported to outperform smaller dense models at fixed parameter and FLOP budgets.The Lottery Ticket Hypothesis [9] helped spur interest in training sparse neural networks [44].Source: §1 · p. 1
  • Dense-to-sparse approaches typically train a dense model first, while sparse-to-sparse training starts at the desired sparsity but has only more recently drawn attention.RigL uses occasional gradient information to update masks, but achieving its theoretical bounds and avoiding dense materialization can be difficult in common frameworks.Source: §1 · pp. 1–2
  • High-performance sparse training is motivated by the resource demands of state-of-the-art models, which can require large clusters that many researchers cannot access.Training compute footprints and energy consumption also raise environmental, moral, and economic concerns [33] [37].Source: §1 · pp. 1–2
  • Top-KAST selects top-magnitude parameters for the forward pass and updates a larger parameter subset, without dense forward passes or dense gradients.An auxiliary exploration loss encourages the mask to adapt during training.Source: §1 · p. 2
  • Top-KAST matches RigL on ImageNet across FLOP budgets and outperforms SET, while achieving state-of-the-art language-modeling performance on enwik8 for small models.The method also aims to remain sparse throughout training.Source: §1 · p. 2

2 Method: Top-KAST

Top-KAST aims to retain a desired final sparsity while avoiding dense parameters and gradients during training. It selects a sparse forward-pass view, updates a larger sparse coordinate set to explore masks, and adds a regularizer and implementation choices to support that process.

2 Method: Top-KAST

  • Top-KAST targets a desired final sparsity with overheads near fixed-topology sparse training, unlike dense-to-sparse methods that do not meet the overhead criterion.The method is proposed as an alternative that also seeks high accuracy for a given training-FLOP budget and easier integration into existing frameworks.Source: §2 · pp. 2–3

2.1 Sparse Forward Pass

  • At each step, α_t retains the top D-proportion of θ_t by magnitude, with selection performed per layer; this sparse view is used in both passes while updates apply to θ.The sparse view is reconstructed deterministically from θ, and magnitude selection is presented as an inexpensive estimate of parameter contribution.Source: §2.1 · p. 3

2.2 Sparse Backward Pass

  • The update set B includes active coordinates A plus the next-largest M-proportion, enabling sparse gradients and mask exploration before the active set stabilizes.Training is described as an exploratory stage with changing active sets followed by refinement around a stable mask; no θ-sized dense vector is needed in either pass.Source: §2.2 · p. 4

2.3 Exploration Regularisation Loss

  • The exploration regularizer penalizes updated weights, with stronger penalties on B \ A scaled inversely with D to reduce near-tied mask changes, especially at high sparsity.The authors note that the approach is not strictly valid coordinate descent on either α or θ.Source: §2.3 · p. 4

2.4 Implementation of Top-KAST

  • Top-K selection each iteration may add overhead or require dense-model memory, but CPU-side selection and updates every 100 steps are described as ways to reduce costs.The paper reports comparable results when the operation is performed every 100 steps.Source: §2.4 · p. 5

3 Related Work

Prior sparse-training methods range from dense-to-sparse pruning to sparse methods based on evolving connectivity or constrained gradient updates. Top-KAST is presented alongside these approaches as a sparse-to-sparse alternative, with comparable performance to RigL while alleviating implementation issues.

Dense-to-sparse methods

  • Dense-to-sparse methods such as magnitude pruning use dense weight or gradient information during training, limiting sparse-model scale to the largest trainable dense model.The paper names L0 regularization, variational dropout, discovering neural wirings, and soft weight threshold reparameterization as examples.Source: §3 · p. 5

Evolutionary methods

  • Earlier TWEANN research evolved sparse networks, but largely did not focus on sparsity's potential benefits for model size and efficiency.The passage also notes more recent work evolving sparse topologies.Source: §3 · p. 5

Deep Rewiring

  • Deep Rewiring fixes weight signs and zeros weights when signs would flip, while using a random walk over weights and connectivity.Despite theoretical convergence proofs, its practical performance is described as lagging behind later work.Source: §3 · p. 5

Sparse Evolutionary Training (SET)

  • Sparse Evolutionary Training (SET) drops weights by magnitude and adds connections randomly, whereas Top-KAST uses gradient information for additions.The authors report Top-KAST's gradient-based additions as more performant than random additions.Source: §3 · p. 5

Dynamic Reparameterization

  • Dynamic Reparameterization shifts parameter budgets between layers but does not constrain the resulting FLOPs, which may increase.The method aims to place capacity where it is most effective.Source: §3 · p. 5

Rigging the Lottery (RigL)

  • RigL uses occasional dense-gradient calculations to choose parameters to wake up; Top-KAST performs comparably while alleviating implementation issues.The paper notes that RigL's theoretical sparsity-proportional cost is difficult to achieve in current frameworks.Source: §3 · pp. 5–6

4 Experiments: ImageNet

On ImageNet, Top-KAST remains sparse throughout training while matching or exceeding several comparison methods under the reported settings. The ablations examine how extra backward-pass units and mask dynamics relate to performance and exploration.

4 Experiments: ImageNet

  • Top-KAST is comparable to dense pruning methods at constant FLOPs and outperforms SET and static random sparsity patterns while remaining sparse throughout.The reported comparisons include 90% and 95% sparsity and settings that relax the assumption of dense first and last layers.Source: §4 · p. 6
  • Top-KAST performs comparably to RigL per FLOP, surpasses it at 98% and 99% sparsity, and requires only the bucket B size as a hyperparameter.Matching RigL's performance by average backward sparsity requires slightly higher densities for Top-KAST; RigL's periodic dense-gradient calculations are difficult to implement efficiently in current frameworks.Source: §4 · p. 6

4.1 Ablation studies

  • Randomly selecting extra backward-pass units performs better at 90% sparsity but worse at higher sparsities than selecting the next-highest-magnitude units.The authors argue that targeting high-magnitude units matters more in very sparse settings, where random selection may take many iterations to reach relevant weights.Source: §4.1 · p. 7
  • Removing exploration units from the start harms performance, while 5,000 training steps with them considerably boost it and stopping at 16,000 steps recovers most benefits.The results support an exploration phase followed by refinement on a mask that stays more or less constant during the latter half of training.Source: §4.1 · p. 7
  • The mask stabilizes over time, and only about 5% of initially unused reservoir units ever activate, mostly early in training.The mask-change analysis uses a run with 80% forward sparsity and 50% backward sparsity.Source: §4.1 · p. 7, Figure 3 · p. 7

5 Experiments: Language Modeling

Top-KAST is evaluated on Enwik8 and WikiText-103 with Transformer-XL language models. The reported results retain performance at 80% sparsity on both benchmarks, with performance drifting beyond 90% sparsity on WikiText-103.

5 Experiments: Language Modeling

  • The experiments use Enwik8, a character-level benchmark, and WikiText-103, a word-level benchmark, with the long-range Transformer-XL variant.The paper motivates language-model sparsification by prior observations that performance improves with increased data and model parameters.Source: §5 · p. 8
  • On Enwik8, Transformer-XL + Top-KAST achieves state-of-the-art performance among previously published models under 60M parameters and remains comparable to the 24-layer dense baseline through 80% sparsity.The dense baseline obtains 0.99 bits-per-character (BPC); Table 2 reports test BPC for small models.Source: §5 · p. 8, Table 2 · p. 8
  • On WikiText-103, Top-KAST achieves 80% sparsity with minimal performance degradation, but performance begins to drift beyond 90% sparsity.The 16-layer dense baseline obtains 18.3 test perplexity; Table 3 reports test perplexity across forward-backward sparsities.Source: §5 · p. 8, Table 3 · p. 8

6 Conclusion

Top-KAST keeps training sparse in both forward and backward passes, using high-magnitude weights for the forward pass and additional exploration weights for the backward pass. The authors report competitive ImageNet results and state-of-the-art language-modeling results for small models, while noting that sparse hardware integration remains unfinished.

6 Conclusion

  • Top-KAST stays sparse in both forward and backward passes and is easily implemented with modern neural-network packages.It keeps the highest-magnitude weights for the forward pass and an additional set of exploration weights for the backward pass.Source: §6 · pp. 8–9
  • On ImageNet, Top-KAST achieves performance comparable to dense-to-sparse methods and exceeds several sparse-to-sparse methods while remaining sparse.The authors also report successful sparsification of Transformers for language modeling and state-of-the-art results for small models: 1.00 bpc at 55M parameters versus a 0.99 baseline at 277M.Source: §6 · pp. 8–9
  • Further work is needed to integrate Top-KAST fully with sparse hardware and appropriate sparse kernels.Source: §6 · p. 9

Broader Impact

The authors argue that sparse training could support larger models on a given device and reduce training and deployment resource use, while recommending investigation of possible differential sensitivity to data imbalances.

Broader Impact

  • Sparse training could let practitioners train larger models on existing hardware and reduce computational and energy costs, but extensions should assess sensitivity to class imbalance relative to dense counterparts.Source: Broader Impact · p. 9

A Language Modeling

The appendix describes Transformer-XL baseline settings and compares Top-KAST with magnitude pruning on a smaller Enwik8 model because the large-model pruning implementation required additional memory. The comparison finds pruning slightly ahead when Top-KAST uses a dense backward pass, while Top-KAST remains competitive with sparse backward passes.

A Language Modeling

  • The dense Transformer-XL baseline uses the large-model specification from Dai et al. [6], with 285M parameters, and its hyperparameters are listed in Table 4.Training uses 4,000 warmup steps from 1e-7 to 2e-4 followed by cosine decay.Source: A Language Modeling · p. 12, Table 4 · p. 12
  • Magnitude pruning could not be implemented for the large Transformer-XL model because of its additional memory requirements, so the comparison uses a 69M-parameter model.The smaller model differs by using d_model = 512, d_ff = 1536, and 8 heads.Source: A Language Modeling · p. 12
  • Pruning slightly outperforms Top-KAST when Top-KAST uses a dense backward pass, while Top-KAST remains competitive with sparse backward passes.Table 5 reports test perplexity for the smaller Transformer model.Source: A Language Modeling · p. 12, Table 5 · p. 12

B ImageNet

ImageNet experiments use ResNet-50 with a specified training schedule, generally keeping the first and last layers dense while also evaluating an all-layers-sparsified setting.

B ImageNet

  • ResNet-50 is trained for 32,000 steps with batch size 4,096, learning rate 1.6, and drops by 0.1 at epochs 30, 70, and 90.The learning rate ramps linearly for five epochs; Top-KAST uses weight decay 1e-4 across a range of forward and backward sparsities.Source: B ImageNet · p. 12
  • The experiments generally keep the first and last layers dense, but also consider sparsifying every layer.The all-layers-sparsified setting is labeled as using 1x training.Source: B ImageNet · pp. 12–13
  • An original publication’s parameter count was corrected from 255M to 285M.The passage does not specify which model the count refers to.Source: B ImageNet · p. 12

C Implementation of RigL and Top-KAST

The section contrasts the implementation burden of RigL’s dense-gradient Top-K with Top-KAST’s sparse training path, and gives pseudocode for selecting forward and backward parameters. It also describes an implementation that refreshes Top-K every 100 steps.

C Implementation of RigL and Top-KAST

  • RigL’s dense-gradient Top-K must occur inside gradient calculation, requiring chunked, serial computation and substantial gradient-code rewrites that may affect performance.The proposed path passes locations and a step count through the gradient calculation; chunks are all-reduced before updating the running Top-K.Source: C Implementation of RigL and Top-KAST · p. 13
  • With Top-K refreshed every 100 steps, index and weight communication is occasional, and CPU-parallel Top-K avoids requiring data or forward passes on the accelerator.The accelerator can use only the sparse weights; Table 6 is identified as reporting results at different Top-K frequencies, but no values are provided here.Source: C Implementation of RigL and Top-KAST · p. 13, Table 6 · p. 13

D Pseudocode

  • Top-KAST selects separate X% forward and Y% backward parameter sets; its pseudocode updates backward parameters and uses explicit masking with dense kernels.A truly sparse implementation requires custom sparse kernels.Source: D Pseudocode · pp. 13–14

Related papers

Loading 2106.03517v1…