Always-Sparse Training with Guided Stochastic Exploration

I was quite confused as to why the authors were using Sparse-dense matrix multiply, instead of SpGEMM. Then I realized that SpGEMM would show up in convolutional filter pruning, whereas Sparse-dense would show up in connection weight pruning.

I am still very worried about the plots in Appendix L on evaluating this relationship. For example, in both plots, the x-axis starts at 6 ticks off from 90% sparsity. So the left-most datapoint is already at 30% sparsity.

Here's the problem, from what I can see. Looking at solely the largest (purple) evaluation. If you look at the less massively parallel architecture (CPU) with much less SIMD, the point at which Sparse-dense becomes faster than dense-dense is much earlier than in the V100 architecture. On the CPU it's around 91% sparse. For the GPU it's around 97-98% sparse. Extrapolating this to more modern hardware such as H100 or A100, is the relationship going to be even worse? What about extrapolating to multi-gpu execution?

My specific comments

• If the purpose of this paper is to reason or advance the field of "biologically plausible neural networks" that may have applications to neuroscience, could the authors cite why this paper would be valuable to the neuroscience community? Could the authors make that argument iron-tight and clear before it is considered under the peer review setting?

• If the purpose of this paper is to connect the field of Artificial Neural Networks with "biological neural networks" in brains of living creatures, could the author expound on this greatly? One place to start would be how actual neurons in actual brains having spiking activity called spike trains, and don't behave much like artificial neural networks at all. Spiking neural networks are also a good place to start on this path.

Was to broadly evaluate this paper within the context of the intersectional field of ANNs and neuroscience, which from what I understand is a well established field. It could always be the case that developments in ANN pruning can improve our understanding of neuroscience and vice versa.

Given that the authors say that this is not the case. After looking at Appendix L, here are my followup thoughts:

• This is very much tacked on, where the authors move the actual difficult question to "outside the scope of this work."
• The reviewer believes that this should be the scope of the work, being able to show practical improvements.
• The authors themselves say it requires 90%+ or higher sparsity to achieve any speedup at all. Do the authors believe this tradeoff is worth it?
• Note that sparsifying matrices naturally degrades the performance of the neural network. So following this line of reasoning, I would have to sparsify up to 95%? 99%? in order to achieve a reasonable speedup. How much would the performance degrade by that point?
• It's very difficult to reason about the plots in Appendix L without also seeing, in the same plot, the performance degradation for that amount of speedup.

If I have missed something, I am happy to withdraw my comment.

Thank you for taking the time to review our paper. We respectfully disagree with your assessment. There is an established line of peer-reviewed work that aims to deliver the highest accuracy for a given FLOPs budget by training sparse models. In this line of work, our method improves upon both the FLOPs and the accuracy, thus yielding a strict improvement over previous methods. Our time complexity analysis further supports this claim.

The vast majority of the DST literature uses FLOPs to compare the efficiency of various methods. This is because support for efficient sparse implementations of common neural network operations is currently missing in machine learning frameworks, as noted in our limitations section, making wall-time experiments challenging. However, we also explicitly highlight that real speedups, even with unstructured sparsity, can be obtained at high levels of sparsity as demonstrated in Appendix L.

We want to stress that we make no claims throughout the paper about any connections between our method and biologically plausible neural networks. We mention the sparsity level of some biological brains to indicate that such high levels of sparsity exist. However, even at the levels of sparsity that are common within the sparse ANN literature, our method provides improvements in FLOPs as shown in Figure 4.

While we disagree with your assessment, we still greatly appreciate that you took the time to review our work and we hope that our explanation helps to clarify our approach.

We appreciate the reviewer's timely reply. We fully agree that there is an opportunity to establish connections between the presented work and neuroscience. With our previous comment we wanted to clarify that we do not investigate such connections explicitly in our current work, but we do see the value in what our method enables: efficiently training large models with the sparsity of biological neural networks. We are eager to discover the insights that will come from investigating these connections.

