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

Authors: andi  han
jiaxiang  li
wei  huang
mingyi  hong
akiko  takeda
pratik  jawanpuria
bamdev  mishra
riken  aip
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 14:18:21

Abstract: Large language models (LLMs) have shown impressive capabilities across various tasks. However, training LLMs from scratch requires significant computational power and extensive memory capacity. Recent studies have explored low-rank structures on weights for efficient fine-tuning in terms of parameters and memory, either through low-rank adaptation or factorization. While effective for fine-tuning, low-rank structures are generally less suitable for pretraining because they restrict parameters to a low-dimensional subspace. In this work, we propose to parameterize the weights as a sum of low-rank and sparse matrices for pretraining, which we call SLTrain. The low-rank component is learned via matrix factorization, while for the sparse component, we employ a simple strategy of uniformly selecting the sparsity support at random and learning only the non-zero entries with the fixed support. While being simple, the random fixed-support sparse learning strategy significantly enhances pretraining when combined with low-rank learning. Our results show that SLTrain adds minimal extra parameters and memory costs compared to pretraining with low-rank parameterization, yet achieves substantially better performance, which is comparable to full-rank training. Remarkably, when combined with quantization and per-layer updates, SLTrain can reduce memory requirements by up to 73% when pretraining the LLaMA 7B model.

Main paper: 
 introduction: 
 Large foundation models have achieved tremendous success in various domains, including linguistics, computer vision and biology. In particular, large language models (LLMs), such as the GPT series [39, 5] and the LLaMA family [51, 52] have reshaped the perception of how machine understands human languages. The predominant success of these models is primarily due to the model size, usually scaling to hundreds of billions of parameters. The scaling laws seem to suggest the capacity of LLMs grows with the model size [25] , but nonetheless requiring massive amount of resources for pre-training, storing, and fine-tuning. Particularly, memory requirement for training an LLM imposes a hard barrier for model deployment on commercial GPUs. For example, the LLaMA 7B model requires a minimum memory cost of approximately 42G under 16-bit floating point, including 14G of parameter state and 28G of optimizer state for momentum-based optimizers, like Adam [59, 28] .Building an LLM (from scratch) for downstream tasks typically involves two phases, i.e., pretraining and fine-tuning. The goal of pretraining is to capture general language patterns and semantics, enabling the model to acquire useful representations of words and sentences. Common pretraining objectives include masked language modeling [26] , next token prediction [39, 40] , etc.38th Conference on Neural Information Processing Systems (NeurIPS 2024). Fine-tuning then tailors the learned model representations from pretraining to downstream tasks, adjusting its weights to enhance performance on specific objectives. Pioneered by LoRA [21] , recent works have popularized low-rank finetuning of a given pretrained model (W 0 ), where W 0 is generally full-rank (i.e., pretrained without any constraints). The premise is that LLMs usually adapt to downstream tasks in a lowdimensional subspace, which allows to parameterize the update by low-rank factors. Low-rank finetuning requires minimal trainable parameters and significant reduction of memory and computation resources [10, 12] . A number of works [14, 53, 18, 29, 30, 34, 3] have emerged to further improve the efficiency and adaptation capacity of LoRA.While most of the works have focused on exploiting low-rank structure for fine-tuning, only a few [27, 24, 43, 47] have considered pretraining with low-rank weights. It has been observed that the performance of low-rank training often lags behind full-rank training despite the great potential for improving training and memory efficiency [47, 59] . This is because neural networks often exhibit full-rank structure in the weights and imposing low-rank restrictions could significantly limit their representation power. Hence, recent works have explored full-rank training with low-rank updates. For instance, ReLoRA [32] periodically restarts LoRA, where the low-rank updates are merged with the weights from the last period. However, ReLoRA also requires a warm-start full-rank training to achieve competitive performance [32] . GaLore [59] takes a different route by enforcing a low-rank structure not on the weights but on the gradients. This allows the Adam optimizer states to be stored in a low-dimensional space. While being memory efficient, GaLore is not parameter efficient because it still performs full parameter update with "projected-back" low-rank gradients.Parameter efficiency is a desirable property post-pretraining for model deployment, fine-tuning, and model storage. On the other hand, memory efficiency is necessary for training models with lower hardware requirements. Despite the importance of both parameter and memory efficiency, these two goals are often pursued independently. While low-rank models achieve both parameter and memory efficiency, as discussed earlier, they do not perform well in general [47, 59] . Therefore, a natural question is: how can we adapt low-rank training to achieve comparable performance as full-rank training while maintaining both parameter and memory efficiency?Contributions. In this work, we answer the above question by directly parameterizing the weights as low-rank plus sparse factors for pretraining. It should be noted that both low-rank and sparse factors individually facilitate parameter efficiency. Furthermore, their combination (usually) ensures that the final pretrained model is of high rank. Existing strategies for sparse learning usually involve prune-and-growth [13, 2, 58, 49] that iteratively train, prune, and grow neurons. Such a strategy is usually not memory efficient due to the need of storing (and learning) a support and a dense weight matrix. In contrast, we motivate and adopt a simpler strategy of fixing a uniformly random support (for the sparse factor). This allows to only store indices and values for memory efficient training, which scales with the number of nonzero entries. We show such a simple approach allows to further reduce the memory consumption during pretraining compared to ReLoRA [32] and GaLore [59] without sacrificing performance. We show this through an extensive set of experiments on the LLaMA language models with varying model size from 60M up to 7B parameters. We call our proposed sparse plus low-rank pretraining algorithm as SLTrain. In Figure 1 , we observe that SLTrain obtains perplexity score comparable to full-rank model with considerable memory and parameter efficiency.We end this section by noting that the idea of marrying low-rank and sparse factors has been explored for robust matrix recovery [6, 57, 4] , attention matrix approximation [7] , and neural network compression [31] . However, it is introduced for pretraining LLMs for the first time in our work. 
