SLTrain: a sparse plus low-rank approach for parameter and memory efficient pretraining

Summary

In this work, the authors introduce SLTrain, a novel method for pre-training large language models (LLMs) that combines sparse and low-rank matrix structures to enhance parameter and memory efficiency. The low-rank component is learned via matrix factorization, while the sparse component is achieved by uniformly selecting the sparsity support at random and learning only the non-zero entries with the fixed support. This method significantly reduces the memory requirement for LLM pre-training while achieving on-par performance.

Strengths

1. This paper introduce SLTrain, serving as the first method to pre-train LLMs with sparse + low rank matrices. 2. The authors provide extensive experimental results demonstrating the effectiveness of SLTrain across various model sizes, from 60M to 7B parameters. The comparison with state-of-the-art methods such as ReLoRA and GaLore is thorough and highlights SLTrain's advantages in memory reduction and parameter efficiency. 3. The authors provide a solid theoretical motivation for combining sparse and low-rank components. The empirical analysis of singular values and the distribution of residuals supports the feasibility and effectiveness of this approach.

Weaknesses

1. As the model size scales up, the perplexity score gap between SLTrain and Full-Rank increases, whereas GaLore maintains a more consistent performance at scale. Consequently, the reduction in parameters with SLTrain leads to suboptimal performance and potential scalability challenges. 2. From Table 3, we observe that the memory efficiency of SLTrain does not translate into faster training speeds. This is due to the inclusion of sparse components during training. Additionally, I am curious about its efficiency during inference. For inference, we have two options: using sparse and low-rank matrices, which should reduce memory usage but increase inference time, or using a larger dense weight matrix, similar to full-rank inference. Therefore, I am interested in the memory and time requirements for both inference modes. 3. I am curious about the potential of using structural sparse patterns, such as butterfly matrices [1], to enhance training performance and efficiency. These structural sparse matrices can be computed in a more efficient way and has better expressivity. Reference: [1] Dao, Tri, et al. "Pixelated butterfly: Simple and efficient sparse training for neural network models." arXiv preprint arXiv:2112.00029 (2021).

Questions

See Weaknesses.

Rating

5

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

See Weaknesses.

Summary

The submission proposes an approach to reduce memory and computational overhead in training large neural networks. It combines sparse training with low-rank adaptations to achieve efficient training without significant performance degradation. The paper includes an evaluation of SLTrain compared to full-rank, low-rank, ReLoRA, and GaLore models across multiple model sizes and tasks of the LLAMA model family.

Strengths

1. Ablation Study: The paper yields a good overview of the structure of Llama and analyzes the singular spectrum of several layer matrices, to yield a solid motivation for the sparse + low-rank trick 2. Comparative Analysis: The paper provides a detailed comparison with other methods such as full-rank, low-rank, ReLoRA, and GaLore for Llama 3. Implementation Details: The methodology is well-documented, with clear descriptions of how parameters and optimizer states are managed.

Weaknesses

1. Main concern in this work is in the paragraph starting in line 172. The authors aim for low-rank + sparse pretraining of LLMs to save memory. However, their proposed evaluation of a sparse plus low-rank layer is (AB + S)x, where + denotes a sparse matrix add. This implies that the resulting full matrix Y=AB + S is stored in memory, (if only temporarily) which defeats the purpose of a low-rank formulation. The authors mention “gpu-friendliness” as the reason for this choice, mistaking GPU throughput for real efficiency gain. - In that regard: How exactly is the “actual memory footprint” in paragraph (line 315) measured? Rigorous details are needed here to make the contribution credible. 2. The paper is very experimental without theoretical justification of the method (which itself is fine), thus an increased focus on implementation details and actual memory consumption of the implementation is expected. 3. Results are specific for Llama. The authors consider speficifally the Llama model family in this study. Given the computational effort to train LLMs, this is fine, but is should be clearly stated in the scope, which mentions general LLMs. Have the authors reasonable proof that their method, and mostly, the study of the singular spectrum extends to other LLMs?

Questions

See above

Rating

5

Confidence

4

Soundness

2

Presentation

2

Contribution

3

Limitations

yes

Summary

The authors propose SLTrain that performs a low-rank factorization of the weights as well as a sparse matrix of factors that represents which parameters to update. The authors show that their method can achieve significant memory savings compared to GaLore while retaining performance.

Strengths

- the paper is well written and easy to read - the method is very simple to integrate to the model as randomly generating sparse factors is easy to initialize - the savings look significant compared to GaLore

Weaknesses

- no strong theoretical justification why the sparse method works. - it is not clear how much sparsity is needed for a specific task. Is the degree of sparsity impactful on the performance? - what if we only parametrize the weights with BA and not BA+S, wouldn't that decrease the number of parameters? and how well will it perform? - wouldn't it be more useful if we regenerate the random sparse factors every couple of iterations to have a more effective pretraining? - how much do the results depend on the randomness of the sparse factors? can we have multiple runs and a standard deviation to see if that randomness play a big role in the results? - comparing GaLore and SLTrain is not really apple-to-apple as GaLore is based on the gradients and SLTrain is based on the weights. What if you had GaLore+SLTrain where gradients and weights are both made more efficient?

Questions

Please address the weaknesses above

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

yes

Response to Rebuttal

Thank you for your rebuttal. Regarding W1, I still observe a non-negligible performance gap between SLTrain and Full Rank pre-training. In Table 2 of the PDF, the perplexity is slightly higher for SLTrain compared to Full Rank, while the memory usage is slightly lower. Given this, I believe that Full Rank pre-training might still be preferred for its better performance. I would expect SLTrain to achieve a similar performance to be truly impressive. Thank you for your responses to W2 and W3. I will maintain my current score.