A quick aside on terminology, we will use MM to denote dense-dense matrix-multiply, SpMM for sparse-dense matrix-multiply, and SpSpMM for sparse-sparse matrix-multiply.

While in the limitations section we state that "the significant engineering required to develop an optimized [sparse] GPU kernel falls outside the scope of this work", this does not in any way mean that we consider it outside the scope of our work to present evidence for the practical improvements of our method, quite the contrary. We present the SpMM benchmark in Appendix L and extend the CIFAR experiments to 98% sparsity exactly because we are aware of limitations within the sparse training literature which heavily relies on FLOPs comparisons rather than showing practical speed improvements.

It is important to note that the benchmark in Appendix L does not represent the limit of SpMM efficiency, it should rather be interpreted as a snapshot of the current state of SpMM in machine learning frameworks (PyTorch in particular). It is entirely plausible that advances in both hardware and software will improve the efficiency of SpMM, see for instance two recent works [1, 2]. Perhaps the reviewer was referring to the COO sparse matrix format results in Appendix L, which are indeed less efficient on the GPU, breaking even around 98% sparsity. However, using the CSR sparse matrix format, both the CPU and GPU start showing speedups from 90% sparsity. Moreover, at 99% sparsity, SpMM is around 10 times faster on both the CPU and GPU, compared to MM. Despite concerns by the reviewer, there is no reason to believe that these results will be affected by multi-GPU execution using either data or model parallelism.

Even with the current state of SpMM, we can thus conclude that real speedups are possible. If we consider the reported VGG-16 results on CIFAR-10, at 98% sparsity a 3.9% drop in accuracy is obtained compared to the dense model. However, 98% sparsity would result in close to 10 times faster training. We believe that there certainly are scenarios where trading 3.9% accuracy for an almost 10 times speedup is entirely reasonable.

We greatly appreciate the reviewer's willingness to engage in dialogue about the presented work, and the opportunity to underscore important considerations regarding our work.

[1] Huang, Guyue, et al. "Ge-spmm: General-purpose sparse matrix-matrix multiplication on gpus for graph neural networks." SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020.

[2] Gerogiannis, Gerasimos, et al. "SPADE: A Flexible and Scalable Accelerator for SPMM and SDDMM." Proceedings of the 50th Annual International Symposium on Computer Architecture. 2023.

We appreciate the reviewer's insightful comments and have carefully considered the concerns raised. We acknowledge that many of the DST methods, including ours, are designed to benefit the training of large models at high levels of sparsity. However, performing those experiments is extremely costly and would severely limit the reproducibility of the results. It is important to emphasize that the results on FLOPs presented in Figure 4 already take into account that RigL computes the dense gradients periodically. Despite this, our method is able to reduce the overall training FLOPs, especially at high levels of sparsity.

Thank you for bringing the interesting Powerpropagation work to our attention. Both In-Time Over-Parameterization (ITOP) and Powerpropagation require a DST method and are therefore orthogonal improvements. They could thus be combined with our method without undermining our core contribution. ITOP performs an extensive hyperparameter search while Powerpropagation proposes to raise the weights (element-wise) to a constant positive power, which they show results in a rich-get-richer dynamic of the weights such that the large magnitude weights get larger while the small magnitude weights get smaller. They state that this could be useful in combination with sparse training methods as they would benefit from a regularizer which promotes weight sparsity.

Let us clarify the notation used for the sampling distributions. The vectors $f^{[l]}$ and $g^{[l]}$ contain the probability that each input/output is sampled in the $l$-th layer. Let's say layer $l$ has $n$ inputs and $m$ outputs, then $f^{[l]} \in R_{+}^{n}$ and $g^{[l]} \in R_{+}^{m}$, that is, they are non-negative real vectors. Sampling one index from $f^{[l]}$ and one from $g^{[l]}$ fully specifies a weight. Each weight is thus sampled by the probability specified by the outer product $f^{[l]} \otimes g^{[l]}$, which is never explicitly computed. The computation of $f^{[l]}$ and $g^{[l]}$ is performed with information local to the $l$-th layer. With GraBo, for example, the (element-wise) absolute of the activations $|h^{[l-1]}| \in R_{+}^{n \times b}$, where $b$ is the batch size, are summed over the batch size in the numerator and summed over both the batch size and the inputs in the denominator to ensure that $f^{[l]}$ is a proper probability distribution. The same goes for the definition of $g^{[l]}$ but it uses the (element-wise) absolute of the gradients at the outputs $|\delta^{[l]}| \in R_{+}^{m \times b}$ instead. We will adjust the explanation of our method in the paper accordingly to improve the clarity.