background on low-rank pretraining: 
 Existing pretraining works [24, 43] have explored low-rank parameterization of the layer weights directly as W = BA. However, it has been empirically observed that vanilla low-rank parameterization suffers from large performance degradation because of the limited representation capacity [47, 32, 59 ]. Hence, motivated from low-rank adaptation (LoRA) [21] for fine-tuning, for pretraining, ReLoRA [32] suggests to parameterize the layer weights asEQUATIONwhere m represents the number of low-rank factors. This parameterization results in an overall high-rank update compared to LoRA because the sum of low-rank matrices is generally a higher rank matrix. The optimization is performed by training B s , A s iteratively, merging W s ← W s-1 + B s A s , and then restarting the optimization for B s+1 , A s+1 . A key drawback of ReLoRA is that it stores the full-rank matrix W s throughout the training and inference stages. Hence, it is memory intensive and not parameter efficient. While ReLoRA performs sequential low-rank updates in (1), a recent work [22] has explored parallel low-rank updates and merging them for pretraining.A more recent work, GaLore [59] , imposes low-rank structure on the gradient. Specifically, GaLore still optimizes full-rank weights and computes full-rank gradients G t at iteration t, but updates Adam moments M t , V t in a low-dimensional space, i.e.,M t ← β 1 M t-1 + (1 -β 1 )P ⊤ t G t , V t ← β 2 V t-1 + (1 -β 2 )(P ⊤ t G t ) 2 W t+1 ← W t -η P t M t /( V t + ϵ),where P t is a projection matrix constructed by taking the largest left singular vectors of G t . To reduce computational cost, P t is computed every several iterations and is stored in the middle. Although being memory efficient (as M t and V t are computed in the smaller dimension), GaLore is not parameter efficient due to computation of P t M t for updating W t . 
