Sparse maximal update parameterization: A holistic approach to sparse training dynamics

Summary

This work (similarly to \muP method) proposes a specific parameterization for weights, gradients and updates such that the optimal training hyperparameters (i.e. learning rate) are transferrable across different sparsity levels. The approach is validated on the task of language modeling within wide range of sparsities.

Strengths

* The proposed approach is theoretically sound and intuitive. * SμPar appears to produce a stable optimal learning rate with a wide range of sparsities, whereas the values for standard parameterization and dense \muP have to be tuned according to the level of sparsity. * SμPar allows one to achieve strong results for larger, outperforming the tuned SP and \muP configuration, given a small proxy tuned model. With significantly smaller tuning cost one can achieve better results.

Weaknesses

Whereas all the components in SμPar seem to be important, the importance of each individual parameterization is not studied. It could be the case that SμPar for weights while keeping \muP parameterization for gradients and learning rate may suffice and produce the same performance. For the sake of completeness, I would suggest ablating individual terms. Prior [1], [2] work shows that proper weight initialization (scaled according to sparsity) already improves training dynamics.

Questions

Does the SμPar transfer to other domains in the same way? For example, one could try pretraining a vision transformer in a setup similar to [3]. Sparsity scaling laws paper [3] showed that given sufficiently long compute, sparse models turn out to be compute-optimal. Given that SμPar finds better hyperparameters, could one expect that higher sparsity levels may be favored compared to the recipe used in the aforementioned work? It seems like the value of train loss decreases with the increase of sparsity (i.e., sparse models are better given the same number of parameters), whereas in Figure 10, the performance of SμPar loss seems to be independent of sparsity. I would expect an increase in the quality gap for longer training according to sparsity scaling laws. --- [1] Evci, Utku, et al. "Gradient flow in sparse neural networks and how lottery tickets win." Proceedings of the AAAI conference on artificial intelligence. Vol. 36. No. 6. 2022. [2] Liu, Zhuang, et al. "Rethinking the value of network pruning." arXiv preprint arXiv:1810.05270 (2018). [3] Frantar, Elias, et al. "Scaling laws for sparsely-connected foundation models." arXiv preprint arXiv:2309.08520 (2023).

Rating

7

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

See Weaknesses and Questions. Would be interesting to see whether the introduced parametrization is optimal for training methods with dynamic sparsity (RigL [1], Top-Kast [2], AC/DC [3]). --- [1] Evci, Utku, et al. "Rigging the lottery: Making all tickets winners." International conference on machine learning. PMLR, 2020. [2] Jayakumar, Siddhant, et al. "Top-kast: Top-k always sparse training." Advances in Neural Information Processing Systems 33 (2020): 20744-20754. [3] Peste, Alexandra, et al. "Ac/dc: Alternating compressed/decompressed training of deep neural networks." Advances in neural information processing systems 34 (2021): 8557-8570.

Summary

A parameterization method is proposed to ensure that the sparse models with different weight sparsity share the same set of optimal hyperparameters. The sparse models with S$\mu$PaR achieve lower validation accuracy using the same hyperparameters the dense model uses.

Strengths

- The paper is well-written and easy to follow. The experimental results support the claim well. - The proposed method links the maximal update and sparse network training, which is an interesting found.

Weaknesses

- The proposed method is discussed and evaluated only with the random unstructured sparsity pattern, which is hardly used in practice to speed up the DNN training and inference. - The theoretical contribution is somewhat weak. The proposed parameterization method is simple extension of maximal update parameterization to the sparse version, which I don't see significant technical improvement. - The experimental results should include the accuracy comparison between S$\mu$PaR and other sparse weight initialization methods (e.g., [1]). Although the proposed method was compared with SP and $\mu$P, the superiority of S$\mu$PaR is somewhat predictable since the other two are not designed for the sparse weights. I am curious how much improvement S$\mu$PaR can make from the previous sparse weight initialization schemes. - The scaling factors for the weight update need to be derived for each optimizer, which might limit the usability of S$\mu$PaR in practice. [1] Lee, Namhoon, et al. "A signal propagation perspective for pruning neural networks at initialization." arXiv preprint arXiv:1906.06307 (2019).

Questions

Major questions - Is the proposed parameterization method generalizable to any structured sparsity patterns (e.g., channel pruning) for the density-invariant hyperparameters? As mentioned in the paper, the unstructured sparsity is extremely hard to speed up. Discussing S$\mu$PaR on more structured sparsity patterns, such as channel pruning, will make the paper stronger. - Why does S$\mu$PaR stabilize the sparse network training even after multiple training steps? How is S$\mu$PaR technically different from the previous sparse weight initialization methods? Also, what is the accuracy difference between S$\mu$PaR and other sparse weight initialization methods? Minor questions - Line 113 about the weight update $\mathbf{X}^{l+1} + \Delta \mathbf{X}^{l+1}$ is confusing because the weight update is denoted by $\Delta \mathbf{W}^l$ in Line 105. - Is there any connection between the proposed method and the iterative magnitude pruning (IMP)used in Lottery Ticket Hypothesis [2]? Can S$\mu$PaR improve the final performance of IMP? [2] Frankle, Jonathan, and Michael Carbin. "The lottery ticket hypothesis: Finding sparse, trainable neural networks." arXiv preprint arXiv:1803.03635 (2018).

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have addressed the limitations and broader impacts in the paper.

