CoLA: Compute-Efficient Pre-Training of LLMs via Low-Rank Activation

Many previous works have observed the low-rank structure of model activations in deep neural networks Cui et al. (2020); Huh et al. (2021). We also observe this phenomenon in LLMs, i.e. the effective rank of the activations is much smaller than their original dimensionality. To quantify this, we define the effective rank $r(\alpha)$ of activation as the minimal number of singular values needed to preserve an $\alpha$-fraction of the total spectral energy. Formally:

$$r(\alpha)\;=\;\min\left\{k\;\middle|\;\frac{\sum_{i=1}^{k}\sigma_{i}^{2}}{\sum_{i=1}^{n}\sigma_{i}^{2}}\;\geq\;\alpha\right\},$$

(1)

where $\sigma_{1},\sigma_{2},\ldots,\sigma_{n}$ are the singular values of the activation matrix, and $0<\alpha\leq 1$ is the desired ratio of preserved information. As shown in our experiments, the rapid decay of singular values [Fig. 2a] leads to much smaller $r(\alpha)$ compared to the full dimension [Fig. 2b]. This highlights the significant low-rank nature in the activations of pre-trained LLMs. More results showing the same pattern can be found in Appendix 11.

Motivated by the observed low-rank nature of LLM activations, we propose to enforce explicit low-rank activation via injecting non-linearity between factorized weight matrices.

Let $\mathbf{W}\in\mathbb{R}^{d_{\text{out}}\times d_{\text{in}}}$ be the weight matrix of an arbitrary linear layer followed by a nonlinear activation in the transformer architecture:

$$\mathbf{h}=\sigma\left(\mathbf{Wx}\right),\;\text{with}\;\mathbf{x}\in\mathbb{R}^{d_{\text{in}}}.$$

(2)

We replace $\mathbf{W}$ by two low-rank matrices $\mathbf{A}\in\mathbb{R}^{r\times d_{\text{in}}}$ and $\mathbf{B}\in\mathbb{R}^{d_{\text{out}}\times r}$, where rank $r<\min(d_{\text{in},\text{out}})$ is a hyper-parameter, and inject a non-linear activation $\sigma$ in the middle. This modification results in a transformation consisting of $(\text{linear}\circ\text{non-linear}\circ\text{linear})$ operations:

$$\mathbf{h^{\prime}}=\mathbf{B}\,\sigma(\mathbf{A}\mathbf{x}).$$

(3)

During the backward step, we simply apply the chain rule to compute gradients w.r.t each of the low rank factors as

		$\displaystyle\nabla_{\mathbf{B}}=\nabla_{\mathbf{h}^{\prime}}\mathbf{z}^{T},\nabla_{\mathbf{z}}=\mathbf{B}^{T}\nabla_{\mathbf{h}^{\prime}},\nabla_{\mathbf{o}}=\nabla_{\mathbf{z}}\odot\sigma^{\prime}(\mathbf{o}),$		(4)
		$\displaystyle\nabla_{\mathbf{A}}=\nabla_{\mathbf{o}}\mathbf{x}^{T},\nabla_{\mathbf{x}}=\mathbf{A}^{T}\nabla_{\mathbf{o}},$

where $\mathbf{o}=\mathbf{Ax},\mathbf{z}=\sigma(\mathbf{o})$, $\odot$ denote the element-wise product. We empirically find that keeping the original nonlinearity on top of Eq. (3) does not harm the performance, nor necessarily brings benefits. However, applying Eq. (3) to all linear layers regardless of whether being followed by nonlinearity is crucial to boost model performance. We refer more details to the ablation study in Appendix 11.

Fig. 4 shows the architecture of each transformer block when adopting CoLA into the LLaMA architecture. We highlight the fact that only the original linear layers and (if any) their follow-up non-linear transformation are modified to the CoLA formulation. Other computations such as the scaled-dot product of the self-attention, as well as residual connections and the element-wise product of LLaMA’s feed-forward layers, remain unchanged.

We analyze and compare the computational complexity of CoLA with other efficient pre-training methods based on the LLaMA architecture. We adopt a similar notion from Kaplan et al. (2020), where a general matrix multiply (GEMM) between an $M\times N$ matrix and an $N\times K$ matrix involves roughly $2MNK$ add-multiply operations. We denote the model inner width as $d$, and the inner width of the feed-forward layer as $d_{\text{ff}}$. For simplicity, we only show non-embedding calculations of a single sequence with token batch size of $n$ for each decoder layer. This is because the total computation scales only linearly with the number of layers $n_{\text{layer}}$ and the number of sequences $n_{\text{seq}}$. Furthermore, lower-order cheap operations of complexity $\mathcal{O}(nd)$ or $\mathcal{O}(nd_{\text{ff}})$ are omitted, such as bias, layer norm, non-linear function, residual connection, and element-wise product.