sltrain: proposed sparse plus low-rank pretraining: 
 In order to achieve both parameter and memory efficiency, we propose to adapt low-rank parameterization by introducing a sparse factor. We model the weight matrices as a sum of sparse and low-rank matrices. Our proposed modeling is referred to as SLTrain. Below, we discuss the motivation, modeling details, and practical considerations for implementing SLTrain.Both low-rank and sparsity are parsimonious modeling strategies for exploring low-dimensional weight matrices. The low-rank component aims to learn the low-dimensional bases or eigenspaces of the weights. The sparse component, on the other hand, identifies effective neuron-wise interactions and disregards non-expressive ones. In linear algebra terms, the low-rank component enforces sparsity of singular values, whereas the sparse component enforces sparsity of individual entries. In general, low-rank matrices are not sparse, and sparse matrices are not necessarily low-rank [6] . These concepts provide complementary information that should be explored simultaneously.Despite that low-rank modeling alone can have limited expressivity due to the low-rank structure it imposes, we show in the below proposition that low-rank plus a uniform sparse matrix with only Ω(log n/n) number of entries is full-rank with high probability. Proposition 1. Consider a matrix S ∈ R n×n with support S sampled uniformly at random with probability δ ∈ (0, 1), i.e., P[(i, j) ∈ S] = δ, for all i, j ∈ [n]. Suppose δ = Ω(log n/n), then with probability at least 1 -O(1/n), BA + S is full rank for arbitrary randomly generated B, A.To further motivate the sparse plus low-rank modeling, in Figure 2 , we illustrate different statistics from weight matrices of a pretrained LLaMA 60M model on C4 dataset (introduced later in Section 5). In Figure 2 (a), we plot the singular values of weight matrices of different layers. The plot exhibits a fast decay of the singular values followed by a more stable decay of (smaller) singular values. This suggests that the top subspaces can be effective in model compression, and therefore, builds a case for low-rank modeling. However, the tail singular value distribution shows that low-rank modeling purely may not be sufficient. In order to better understand the tail part, in Figure 2 (b), we visualize the magnitude of the attention output weight matrix before and after we extract the top r-dimensional subspaces (r = 128) for the last attention layer. It is apparent that, after removing the top subspaces, both the magnitudes and the variation of the entries present in the residual matrix become smaller.Plotting the magnitudes of the entries in Figure 2 (c) we see that 97% of the entries have a magnitude below 0.04. In Appendix B and C, we provide such visualizations for other layers of LLaMA 60M and Pythia 70M to further corroborate the findings. Overall, the figures suggest that a sparse matrix with random support can approximate the residual well given the magnitudes do not vary too much across the entries.In Table 1 , we perform an ablation study that verifies the feasibility of using a random sparse support for approximating the residual matrix. Specifically, we take L 0 as the best rank-r approximation (r = 128) for the pretrained weight matrix W 0 and evaluate the perplexity score (PPL) on the validation set. We see that compared to the full-rank pretrained model, low-rank approximation L 0 suffers from a drastic performance drop. We also augment the low-rank approximation L 0 with either top 3% or random 3% of entries of the residual matrix, which we label as top sparse or random sparse pruning, respectively. We observe that L 0 plus top sparse pruning performs better compared to L 0 plus random sparse pruning. Nonetheless, both the performance is poor. We further evaluate fixing the low-rank approximation (to L 0 ) and only optimizing the sparse components with either top support or random support (both run for five times). Averaged PPL (over the five runs) for both the approaches improve and are comparable. This shows that fixing a random support for the sparse factor is a promising strategy from both efficiency and performance point of view. We explore learning both the sparse and lowrank factors in the next section.Building on Section 3.1, we propose to parameterize weight matrices W ∈ R d×p asW = BA + S,where B ∈ R d×r , A ∈ R r×p are low-rank factors with r < min{d, p} being the rank parameter and S ∈ R m×n is a sparse matrix. The number of non-zero entries (nnz) in S is determined by the sparsity level parameter δ ∈ (0, 1), i.e., nnz(S) = δdp. So, the total number of parameters for the proposed parameterization is (d + p)r + δdp, which is much smaller than the full-rank layer parameters dp when we choose δ ≪ 1. In addition to being parameter efficient, the optimization states also cost less memory and scales with the number of trainable parameters. Finally, we note that the overall rank of W will generally be high due to the presence of the sparse factor S, based on Proposition 1.The performance of such a parameterization highly depends on whether there exists an implementation that is both computation and memory efficient. Nevertheless, modern GPU hardware is not suited for sparse tensor multiplication Sx for given input x, as well as its gradient, especially when S presents an unstructured sparsity pattern [7] . This causes increased computational bottleneck despite showing memory advantage. Thus, existing works on sparse network and training mostly rely on learning and storing a parameter mask (i.e., support) [48, 15, 33] by letting S = M ⊙ U , where M ∈ {0, 1} d×p is a binary mask and U ∈ R d×p is a dense parameter. This allows to exploit GPU accelerator for dense matrix computation. However, masking requires to store both the support and a dense parameter for training, which significantly increases the memory cost.In this work, we achieve memory efficiency by representing S in terms of its indices and values, i.e., (I, V) ∈ R 2nnz(S) . This is possible because we randomly (and uniformly) fix the support a priori. The motivation for using a random (but fixed) support comes from the usefulness of random support in Table 1 . This ensures the memory scales with only the sparsity in S (i.e., the support size) rather than the full size of S. Further, the forward pass involves computingBAx + Sx = (BA ⊕ I V)x,where we denote W ⊕ I V as scatter-adding V to W at the indices specified in I. Because this operation results in a dense matrix, sparse matrix multiplication is avoided. Hence, this is GPU friendly without requiring to store a binary mask.We remark that despite we require to compute a dense matrix, which has the same size as the full-rank matrix, we never store it for backpropagation. In particular, we can compute the gradient with respect to B, A, V, and input x asEQUATIONwhere we let z = (BA ⊕ I V)x and L denotes the loss function. We also denote W I as gathering the values of W at indices I. In other words, we only need to store B, A, I, V for backpropagation. This is illustrated in Algorithm 1 where we define a customized linear layer in SLTrain. We highlight that such a parameterization is agnostic to the chosen optimizers and can easily be integrated with any optimizer including Adam.In comparison with the recent pretraining works based on low-rank factors/gradients, SLTrain is more parameter and memory efficient than ReLoRA [32] and GaLore [59] as it only optimizes the low-rank and sparse factors without the need for storing full-rank matrices.Initialization and scaling. We consider LoRA type of initialization for low-rank factors, i.e., Kaiming initialization [19] for A factor and zero initialization for B factor.Algorithm 1 SLTrain for linear layer 1: Input: x, B t , A t , (I, V t ). 2: def forward(x, B t , A t , I, V t ): x, B t , A t , I, V t ← saved_tensor 8:Compute gradient as in (2). 9:return ∇ x L, ∇ Bt L, ∇ At L, None, ∇ Vt LFor sparse factor, we adopt uniform initialization for the values V in the range of[-1/ √ d in , 1/ √ d in ],where d in denotes input feature dimension. We choose the sparse support I uniformly at random up to the desired sparsity level δ. Furthermore, in order to balance the contribution of low-rank factor and sparse factor, we follow LoRA [21] to scale the lowrank factors by α/r where the balancing parameter α is a hyperparameter. This hyperparameter, along with the stepsize has a joint effect on the training speed of low-rank versus sparse factors.Regularization and preconditioning. It is expected that the optimization of low-rank factors can cause instability when using larger stepsize or larger balancing parameter α, an issue already present in low-rank training [43] . This is primarily due to the multiplicative updates of B, A simultaneously. Existing solutions, such as orthogonal constraints or regularization [43] , preconditioning [50, 23, 56] , can be easily combined with the proposed modelling for more stable convergence.Integration with other techniques. Since the proposed sparse plus low-rank approach pursues memory saving from the perspective of reparameterization, SLTrain can be easily integrated with optimizer-based techniques for further improving memory efficiency, including quantization [9] (that uses lower-bits for storing moment states without sacrificing the performance), per-layer weight updates [36] (that updates the parameters along with backpropagation), and activation checkpointing (that recomputes the activation states instead of storing them). In addition, SLTrain can be even combined with low-rank gradients in GaLore [59] for low-rank factors. This can further reduce the memory footprint, especially for larger models where the rank r is set to be high. On the other hand, because we use a simple strategy of fixed-support sparse learning, it may be beneficial to combine with different techniques for dynamic support learning [2, 20] . 
