Top-KAST: Top-K Always Sparse Training
This addresses the computational bottleneck for scaling up large models in ML/AI, making massive sparse models more accessible, though it appears incremental relative to existing sparse training methods.
The paper tackles the problem of training sparse neural networks without instantiating dense parameters or gradients, proposing Top-KAST which maintains constant sparsity throughout training. Results show it performs comparably or better than previous methods on ImageNet while fully preserving sparsity, and it enables sparse language models with tens of billions of parameters to run with significantly fewer resources.
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.