We show the detailed cost of the full-rank training in Table. 2. Notice that we apply the $2\times$ rule when calculating the backward cost. This is because for each forward GEMM that Eq. (2) describes, two GEMMs are needed to compute gradients for both the weight matrix $\mathbf{W}$ and the input $\mathbf{x}$, and are of the same cost the forward GEMM, i.e.,

$$\nabla_{\mathbf{x}}=\mathbf{W}^{T}\nabla_{\mathbf{h}},\nabla_{\mathbf{W}}=\nabla_{\mathbf{h}}\mathbf{x}^{T}.$$

(5)

We apply the same analysis to all the following pre-training methods:

• •

  LoRA/ReLoRA Hu et al. (2021); Lialin et al. (2023): $\mathbf{h}_{\text{LoRA}}=\mathbf{W}_{0}\mathbf{x}+\mathbf{BAx}$, with fixed $\mathbf{W}_{0}$.

• •

  SLTrain Han et al. (2024): $\mathbf{h}_{\text{SLTrain}}=\mathbf{BAx}+\mathbf{Sx}=(\mathbf{BA}\oplus_{\mathcal{I}}\mathcal{V})\mathbf{x}$, where $\oplus$ denotes the scatter-add operator, $\mathcal{I}$ and $\mathcal{V}$ are the indices and values of non-zero elements in the sparse matrix $\mathbf{S}$.

• •

  GaLore Zhao et al. (2024): $\mathbf{R}_{t}=\mathbf{P}_{t}^{T}\mathbf{G}_{t},$ $\tilde{\mathbf{G}}_{t}=\mathbf{PN}_{t}$, where $\mathbf{P}_{t}$ projects the gradient $\mathbf{G}_{t}$ onto a low-rank space, and then projects it back when updating the full-rank weight $\mathbf{W}$.

We summarize the computational costs of these methods in Table 3 and observe that the costs of SLTrain and GaLore are lower bounded by full-rank training, while (Re)LoRA is lower bounded by CoLA when choosing the same rank. In contrast, CoLA reduces the computation from full-rank training when $r<0.62d$, assuming $d_{\text{ff}}\approx 2.5d$ in LLaMA-like architecture. The default rank choice is set to $r=\frac{1}{4}d$, leading to a reduction in compute to about half of the full-rank training. We refer all details of compute analysis to Appendix B.

We assume a common notion that training modern transformers with Adam (or AdamW) involves four key memory components Zhao et al. (2024); Han et al. (2024): model parameters, gradients, optimizer states, and activations (intermediate forward-pass results). Gradients consume $1\times$ model size, Adam consumes $2\times$, and activations typically consume $1\sim 4\times$, depending on the batch size.

We focus on the scenario where the memory cost determined by the model size is not in the extreme limit of the GPU. We argue that this is rather realistic, since the model size and the minimum required tokens should scale up simultaneously during pre-training Kaplan et al. (2020); Hoffmann et al. (2022); Krajewski et al. (2024); Kumar et al. (2024). A tiny batch size on a single GPU would be impractical. Therefore, we analyze memory usage on a 40-GB A100 or a 94-GB H100 GPU with a fairly large sequence batch size. Fig. 5 shows that activations dominate memory usage in this setup. Although our method overall consumes less memory than full-rank training (largely due to parameter efficiency), vanilla CoLA allocates more memory to activations. However, we argue that our unique architecture can significantly enhance existing memory-saving techniques and consume only a fraction of the original cost.

Gradient checkpointing (GCP) Chen et al. (2016) is a system-level technique that reduces memory usage by selectively storing (“checkpointing”) only a subset of intermediate activations during the forward pass. When the backward pass begins, the missing activations are recomputed on the fly instead of being stored in memory, thereby lowering the memory cost. A vanilla (also the most effective) implementation of GCP in LLM pre-training is to save merely the input and output of each transformer block, and re-compute everything within each block during the backward step. Some works have investigated the optimal selection of checkpoints through both empirical and compiler view Feng and Huang (2021); He and Yu (2023). Such techniques can also be developed for CoLA, and are beyond the scope of this paper.

Motivated by the bottleneck structure of CoLA, we implement CoLA-M as saving only the low-rank activations (red circles in Fig. 4), and re-compute the up projections, and (if applicable) the self-attention (painted in sketch in Fig. 4) during the backward pass. This reduces the re-computation cost to half of the CoLA forward. We analyze the memory and re-computation cost using the same notions as in Section 3.3 and denote $h$ as the number of attention heads. We further simplify the analysis under LLaMA architecture by uniformly assuming $d_{\text{ff}}\approx 2.5d$. The memory and re-computation overhead are shown in Table 4. We refer the detailed analysis to Appendix C.