related works: 
 Low-rank fine-tuning and training. Building on the idea of LoRA [21] that parameterizes the update as low-rank factors, i.e., ∆W = BA, ROSA [14] dynamically adapts subspaces for training, where the subspaces are selected by taking SVD of the current weight matrices. Chain of Lora [53] decomposes the low-rank update into a sequence of small-size matrix product, ∆W = k j=1 B j A j . NOLA [29] parameterizes the two small matrices as linear combination of two sets of random basis respectively, B = m i=1 α i B i , A = n j=1 β j A j where A i , B j are fixed random matrices. NOLA optimizes over the coefficients, thus further improving parameter efficiency. VeRA [30] considers a similar parameterization as NOLA where B = diag(b) B, A = diag(a) A for fixed random matrices B, A. DoRA [34] decomposes pre-trained weights into magnitude and directional components and separately fine-tune with LoRA adopted for directional update. SoRA [11] introduces a dynamic rank adaptation strategy for tuning LoRA rank. ResLoRA [46] adds a residual path for LoRA adaptors. For pretraining, in addition to ReLoRA [32] and GaLore [59] , Flora [17] demonstrates LoRA updates approximate random projection of gradient, and by resampling the random projection, high-rank training can be achieved. LTE [22] adopts the similar idea of parameterizing a high-rank matrix through summation of low-rank matrices and adopts the parallel train-and-merge strategy as opposed to sequential in ReLoRA [32] .Sparse fine-tuning, training and sparse networks. Sparse fine-tuning/training aims to selectively update the weights with others fixed [48, 1, 2, 49, 15, 33] . This usually entails choosing a proper subset of parameters either randomly [49] , or based on approximate Fisher information [1] , magnitude of the change [48] , gradients and momenta [2] , as well as by learning a parameter mask (for storing support) with sparsity regularization [15] . On the other hand, sparse networks, also known as model pruning, directly search for a minimal architecture [16, 35] by removing redundant weights. We refer to the survey [20] for complete discussions of sparse network pruning.Sparse plus low-rank. Decomposing a matrix into the sum of low-rank and sparse matrix is a classic optimization problem for matrix recovery [6, 54, 4] . Recently, some works also consider harnessing both low-rank structure and sparsity for neural network compression. Scatterbrain [7] considers approximating the attention matrix for faster inference with sparse plus low-rank factors. More specifically, given Q, K, V ∈ R n×d The main aim is to efficiently approximate exp(QK ⊤ )V , which suffers from quadratic complexity in sequence length n. Hence, [7] proposes to leverage a random feature map ϕ :R d - → R m , defined as ϕ(x) = 1 √ m exp(W x -∥x∥ 2 /2) with entries of W sampled from Gaussian distribution N (0, 1), which defines a low-rank approximation ϕ(Q)ϕ(K) ⊤ ≈ exp(QK ⊤ ). Then a sparse matrix is constructed based on locality sensitivity hashing with the non-zero entries S i,j = exp(QK ⊤ ) i,j -ϕ(Q) ⊤ i ϕ(K) j . However, the aim of [7] is to approximate the attention matrix to reduce the computational cost while we aim to achieve memory efficiency by directly parameterizing the weight matrix. More specifically, in the context of selfattention where Q = XW Q , K = XW K , V = XW V , we directly parameterize each projection matrix W Q , W K , W V as low-rank plus sparse factors, e.g., W Q = BA + S. In addition, LoSparse [31] proposes to decompose the pretrained weights into low-rank plus sparse factors for structured Table 2 : Validation perplexity (PPL(↓)), number of parameters in millions (Param), and estimated total memory cost in G (Mem). The perplexity results for all the baselines are taken from [59] . For SLTrain, we use the same rank as other baselines and fix δ = 0.03. compression. Nevertheless, they consider optimizing the sparse matrix via iterative thresholding, which requires to store the full-size sparse matrix. We instead consider directly optimizing the sparse matrix on its non-zero entries for memory-efficient pretraining.Memory efficient training. To overcome the memory limitation of LLMs, many techniques have been proposed, such as reduced-precision, quantization [9, 10] , gradient checkpointing [42] and gradient accumulation [8] , and row-sum/column-sum of second-order statistics in Adafactor [45] , among many others. As has already been discussed, the proposed sparse plus low-rank parameterization is orthogonal to these developments where the techniques can be easily integrated for further memory reduction. 
