Gradient Flow in Sparse Neural Networks and How Lottery Tickets Win

Abstract

Sparse Neural Networks (NNs) can match the generalization of dense NNs using a fraction of the compute/storage for inference, and have the potential to enable efficient training. However, naively training unstructured sparse NNs from random initialization results in significantly worse generalization, with the notable exceptions of Lottery Tickets (LTs) and Dynamic Sparse Training (DST). In this work, we attempt to answer: (1) why training unstructured sparse networks from random initialization performs poorly and; (2) what makes LTs and DST the exceptions? We show that sparse NNs have poor gradient flow at initialization and propose a modified initialization for unstructured connectivity. Furthermore, we find that DST methods significantly improve gradient flow during training over traditional sparse training methods. Finally, we show that LTs do not improve gradient flow, rather their success lies in re-learning the pruning solution they are derived from — however, this comes at the cost of learning novel solutions.

Unmasking the Lottery Ticket Hypothesis: What's Encoded in a Winning Ticket's Mask?

“…Our work identifies the mechanism by which IMP finds matching solutions: the algorithm maintains the information from the dense network about the loss landscape by encoding this information into the mask. Evci et al (2022) report similar findings but for a different setting. In their work, they construct sparse masks from a pruned solution (sparse network trained to convergence), where the latter is obtained through gradual magnitude pruning (GMP) throughout training (Zhu & Gupta, 2017) as opposed to IMP.…”

“…The two closest related works to our results are and Evci et al (2022). Both works consider linear mode connectivity between two networks of the same sparsity.…”

“…compare two networks trained from the same rewind point at the same pruning level but with different SGD noise. Evci et al (2022) find a pruning solution and mask using Gradual Magnitude Pruning (GMP; Zhu & Gupta (2017)) instead of IMP and apply the mask to the rewind point to obtain a lottery ticket solution. They then observe that these two sparse solutions are linearly mode connected, and close in function space.…”

Training Recipe for N:M Structured Sparsity with Decaying Pruning Mask

“…3) From-scratch with learned one-shot pruning pattern (Figure 1h) [13,11], which determines the sparsity pattern from the trained dense version and trains a sparse model from scratch. 4) From-scratch while learning sparsity pattern (Figure 1i) [51,6,14,26,9,4,34,58,10], which trains a sparse model from scratch while learning sparsity patterns simultaneously.…”

“…The irregularity of the sparsity pattern makes it challenging to be effectively leveraged by the dense accelerators such as GPU and TPU. The sparsified models often ends up with similar or worse performance (because of the extra complexity to compress and decompress the parameters) than their dense counterparts [2,32,43,21,30,15,59,50,10].…”