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

Thank you for acknowledging the strengths of our work and providing many constructive feedback.

1. (W1) 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.

In Table 2, the score gaps between SLTrain and Full-Rank are 0.09, 1.68, 0.62, and 0.58 for 60M, 130M, 350M, and 1B models, respectfully. The corresponding difference between GaLore and Full-Rank scores are 0.82, 1.00, 0.15, and 0.08, respectively. Hence, for both SLTrain and GaLore, the difference from Full-Rank model increases from 60M to 130M and subsequently decreases. This suggests that the general trend is similar for SLTrain and GaLore. Hence, we respectfully disagree with the reviewer's comment.

Notably, the gap between SLTrain and full-rank training can be further reduced by increasing the sparsity level from the outset or by continual pretraining with an additional sparse factor. This approach does not introduce significant memory overhead but enhances performance to levels comparable with full-rank training. Evidence for this is presented in Table 5 of the main paper, where increased rank/$\delta$ corresponds to improved performance.

To validate this claim on larger models, we conducted additional experiments, increasing the sparsity $\delta$ from 0.03 to 0.05 for training LLaMA models with 350M and 1B parameters. The results, detailed in Table 2 of the supplementary one-page PDF, demonstrate that increasing $\delta$ to 0.05 reduces the performance gap while maintaining memory efficiency relative to full-rank training. Furthermore, we highlight that larger models allow for a greater increase in $\delta$ due to the more significant memory gap.

2. (W2) Comparison of low-rank plus sparse and dense matrix during inference in terms of memory and time.

In Table 1 of the one-page PDF, we have included a comparison of SLTrain and full-rank in terms of inference memory and runtime on LLaMA 130M up to 7B (with the same configuration as in the main paper). We can explicitly observe the trade-off in terms of memory and computation. In particular, as the model size increases, the percentage of memory savings becomes more pronounced, while the computational cost increases is less obvious.

Lastly, we would like to highlight that the computational efficiency of SLTrain highly depends on the implementation as well as the associated hardware. Traditional GPUs are usually poorly-designed for unstructured matrix multiplication. Thus to appropriately leverage the GPU power, we first compute $W = BA +S$ and then multiply by the input with dense operations. It should, however, be noted that we only declare a temporary variable for storing $BA + S$ for each linear layer and when it is executed, it will be freed and be replaced by $BA + S$ in the subsequent linear layer. This only results in a slight increase in memory compared to low-rank parameterization. On the other hand, this is still noticeably smaller in memory compared to full-rank parameterization (such as GaLore) where weights of all layers simultaneously occupy memory. This is a strategy that balances memory as well as computation. However, we believe we can match the computational efficiency of the dense model by exploiting many recently introduced sparsity-friendly hardware, such as [1]. Such a hardware is known to match the computational efficiency even for unstructured sparsity.

[1] Thangarasa, V., Saxena, S., Gupta, A., and Lie, S. Sparse-IFT: Sparse Iso-FLOP Transformations for Maximizing Training Efficiency. In ICML 2024.

3. (W3) Potential of using structural sparse patterns.

Thank you for your suggestion. We also believe this is interesting to explore and have already discussed this possibility in the concluding remarks in the main paper.

We sincerely thank you for your feedback. We would like to take this opportunity to address your concerns and questions individually. We hope our clarifications would lead to your re-evaluation of the contribution of this work.

1. (W1) (1) $BA + S$ requires full matrix to be stored in memory (if temporarily) which defeats the purpose of low-rank formulation. (2) How exactly is the “actual memory footprint” in paragraph (line 315) measured?

(1) First, we would like to reiterate that the choice of computing $BA + S$ first before multiplying $x$ is due to sparse multiplication being poorly-supported in most GPUs, thus resulting in higher computational cost if computed separately. However, we only declare a temporary variable for storing $BA + S$ for each linear layer and when it is executed, it will be freed and be replaced by $BA + S$ in the subsequent linear layer. This only results in a slight increase in memory compared to low-rank parameterization. On the other hand, this is still noticeably smaller in memory compared to full-rank parameterization (such as GaLore) where weights of all layers simultaneously occupy memory.

To verify the above claims, we perform an additional experiment, comparing the actual maximum memory consumption for the proposed SLTrain linear layer (Algorithm 1, $(BA + S)x$) and standard full-rank linear layer ($W x$) and low-rank linear layer ($BA x$) in a feedforward neural network. We include the results for both forward and backward pass where we vary the number of layers in Figure 1 of the additional one-page supplementary PDF. Specifically, we set the input, hidden and output size to be $2048$ and $r = 128$ with $\delta = 0.03$. From the figure, we observe that as the number of layers increases, the reduction in memory of SLTrain becomes more evident compared to full-rank model. On the other hand, the memory overhead of SLTrain compared to low-rank model is only marginal. In terms of computational cost, we see compared to full-rank model, SLTrain only requires slight computational overhead, which is due to the scatter-adding operation.

Hence, our proposed $BA+S$ modeling and computation is memory efficient, which does not defeat the purpose of low-rank formulation.