experiments: 
 This section validates the effectiveness of low-rank plus sparse structure for pretraining and finetuning large language models. All the experiments are run on NVIDIA A100 GPUs. The code is available on https://github.com/andyjm3/SLTrain.Following [32, 59] , we consider pretraining the LLaMA language models [51] with pre-normalization, RMSnorm [55] , and SwiGLU activation [44] . We train LLaMA LLMs on C4 (Colossal Clean Crawled Corpus) dataset [41] , which is specially designed for pretraining. The training is performed without data repetition and we consider LLaMA with varying model sizes from 60M up to 7B parameters.Baselines. We compare our SLTrain with the baselines which exploit low-rank structures.• Full-Rank: This is the vanilla baseline that pretrains with full-rank weights using the Adam optimizer. • Low-Rank [24] : This is the low-rank parameterization of weights by factorizing W = BA where optimization is on B, A. • ReLoRA [32] : ReLoRA periodically restarts LoRA [21] by merging the learned low-rank adaptors with layer weights and reinitializing the optimizer state and learning rate. • GaLore [59] : GaLore explores low-rank structure for the gradients rather than for the parameters.We implement SLTrain with Adam by reparameterizing the weights from all linear layers, including fully-connected layers as well as query, key, value projection layers. The remaining parameters are updated with full-rank parameterization. This is consistent with the setup used in [21, 32, 59] .Hyperparameters. For SLTrain, we fix the rank r to be the same as the baselines and fix sparsity ratio δ = 0.03 across all the model sizes except for LLaMA 7B where we choose δ = 0.05, which achieves a good balance between efficiency and performance. We tune and fix the stepsize to be 0.003 and tune α in the range of [8, 16, 32] for the LLaMA 60M, 130M, 250M, 1B models. We fix the other parameters to their default settings. In particular, we choose α = 32 for the LLaMA 60M model and α = 16 for the 130M and 350M models and α = 8 for 1B. For the LLaMA 7B model, we choose stepsize to be 0.0005 and α = 8. Except for the 7B model, we directly inherit the perplexity results from [59] and thus do not need to tune the hyperparameters from the baselines. We ensure the comparison is fair based on the training token number.Memory cost estimation. We compare the proposed SLTrain with the low-rank baseline models in terms of estimated memory consumption. Following [59] , we compute memory estimates with bfloat16 format, where each floating point number occupies 2 bytes. We remark that SLTrain stores the indices with int64 format, which occupies 8 bytes per digit. The memory cost for a training algorithm consists of the parameter memory and optimizer state memory. The parameter memory refers to the memory occupied by storing parameters, and the optimizer state memory refers to the memory required to store the first and second-order moment statistics, e.g., in Adam. Table 2 reports the total estimated memory cost for each method. The detailed breakdown of memory estimation can be found in Appendix F.Perplexity vs efficiency. In Table 2 , we compare the performance of different methods in three aspects: perplexity score, parameter size, and memory cost. We observe that SLTrain performs comparatively as the full-rank training and GaLore [59] while achieving further reduction in parameter size and memory cost. In addition, SLTrain only adds a small parameter and memory overhead to the low-rank parameterization, yet significantly improves the perplexity score. Hence, learning the additional sparse factor indeed helps in strengthening the representation capacity of SLTrain. This intuition is also validated in Figure 10 (Appendix D ), where we plot the singular value distribution for different weight matrices. Due to the presence of the sparse factor, we observe that the spectrum of the SLTrain weight matrices gets enhanced beyond r = 128. Measuring actual memory footprint. In Figure 3 , we record the actual memory footprint of different methods across various model sizes, on a single A100 80G GPU. We measure the memory of 8-bit SLTrain with per-layer weight update using a single batch size and bfloat16 data type. Gradient checkpointing is not used for any method. The baselines include Adam trained on full-rank weights, 8-bit Adam, and 8-bit GaLore with per-layer weight. From the figure, we see that SLTrain achieves memory reduction by 51%, 58%, 73% for 350M, 1B, 7B models, respectively. Notably, compared to state-ofthe-art memory-efficient method GaLore [59] , SLTrain reduces the memory requirement by 29%, 34%, 17% when pretraining 350M, 1B, 7B models, respectively.Measuring throughput. We measure the throughput of SLTrain for pretraining on the LLaMA 350M and LLaMA 1B models with a token batch size of 256 on 1×80G A100 GPU and 4×80G A100 GPUs, respectively. The throughput is averaged over 5000 training steps. We observe in Table 3 that the throughput token of SLTrain is slightly lower than the full-rank and GaLore baselines. This is mainly due to the retrieving and setting for the sparse entries. We believe more efficient implementation can be developed in this regard.Scaling to LLaMA 7B model. For pretraining the LlaMA 7B model, due to the resource constraints, we only compare SLTrain with GaLore, implemented with 8-bit Adam [9] on 4×80G A100 GPUs, without per-layer weight updates nor gradient checkpointing. 1 . We directly use the training scripts of GaLore. 2 In Table 4 , we observe that SLTrain performs comparably to GaLore in terms of perplexity and throughput while achieving significant memory reduction of 26% per GPU device.Inference memory and throughput. In Table 5 , we compare the inference memory usage and throughput between SLTrain and the full-rank model across various model sizes, ranging from LLaMA 130M to 7B. A clear trade-off between memory and computational cost is observed. Specifically, as the model size increases, the percentage of memory savings becomes more pronounced, while the corresponding increase in computational cost is less significant. This underscores the growing advantage of SLTrain when employing larger models. Varying random sparse support. In contrast to common pruning strategies that require careful learning of sparse support, we show that SLTrain works with randomly selected support. In this regard, we perform pretraining on Llama 60M and 130M with five different randomly selected support for the sparse factor S. The convergence plots are shown in Figure 4 , where we see that changing the random support for the sparse factor does not materially affect the performance.How do rank r and sparsity δ affect performance? In Table 6 , we validate the performance of training the Llama 60M and 130M models by varying the hyperparameters r and δ. We notice that, in general, when more parameters are added (corresponding to higher r and δ), the performance is better. However, this performance gain is also accompanied by an increase in memory footprint. In our initial experiments, we also tried the extreme setting with r = 0, i.e., only the S component was learned but with a higher δ value. This setting demonstrated a reasonably good performance. Further investigations in this direction are left for future work.Further comparisons to full-rank performance. In this section, we evaluate the potential of SLTrain to achieve performance comparable to full-rank models by adjusting the sparsity ratio, δ. As shown in Table 7 , we increase δ from 0.03 to 0.1 for the LLaMA 350M and 1B models. Our results indicate that setting δ = 0.1 enables SLTrain to attain perplexity scores similar to those of full-rank models, while maintaining memory and parameter efficiency. Notably, at δ = 0.1, SLTrain reduces the parameter size by 42% for the 350M model and 45% for the 1B model. These results highlight the effectiveness of SLTrain in significantly reducing model size without sacrificing performance.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: The code to reproduce the results is included as part of supplementary material. The datasets are publicly available.• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