Although delicate optimizations of GCP is beyond our scope, we show in Table 5 and in Fig. 7 the quantitative results and scaling behavior of GCP on LLaMA-1B when applying a heuristic checkpointing strategy. We observe that CoLA-M greatly reduces re-computation cost while achieving similar memory saving as vanilla GCP.

We validate our proposed methods by extensively pre-training LLaMA-like LLMs from 60M to 7B scales. Experiments were performed on NVIDIA A100/H100 GPUs. We closely follow the experiment settings of Zhao et al. (2024); Han et al. (2024), and directly compare CoLA with their reported results. We use the C4 dataset Raffel et al. (2020), which is a colossal, cleaned version of Common Crawl’s web crawl corpus. C4 dataset has been widely used for pre-training LLMs. Trainings were done without data repetition on a sufficiently large amount of tokens. We compare CoLA with baselines including full-rank pre-training, ReLoRA Hu et al. (2021), GaLore Zhao et al. (2024), and SLTrain Han et al. (2024), with a focus on methods that explore model efficiency.

We implement CoLA and its memory efficient variant CoLA-M by parameterizing all linear layers into the proposed linear-nonlinear-linear composition [i.e. Eq. (3)], and keep all other parameters and operations unchanged. We use AdamW optimizer and cosine annealing learning rate scheduler Loshchilov and Hutter (2016) with warm-up. We remark that CoLA is NO more sensitive to optimizer-related hyper-parameters. We refer more details to Appendix D.

Table 6 compares our methods and other efficient pre-training techniques in terms of validation perplexity, parameter size, and estimated memory usage of model, gradients and optimizer states. CoLA has the smallest model size, thereby consumes the least memory, and demonstrates on-par validation perplexity compared to the full-rank baselines. Table 7 compares the validation perplexity on the 7B model. We highlight the fact that CoLA outperforms 8-bit Adam/GaLore at their 150k steps (14.61 & 14.65) using only 65k steps. Meanwhile the total memory cost of CoLA-M is less than half of the other methods. We also emphasize our comparison with GaLore and SLTrain. Our proposed method has an overall better performance, whilst further reduces model size on top of the reduction of SLTrain. More importantly, we follow our discussion in Section 3.3 and remark that CoLA has uniformly fewer computations than GaLore and SLTrain when applying the same rank.

We briefly discuss how CoLA might scale differently compared to full-rank training. A comprehensive investigation of this topic is beyond the scope of this work due to the huge computing cost.

Table 8 shows a few experiments on how CoLA might be improved when computing is scaled up. The default rank choice reduces the compute cost to about a half of full-rank training, without significantly hurting the model performance. Meanwhile, if we relax the computing restriction and increase the rank to be greater than one quarter of $d$, then CoLA outperforms full-rank training in all three scales, while still being able to reduce the computing cost. One might argue that full-rank training can also be scaled down to a similar compute of CoLA and might perform similarly. We implement such baselines in Table 8 and refer this setup to “Control". We typically reduce the number of layers or the model width of full-rank models to scale down their computing cost. We find empirically that they will tend to reduce performance quickly and dramatically underperform CoLA.

We further investigate the efficiency of CoLA from a more practical perspective. It is often observed that a theoretically less expensive method can have worse system-level performance due to poorly designed hardware implementation or lack of system-aware optimizations. We show that this is NOT the case for CoLA, by illustrating that its out-of-the-box system performance already significantly outperforms the full-rank training and other efficient training methods. We focus on the actual memory usage and the pre-training throughput.

Fig. 8 shows the measured throughput for pre-training the LLaMA/CoLA 1B model. The sequence batch size is set as 16, which fully utilizes the A100 GPU. Among all these methods, CoLA and CoLA-M are the only two that show higher GPU throughput than the full-rank baseline, while all other methods downgrade throughput due to their computing overhead. In particular, CoLA-M, the memory-efficient CoLA that significantly reduces overall GPU memory, still shows higher training throughput despite the re-computation overhead. Meanwhile, vanilla GCP, which uses a similar idea of trading compute for memory, reduces throughput from the full-rank baseline by 26%. We show further the details of these measurements in Table 9, and also compare their estimated FLOPs. At both 1B and 7B scale, CoLA-M manages to almost halve the computing cost and reduce the memory usage by two thirds, achieving a great balance between computing and memory efficiency. We refer more details of our system profiling to Appendix. F.

We highlight the fact that CoLA not only reduces pre-training resources but also speeds up inference and reduces its memory cost. Table 10 shows that CoLA improves inference throughput by up to $\bf 1.64\boldsymbol{\times}$ while reducing memory cost by up to $\bf 1.67\boldsymbol{\times}$.