Summary

The paper introduces Sparse Maximal Update Parameterization (SµPar), a novel approach designed to address challenges in training sparse neural networks. Sparse networks, despite their computational advantages, often struggle with signal propagation and optimal hyperparameter (HP) tuning across different sparsity levels. SµPar aims to stabilize training dynamics by ensuring that activations, gradients, and weight updates are scale-invariant with respect to sparsity. This method allows for consistent optimal HPs across varying model widths and sparsity levels, thus reducing the tuning costs associated with sparse models. The empirical results demonstrate that SµPar outperforms standard parameterization (SP) and maximal update parameterization (µP) in maintaining stable optimal HPs and achieving better loss metrics across different sparsity levels and model scales.

Strengths

1. The idea of this paper is interesting. The authors select a trending topic with very clear and strong motivation. The preliminary experimental results is very clear and helpful to the readers. 2. The discussion about forward/backward pass and weight updates are very clear. Figures are also clear and good to read. 3. The writing is good and the paper is well structured.

Weaknesses

1. This paper has many supporting experiments reporting the training loss, validation loss, and transfer loss for SµPar. However, loss values are not always reliable for evaluation. The accuracy data is a more convincing data to demonstrate the performance. I wonder what is the accuracy for the experiments. It would make reader easy to compare against other methods. 2. What is the method that is used to prune the network? There are lots of method such as SNIP and GraSP for static sparse training. However, this part is unclear according to the paper. 3. The author claims SµPar is a holistic solution. However, from what the paper has demonstrate (discussion, experiments), the reviewer cannot be convinced. According to weaknesses 2, we don’t know what is the pruning method used in sparse training, and according to the paper, the author also not performs applying SµPar on different sparse training methods. Therefore, we don’t know if SµPar works for different sparse training algorithms. Furthermore, I assume this paper only discuss static sparse training (such as SNIP), without discussing the performance on dynamic sparse training (DST) which is a more reliable and promising sparse training domain (e.g., better accuracy, flexible). Therefore it is not suitable to claim SµPar is a holistic solution. Even the authors don’t make such claim, without sufficient experiments the reviewer mentioned in this part, this paper is still lacks fundamental evidence to demonstrate effectiveness.

Questions

Please refer to weaknesses.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

Please refer to weaknesses.

Summary

This paper studies the effect of various hyperparameters on static sparse training while having a holistic approach. It highlights that there is a correlation between their settings and neural network performance (practically loss function values). The experiments are performed on smaller and larger models, including a large-scale language model.

Strengths

This is an interesting study with a relatively low level of originality (in my opinion). It practically puts together and discuss a bit more various findings from various papers. Also, the paper structure and clarity could have been improved. Currently, it is not easy to understand it. If the paper would be more mature, may have a relatively fair impact in the community.

Weaknesses

I believe that one main weakness is that the networks performance is reported just with loss values. For a more comprehensive view, and deeper understanding other performance metrics (task dependent) shall be also presented for all experiments. With respect to the problem motivation, I believe that the paper exaggerates a bit the difficulty of training static sparse neural networks. Perhaps, for static sparsity, the effect of hyperparameters on network performance is more pronounced, but the alternative dynamic sparsity is quite robust to hyperparameter choice (e.g., sparsity level) and can obtain quite easy performance on par (or even better) with dense neural networks while having much fewer theoretical computational requirements as reported in a number of related works. For a fair overview, the paper shall discuss also comparatively the performance of static sparsity against dynamic sparsity and pruning, or alternatively, to motivate better its settings. The paper readability can be improved with a better structure, more clarity in arguments and qualitative discussions, and more clear statements about the paper novel contributions. (minor) The source code shall be provided to ensure easy reproducibility.

Questions

I don't have questions which may change my evaluation. Still, I consider this study promising, and I hope that my comments will help the authors to improve the next version of the paper.

Rating

3

Confidence

3

Soundness

2

Presentation

1

Contribution

2

Limitations

n/a

Thank you for the rebuttal

Dear authors, Thank you for considering my comments and for preparing an extensive rebuttal. I still believe that all these (including your answers to the other reviewers) has to be integrated in a new version of the paper and that version will still need a new peer-review process. Therefore, I have to keep my original rating. Nevertheless, I take your point about the systematic misuse of dense hyperparameters in sparse training as novel indeed, and for this reason, I will not put myself against acceptance if the other reviewers and the area chair will unanimously support the acceptance of this paper. Best wishes, Reviewer F4Ay

Thank you for your rebuttal

Dear authors, Thank you for your rebuttal and additional experiments. It seems that the performance on DST is not that good. But most of my concerns are addressed. I increase my score to 5