concluding remarks: 
 In this paper, we propose SLTrain for achieving both memory and parameter efficient pretraining of LLMs. SLTrain combines two complementary parsimonious structures, low-rank and sparsity, for learning models with high representation capacity. While low rank is modeled via the product BA of r-dimensional matrices, the support of the sparse factor S is obtained by uniformly random sampling over indices. The matrices B, A, S are learned for different layers via backpropagation. Empirically, we achieve state-of-the-art memory reduction during pretraining while maintaining competitive performance. Although we show results on the LLaMA language models (with varying size from 60M to 7B parameters), we believe SLTrain could also help to improve memory efficiency for vision foundation models and multi-modal foundation models, such as diffusion models [37] and CLIP [38] . The significant improvements shown by SLTrain also motivates future works to understand the theoretical guarantees of training with both low-rank and sparse factors, such as convergence and loss landscape. We hope this work initiates exploration on combination of other parsimonious structures for pretraining such as Kronecker product or structured sparsity (e.g., block-diagonal, group-lasso). Proof of Proposition 1. For each row, the probability that at least one entry is non-zero is 1-(1-p) n due to the independence of the non-zero entry. Due to the independence of uniform selection, the probability that all the rows have at least one entry is (1 -(1 -p) n ) n . By choosing p = 2 log n/n, the probability can be simplified to (1 -e -2 log n ) n = 1 -O(1/n). Similarly, we can apply the same argument for the columns. Taking the union bound shows that with probability at least 1 -O(1/n), S has at least one entry for each row and each column respectively. Further, because the non-zero entries of S are sampled from continuous distribution, the set of matrices such that S becomes low-rank has measure zero. Then taking union bound gives with probability at least 1 -O(1/n), S is full rank.The remaining part is to prove that BA + S is full rank. Because B, A are sampled from a continuous distribution space, the set of such matrices that augments S to form a degenerate matrix has measure zero. Thus, the proof is complete.In the main text, we present singular values and visualizations of the attention output weight matrix for a single layer, both before and after applying the low-rank approximation. Here, we extend this analysis by providing visualizations for additional weight matrices across multiple layers for the LLaMA 60M and 130M models. The figures reveal a consistent pattern, showing a low-rank plus (uniformly) small magnitude structure in the weight matrices of other layers as well. Here 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 9 , where we observe that the residual after removing the best low-rank approximation vary smoothly and has small magnitude. In the head distribution, SLTrain tries to approximate Low-Rank.In Figure 10 , we plot the distribution of singular values of pretrained weights of the LLaMA 60M model on 1.1B tokens. It is interesting to observe the distinct role of both the low-rank (L) and sparse (S) components. In particular, we see a cliff at index 128. This is because the singular values from 1 to 128 is due the learning of low-rank factor and from 129 to 512 is primarily due to the sparse component.To further evaluate the contributions of spectrum from low-rank and sparse components, we decompose the singular values of SLTrain weight matrices into contributions from the low-rank and sparse parts. Specifically, let U ΣV ⊤ = BA + S be the SVD, then we plot diag(Σ), diag(U ⊤ BAV ), diag(U ⊤ SV ), which we respectively call singular values, low-rank singular values, and sparse singular values. In Figure 11 , we see that the low-rank and sparse components contribute to different ends of the spectrum of the weight matrices. This justifies our modeling approach, i.e., by adding a Here we perform an additional experiment by 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 (BAx) in a feedforward neural network. We include the results for both forward and backward pass where we vary the number of layers in Figure 112 . Specifically, we set the input, hidden and output size to be 2048 and r = 128 with δ = 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. Following [59] , we compute memory estimate with bfloat16 format, where each floating point number occupies 2 bytes. For simplicity of estimation, we use the convention that 1G contains 10 9 bytes. Hence the estimation could be different from the estimation in [59] . We summarize the parameter and optimizer state breakdown in Table 8 .For a 60M model, we estimate the memory for each model as follows.• Full-rank: Full rank model has 58.2M parameters, which costs 0.12G and the optimizer state requires to store double size of the parameters, which is 0.23G. • Low-rank: A low-rank parameterization (with r = 128) for selected modules lead to 42.78M parameters (32.78M non-adapted parameters plus 10M low-rank adapted parameters), which consumes a memory of 0.08G. The optimizer state costs double size of memory, which is 0.16G. • ReLoRA: Similar to LoRA, ReLoRA requires to store both the original parameters as well as the adaptors (including both the low-rank adaptors and adaptors for other parameters), which in total has 102.77M parameters with an estimated memory cost of 0.20G. For the optimizer state, ReLoRA stores the moments of only trainable parameters, which is 85.54M with a memory of 0.17G. • GaLore: GaLore requires to store the full set of parameters 58.2M, which is the same as full-rank.For the optimizer state, GaLore stores the moments of projected gradients, which has a size of 78.20M, plus the projection matrix of size 3.67M. In total optimizer state requires to store 0.16G of memory. • SLTrain: For proposed SLTrain (with r = 128 and δ = 0.03), the parameter includes 32.78M base parameters, together with 10M low-rank factors and 0.76M sparse factors, which occupies 0.09G memory (where the sparse factors include 0.76M values in bfloat16 format and 0.76M indices in int64 format). For the optimizer state, the memory cost is double of the parameter size, which occupies 0.17G.For a 130M model, we estimate the memory for each model as follows.• Full-rank: Full rank model has 134.11M parameters, which costs 0.27G and the optimizer state requires to store double size of the parameters, which is 0.54G. • Low-rank: A low-rank parameterization (with r = 256) for selected modules lead to 94.00M parameters (49.17M non-adapted parameters plus 44.83M low-rank adapted parameters), which consumes a memory of 0.19G. The optimizer state costs double size of memory, which is 0.38G. • ReLoRA: Similar to LoRA, ReLoRA requires to store both the original parameters as well as the adaptors (including both the low-rank adaptors and adaptors for other parameters), which in For a 350M model, we estimate the memory for each model as follows.• Full-rank: Full rank model has 367.97M parameters, which costs 0.74G and the optimizer state requires to store double size of the parameters, which is 1.47G. • Low-rank: A low-rank parameterization (with r = 256) for selected modules lead to 185.22M parameters (65.59M non-adapted parameters plus 119.63M low-rank adapted parameters), which consumes a memory of 0.37G. The optimizer state costs double size of memory, which is 0.74G. • ReLoRA: Similar to LoRA, ReLoRA requires to store both the original parameters as well as the adaptors (including both the low-rank adaptors and adaptors for other parameters), which in total has 553.19M parameters with an estimated memory cost of 1.11G. For the optimizer state, ReLoRA stores the moments of only trainable parameters, which is 370.44M with a memory of 0.74G. • GaLore: GaLore requires to store the full set of parameters 367.97M, which is the same as full-rank. For the optimizer state, GaLore stores the moments of projected gradients, which has a size of 282.36M, plus the projection matrix of size 144.04M. In total optimizer state requires to store 0.34G of memory. • SLTrain: For proposed SLTrain (with r = 256 and δ = 0.03), the parameter includes 65.59M base parameters, together with 119.64M low-rank factors and 9.07M sparse factors, which occupies 0.46G memory (where the sparse factors include 9.07M values in bfloat16 format and 9.07M indices in int64 format). For the optimizer state, the memory cost is double of the parameter size, which is 0.78G.For a 1B model, we estimate the memory for each model as follows.• Full-rank: Full rank model has 1339.08M parameters, which costs 2.68G and the optimizer state requires to store double size of the parameters, which is 5.36G. • Low-rank: A low-rank parameterization (with r = 512) for selected modules lead to 609.31M parameters (131.17M non-adapted parameters plus 478.14M low-rank adapted parameters), which consumes a memory of 1.22G. The optimizer state costs double size of memory, which is 2.44G. • ReLoRA: Similar to LoRA, ReLoRA requires to store both the original parameters as well as the adaptors (including both the low-rank adaptors and adaptors for other parameters), which in total has 1948.39M parameters with an estimated memory cost of 3.90G. For the optimizer state, ReLoRA stores the moments of only trainable parameters, which is 1218.62M with a memory of 2.44G.• GaLore: GaLore requires to store the full set of parameters 1339.08M, which is the same as full-rank. For the optimizer state, GaLore stores the moments of projected gradients, which has a size of 866.30M, plus the projection matrix of size 176.16M. In total optimizer state requires to store 2.08G of memory.• SLTrain: For proposed SLTrain (with r = 512 and δ = 0.03), the parameter includes 131.17M base parameters, together with 478.14M low-rank factors and 36.24M sparse factors, which occupies 1.58G memory (where the sparse factors include 36.24M values in bfloat16 format and 36.24M indices in int64 format). For the optimizer the memory cost is double of the parameter size, which is 2.58G.In addition, we estimate the memory for Table 6 by listing out the parameter information for each parameter setting. In addition to pretraining, the sparse plus low-rank factor may also be used for fine-tuning LLMs as: W = W 0 + BA + S, where W 0 is a given pretrained model weight, B and A are low-rank factors, and S is the sparse factor. As in Section 3.2, we learn the fine-tuned weights B, A, and S and term this as SLTrain-FT. The fine-tuning experiments are conducted for RoBERTa base model (with 125M parameters) on GLUE benchmarks. We use r = 8 for all methods and tune the hyperparameters of SLTrain-FT as follows. We tune three hyperparmaters of SLTrain, i.e., sparsity level δ, balancing parameter α and stepsize η. Specifically, we tune δ in the range of [0.005, 0.001] and α in [32, 64, 128] , η in [1e-5, 2e-5, 3e-5, 4e-5, 5e-5]. The tuned hyperparameters are in Table 11 .The results on benchmark datasets are shown in Table 12 . We observe that SLTrain-FT performs competitively as the baselines. While there is no specific memory advantage of SLTrain-FT when a general full-rank W 0 is given, memory reductions can be obtained when W 0 is learned via SLTrain. We provide the source code for reproducing the experimental results reported in the paper. We also summarize some specific configurations that enhance reproducibility.• Datasets. The C4 dataset is publicly available on Huggingface and can be loaded using datasets package https://github.com/huggingface/datasets. The paper proposes a simple strategy that achieves both parameter and memory efficiency of training LLMs. We believe this work would produce positive environmental impact by reducing the energy consumption as well as carbon footprint during pretraining large foundation models.NeurIPS Paper Checklist Justification: We discuss the potential training instability in Section 3.3. We discuss the throughput overhead in Section 5.1. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 3. Theory Assumptions and Proofs some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: The hyperparameters (along with a given range of values that we use to tune), types of optimizers are included in the experiment section, i.e., Section 5.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: We report error bars (in terms of standard deviation) in Figure 4 for pretraining LlaMA 60M, 130M. For larger models, it is time-costly to report the error bars.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: We include a section on broader impact in Appendix I.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not release pretrained language models, image generators, etc. that have high risk of misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: We have properly cited the code and datasets used in the paper in experiment section.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: This paper does not introduce new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: Such experiments are not relevant to this paper.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: This paper does not involve human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.For full-rank model, 8-bit Adam throws out-of-memory error on 4×80G A100 GPUs.The scripts are available at https://github.com/jiaweizzhao/GaLore 
 
 Appendix:  


