Paper: Always-Sparse Training with Guided Stochastic Exploration (ICLR-2024)

Reviewer comment (Experiments):
The paper delivers no actual empirical results useful in practice. As deep learning is a practical and engineering field, it is not really known what is the value of showing improvements in big-O format, when actual implementation of these approaches are in super huge networks such as ChatGPT, GPT4, Diffusion Models, and many other productionized models. There are open source minified models of the above as well that can be run on a desktop or a laptop. The reviewer believes that there are no practical benefits of this approach in that setting either.

Author response:
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.

Reviewer comment (Evaluation):
The paper considers sparsifying weight connections; it is well known that this cannot, in practice, deliver any actual improvement on massively parallel SIMD architectures used in deep learning today. Sparse general matrix-matrix multiplies are known to be, more or less, impossible to speed up in a SIMD setting where the actual observed wall-clock time speedup is anywhere close to the theoretical speedup indicated by FLOP reduction. There are several peer-reviewed works exploring this problem.

Author response:
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.

Reviewer comment (Theory):
The connection between artificial neural networks and actual neural networks in the brains of living creatures is not well thought out, and will be covered in more detail in the "Questions" section of the review.

Author response:
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.

Reviewer comment (Experiments):
While the value proposition of the method is the training of large-scale models that cannot fit onto consumer or high-end GPU, experiments in this work are done at small scale, amenable to researchers and practitioners with limited compute power. The potential advantage of the method could be in reduction of the FLOPs for backward pass, but the baseline RigL requires materialization of dense gradient only on update steps, done typically every 100–1000 steps. Thus, the computational savings are minor. At this scale, a method allowing for higher exploration of sparsity masks is likely to get better performance [1], [2] for a given target sparsity and training time.

Author response:
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.

Reviewer comment (Experiments):
Some more recent works with always sparse training managed to significantly outperform RigL. Powerpropagation+Top-KAST [3] and ITOP-RigL [4] have noticeably better numbers for 90% sparsity on ResNet-50 while adopting the same training procedure.

Author response:
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.

Reviewer comment (Writing):
*Minor*. Notation for $f^{[l]}$, $g^{[l]}$ is confusing. $h_l$, $g_l$ - denote a scalar value - norm of the weights / gradients for a given layer, where $1$ is a vector of all ones with length - `(num_layers,)`. And there is expression in the denominator $1^T |h/g^{(l-1)}| 1$ that would imply, that $h^{(l)}$ is a matrix of size `(num_layers, num_layers)`. I think it would be better to write sum over $l$ in the denominator.

Author response:
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.

Reviewer comment (Reproducibility):
How is the norm of activations / gradient computed? Does one sum all elements in a tensor and take the absolute value, or computes L1, L2 norm. The first approach would yield zero in expectation for symmetric distributions and may not reflect presence of large elements in activation/gradient tensor, whereas the latter two yield the same value after multiplication by sign.

Author response:
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.

Reviewer comment (Novelty):
The novel insights presented in this paper are quite limited. Essentially, the proposed approach involves initially selecting a candidate subset from the weights that were not pruned, followed by weight selection based on existing criteria. Importantly, the same pruning and growth criteria are retained from previous methods.

Author response:
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.

Reviewer comment (Theory):
Regarding the distribution methods employed for sampling, the use of gradient-based distributions appears to contradict the proposed efficiency enhancement. This is because these distributions also require the computation of the full gradient, resulting in a complexity that is identical to RigL [1]. Furthermore, these distributions do not yield any performance improvements over RigL [1], indicating their relatively inconsequential role in this study.

Author response:
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.

Reviewer comment (Evaluation):
The improvements observed in the results are rather modest, typically falling within the range of 0.1% to 0.5%, thus demonstrating limited practical significance.

Author response:
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.

Reviewer comment (Experiments):
If the size of the candidate weights were to expand to encompass all weights, the performance should theoretically match that of RigL. However, Figure 2 in the current presentation only illustrates the scenario where the maximum value of $\gamma$ is set to 2. It would be insightful to have a comprehensive plot that covers a broader range of $\gamma$ values. The existing experimental results seem to suggest that a $\gamma$ value of 1 outperforms RigL (where $\gamma$ approaches infinity). It would be valuable to discern at what point the performance begins to decline for specific $\gamma$ values and provide an explanatory analysis in this regard.

Author response:
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.