Once again, we appreciate your feedback and suggestions, and we hope that our explanation helps to clarify our approach.

Thank you for your valuable feedback on our paper. Regarding your inquiry about the results in Figure 2, the details you are looking for can be found in the 4th column of the figure, where the subset sampling factor is extended up to 100. Interestingly, this column shows exactly the specific behavior you mentioned. In Section 4.2, we explicitly state that "we expect that as the subset size of GSE increases, the accuracy eventually drops back to the level of RigL because the largest gradient magnitude selection then becomes the dominant criteria, making the method more similar to RigL. To this end, we experimented with increasing the number of samples in the last column of Figure 2, note the logarithmic x-axis. The observations align precisely with our expectations." We will update the caption of Figure 2 to explicitly highlight the difference of the x-axis.

While our method adopts commonly used weight magnitude pruning and gradient magnitude growing criteria from the existing literature, our novelty stems from a strategic use of these common criteria. Instead of naively applying them over all possible options, we propose a method that applies the criteria on an informed subset of the weights which we show improves both efficiency and accuracy. Applying common criteria does not diminish the innovation, rather, it grounds our work on theoretically-motivated and peer-reviewed research.

The differences in accuracy are more prevalent at higher levels of sparsity, at 98% sparsity the proposed method outperforms the others by up to 2%. We believe therefore that our method constitutes more than just a marginal improvement, especially when considering that our method is also strictly more efficient.

We acknowledge your concern regarding the sampling distributions and wish to emphasize a crucial aspect central to our paper's premise and the proposed method. Contrary to the impression raised, the design of the sampling techniques ensures that they do not necessitate the computation of gradients for all weights. The core innovation lies exactly in the complete avoidance of dense computations and materialization throughout the entire training process. Our sampling approaches exclusively require the input activations and output gradients of a layer, these are significantly more efficient to obtain than the gradients for all the weights of a layer. This efficiency is attributed to the utilization of sparse backward passes rather than the more resource-intensive dense computations.

Let us clarify the sampling process. The vectors $f^{[l]}$ and $g^{[l]}$ contain the probability that each input/output is sampled in the $l$-th layer. Let's say layer $l$ has $n$ inputs and $m$ outputs, then $f^{[l]} \in R_{+}^{n}$ and $g^{[l]} \in R_{+}^{m}$, that is, they are non-negative real vectors. Sampling one index from $f^{[l]}$ and one from $g^{[l]}$ fully specifies a weight. Each weight is thus sampled by the probability specified by the outer product $f^{[l]} \otimes g^{[l]}$, which importantly is never explicitly computed. The computation of $f^{[l]}$ and $g^{[l]}$ is performed with information local to the $l$-th layer. With GraBo, for example, the (element-wise) absolute of the activations $|h^{[l-1]}| \in R_{+}^{n \times b}$, where $b$ is the batch size, are summed over the batch size in the numerator and summed over both the batch size and the inputs in the denominator to ensure that $f^{[l]}$ is a proper probability distribution. The same goes for the definition of $g^{[l]}$ but it uses the (element-wise) absolute of the gradients at the outputs $|\delta^{[l]}| \in R_{+}^{m \times b}$ instead. We will adjust the explanation of our method in the paper accordingly to improve the clarity.

Thank you again for your valuable feedback. We appreciate the opportunity to clarify these nuances and underscore the careful considerations made in our method's design, which improves upon both the efficiency and accuracy of previous methods.