Decision: There was a disagreement among reviewers on this paper. However, the author pointed out a significant mis-understanding of the dissenting reviewer. Please clarify this issue in the next version.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: Please address the weaknesses above
limitations: limitations: yes
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See above
limitations: limitations: yes
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See Weaknesses.
limitations: limitations: See Weaknesses.

Comments: title: Response to further questions
Response to further comments
Response to Rebuttal
comment: Dear Reviewer AVXH,

We haven't heard back from you after our last response. We would be grateful if you could look at our responses. We believe our responses would have answered all the questions. Hope to see a positive rerating of the submission.

Best,
Authors
Dear Reviewer Rr3S,

We haven't heard back from you after our last response. We would be grateful if you could look at our responses. We believe our responses would have answered all the questions. Hope to see a positive rerating of the submission.

Best,
Authors
Thank you for the rebuttal and for addressing most of my concerns. I have increased my score by 1.
We sincerely thank the reviewers for their constructive feedback, which has significantly strengthened our paper. We believe we have effectively addressed all concerns and questions raised thus far. Specifically:


- We have provided a theoretical justification for the low-rank plus sparse modeling in response to the concerns raised by Reviewer Wf5L and Reviewer Rr3S.


- We have offered detailed explanations and numerical comparisons to clarify how $BA + S$ achieves memory efficiency, addressing the primary concern of Reviewer Rr3S.