(2) For the second question on "actual memory footprint", we measure the max GPU memory allocated (which is the function 'torch.cuda.max_memory_allocated' ) during pretraining.

2. (W2) (1) No theoretical justification. (2) Details on actual memory consumption of the implementation.

(1) We provide a formal justification as follows.

Theorem. Consider a matrix $S \in \mathbb R^{n \times n}$ with support $\mathcal{S}$ sampled uniformly at random with probability $\delta \in (0,1)$, i.e., $\mathbb P[(i,j) \in \mathcal{S}] = \delta$, for all $i, j \in [n]$. Suppose $\delta = \Omega(\log n/n)$, then with probability at least $1- \mathcal{O}(1/n)$, $BA + S$ is full rank for arbitrary randomly generated $B \in \mathbb R^{n \times r}, A \in \mathbb R^{r \times n}$ and for any $r \leq n$.

This claims that while $BA$ itself is low rank and has limited expressivity, augmenting $BA$ with a uniform-support sparse matrix renders it to become full-rank. We will include this theorem in our revised paper.

(2) As discussed in the previous point (W1), we measure the max GPU memory allocated (which is the function 'torch.cuda.max_memory_allocated') during pretraining. As part of supplementary material to the original submission, we also provide the codes that we use to perform those computations. We will include these details in the revised manuscript.

3. (W3) Does the study of singular spectrum extend to other LLMs?

Yes, we believe the observations on singular spectrum also extend to other LLMs. To validate our conjecture, we repeat the analysis of singular spectrum for pretrained Pythia 70M, downloaded from Hugging Face. Specifically, we set the rank $r = 128$ and extract the best rank $r$ approximation of the learned weight matrices. The results are shown in Figure 2 of the supplementary one-page PDF, where we observe that the residual after removing the best low-rank approximation vary smoothly and has small magnitude.

Thank you for the positive and constructive comments on our work.

1. (W1) Lack of theoretical justification.

We motivate the low-rank plus sparse modelling from the empirical observations (Figure 2 in the main text and Figure 2 in the one-page supplementary PDF) that the pretrained weights of LLMs can be well-approximated by a low-rank and a (uniformly) sparse factor. To further justify the modelling, we provide a formal justification as follows.

Theorem. Consider a matrix $S \in \mathbb R^{n \times n}$ with support $\mathcal{S}$ sampled uniformly at random with probability $\delta \in (0,1)$, i.e., $\mathbb P[(i,j) \in \mathcal{S}] = \delta$, for all $i, j \in [n]$. Suppose $\delta = \Omega(\log n/n)$, then with probability at least $1- \mathcal{O}(1/n)$, $BA + S$ is full rank for arbitrary randomly generated $B \in \mathbb R^{n \times r}, A \in \mathbb R^{r \times n}$ and for any $r \leq n$.

This suggests that although $BA$ itself is low-rank, which has limited expressivity, $BA + S$ is full-rank with high probability as long as the support is selected uniformly at random with sufficient non-zero entries. We will include this theorem in our revised draft.

2. (W2) Is the degree of sparsity impactful on the performance?

Yes, the degree of sparsity would impact the performance. We have already shown such results in Table 5 in the main paper where we vary the degree of sparsity ($\delta$). A trade-off exists between the performance and memory consumption, where more parameters (higher memory consumption) in general correspond to better performance.

3. (W3) 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?

The low-rank only approach ($BA$) does not work well in pretraining, e.g., already noted in [59]. In fact that is the motivation for our work (we mention this in Lines 64-65) and that full-rankness is required for effective pretraining. Indeed, $BA+S$ parameterization leads to full-rank weights (as our above theorem suggests) and consistently outperforms $BA$ in scores. See the comparison in Table 2 of main paper, where the baseline "Low-Rank" refers to $BA$ parameterization.

4. (W4) More useful when random sparse factors are regenerated.

We are not sure if regenerating random sparse factors would be useful. Our experiments show that learning of the sparse component $S$ with a fixed-random support is as useful as learning the low-rank part. If we were to regenerate the mask every couple of iterations, then it would mean that we are ignoring the factor learnt $S$ till now and only leveraging the learnt $BA$, which may lead to a decrease in performance. Furthermore, our experiments in Figure 4 suggest that the results remain almost the same with different uniform masks.

5. (W5) 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?

We have already done this. We have validated the influence of randomness for 60M and 130M model by running the models 5 times and presented the result in Figure 4 of the main paper that the randomness does not significantly affect the performance of the two models. In particular, the perplexity for 60M model at 1.1B is 33.91 with a standard deviation of 0.18 and the perplexity for 130M model at 2.2B is 26.01 with a standard deviation of 0.10.

6. (W6) 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?.

First, we highlight that both GaLore and SLTrain aim to do memory efficient pretraining but using different strategies. Therefore, it is not unfair to compare them.

On your suggestion about GaLore+SLTrain, indeed, we have already discussed this possibility explicitly in Line 217, given our proposed strategy is orthogonal to the development of GaLore. However, we believe that this is out of scope for the current submission.