- We have conducted additional experiments demonstrating that SLTrain can achieve comparable performance to full-rank models while significantly reducing parameter size and maintaining memory efficiency, addressing Reviewer AVXH’s concerns regarding the performance gap between SLTrain and full-rank models.


- We have emphasized the parameter efficiency of the proposed model, alongside its memory efficiency, one of the key contributions that tend be overlooked by the reviewers.


- We have addressed all other questions raised by the reviewers.


As the discussion period deadline approaches, we would greatly appreciate a **reconsideration of the scores** if there are no further questions. However, if any additional concerns arise, we remain fully committed to addressing them promptly.
Thank you for the follow-up questions. Please see our responses below.


**1. On gradient memory of $BA + S$.**

No, we never need to store the temporary $(BA + S)$ matrix for gradient computation and this is one of the main points we highlight in Section 3.2 (Line 181 of the main paper). We only need to store the necessary components in order to compute the gradients, i.e., $B, A$ and the indices and values for $S$. We have defined a customized layer specifically for this purpose. Please refer to Algorithm 1 and provided codes in supplementary for more details. 



**2. On combining SLTrain with pipelining.**
 
We have not tried out such pipelined scenarios in the paper. However, a straightforward combination of SLTrain with pipelining may not yield memory savings. For such settings, we would consider the approach of computing $BAx$ and $Sx$ separately albeit with a trade-off on computational cost.   

In summary, while our method is applicable in pipelined scenarios, achieving the full efficiency gains may require additional memory management strategies when layers are evaluated in parallel.






**3. On choice of hyperparameters.**

In practice, we often can try out different combinations to understand the trade-offs better. Please refer to Table 5 of the main paper, where we compare the model's performance across various choices of sparsity parameter $\delta$ and rank $r$. Generally, a higher sparsity parameter allows the model to achieve comparable performance with a lower rank, as you correctly noted.
Thank you for your detailed feedback. Based on your comments, it appears that the parameter efficiency of SLTrain, a central contribution of our paper (alongside memory efficiency), may not have been fully recognized. We would like to emphasize that parameter efficiency is a key aspect of our approach, as highlighted in both the abstract and introduction, and substantiated through various experiments presented in the main manuscript and supplementary material.



To further substantiate our claims, we have conducted an additional experiment involving the training of LLaMA 350M with an increased $\delta = 0.1$. The results are presented in the table below, where we also include a comparison with GaLore to better illustrate the advantages of SLTrain.




|  | PPL| Mem | Param Size|
|----------|--------|--------|--------|
| Full-Rank    | 18.80 |59.34G| 368M  |
| GaLore | 18.95 (+0.8\%)  | 58.35G (-1.7\%)  | 368M (-0\%) |
| SLTrain ($\delta=0.05$)    | 19.24 (+2.3\%)  | 58.00G (-2.2\%) | 200M (-45\%) |
| SLTrain ($\delta=0.1$)  | 18.72 (-0.4\%) | 58.25G (-1.8\%) | 215M (-42\%) |

As shown in the table, SLTrain with $\delta = 0.1$ not only matches but slightly outperforms the Full-Rank model in terms of perplexity, while requiring only 58\% of its parameter size and maintaining memory efficiency. Across all experiments, SLTrain consistently offers the best trade-off between perplexity, memory usage, and parameter size when compared to other baselines.

We believe that parameter efficiency is particularly critical during post-pretraining stages. While many existing works, such as [1], focus on model pruning after pretraining, SLTrain directly trains a smaller model from the pretraining stage, effectively reducing parameter size from the outset.

In summary, the results presented in the table not only demonstrate that SLTrain can achieve performance comparable to the Full-Rank model while maintaining memory efficiency, but also highlight the significant parameter savings it offers.

**We hope this would lead to a re-evaluation of our work.** 


[1] Li, Y., Yu, Y., Zhang, Q., Liang, C., He, P., Chen, W., and Zhao, T. LoSparse: Structured compression of large language models based on low-rank and sparse approximation. In *ICML 2023*.
Out of curiosity: Can the authors elaborate a bit on the choice of the hyperparameters, specifically the sparsity parameter $\delta$ and rank $r$. How can one choose them in practise, and have they observed how they influence each other, e.g. a higher sparsity score means that the models require less rank, or something similar.
Thank you for the answer and clarifying remarks to point 2) and 3).

ad 1): The argument for the memory efficiency of the method is that the full weight matrix is only stored temporarily. Can the authors comment on the parallelizability constraints that this temporal full-matrix constraints bring?
Specifically: The method seemingly generates the full matrix during the layer evaluation of layer $k$, then discards it again to evaluate layer $k+1$? If so, can the authors elaborate how the gradient tape is stored? It must not store the temporary (AB + S) matrix, since then one would observe similar memory footprint as the full model. 
Second, can the authors elaborate on the practicability of combining their approach with, e.g. pipelining, where multiple layers are evaluated in parallel? Is their method still applicable, since in this scenario the temporarily stored AB+S needs to be constructed in many layers simultaneously, potentially leading to memory spikes.
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.