VeLoRA
Overview. Here we address the challenge of compressing intermediate activations tensors while preserving the necessary training dynamics for model convergence. Our memory-efficient algorithm consists of two components: (i) The grouping strategy to divide the original high-dimensional tokens into much smaller sub-tokens; and (ii) Fixed rank-1 projections of these sub-tokens using cheap heuristically initialized principal components. Given a large pre-trained model, we apply these steps to compress the intermediate activations saved during training while preserving most of the original model’s training dynamics. We illustrate the general overview of this pipeline in Fig 1 and show PyTorch-like pseudo-code in Algorithm 1 .

Consider a set of input activations that need to be saved in GPU memory during the forward pass. We denote an element in this set as \mathbf{Z}_{i}=\nabla_{w}f(\mathbf{x}_{i};w)\in{\rm I\!R}^{N\times D} , where N is the number of tokens, and D is the feature depth.
We propose to introduce a simple grouping (reshape) operation that partitions the tokens across the depth dimension: group(\cdot):\mathbf{Z}\in{\rm I\!R}^{B\times N\times D}\xrightarrow[]{}{\rm I\!R}^{B\times ND/M\times M} with M being the new size of each token, now coined a sub-token.
This operation can be described as partitioning a batch of N tokens into a collection of these much smaller non-overlapping ND/M sub-tokens.
Then we project each of the sub-tokens onto a rank-1 subspace. The idea is that this grouping operation enables a much lower-dimensional fixed subspace to be used throughout training without any degradation in model convergence or performance.
We describe the compression steps concisely as follows:

\mathbf{Z}\xrightarrow[]{group(\cdot)}\mathbf{z}\in{\rm I\!R}^{B\times ND/M\times M}\xrightarrow[]{compress(\cdot~{};~{}\mathbf{v})}\mathbf{z}_{p}\in{\rm I\!R}^{B\times ND/M\times 1}, (1)

where \mathbf{z} is used to denote the sub-tokens of \mathbf{Z} and \mathbf{z}_{p} are the compressed sub-tokens. This compression is achieved using the function compress(\mathbf{z}~{};~{}\mathbf{v})=\mathbf{z}\cdot\mathbf{v} , which projects each sub-token onto a one-dimensional sub-space before saving them in memory.
Since M<<D it is more memory efficient to store \mathbf{z}_{p} instead of \mathbf{Z} .
The initialization strategy for \mathbf{v} is important for the performance of the model, however we later show that a simple and cheap average over the first batch of sub-tokens can be very effective.
Finally, the compressed sub-tokens \mathbf{z}_{p} are reconstructed for the gradient calculation during backward pass as follows:

\mathbf{z}_{p}\xrightarrow[]{reconstruct(\cdot~{};~{}\mathbf{v})}\hat{\mathbf{z}}\in{\rm I\!R}^{B\times ND/M\times M}\xrightarrow[]{ungroup(\cdot)}\hat{\mathbf{Z}}\in{\rm I\!R}^{B\times N\times D}, (2)

Here \hat{\mathbf{z}} and \hat{\mathbf{Z}} refer to the reconstructed sub-tokens and tokens of \mathbf{z} and \mathbf{Z} respectively.
The reconstruct function projects the sub-tokens \mathbf{z}_{p} back onto the vector \mathbf{v} as a very coarse reconstruction of \mathbf{z} and it is defined as reconstruct(\mathbf{z}_{p}~{};~{}\mathbf{v})=\mathbf{z}_{p}\cdot\mathbf{v}^{T} .
The overall compression and reconstruction operation is given as proj_{\mathbf{v}}(\mathbf{z})=(\mathbf{z}\cdot\mathbf{v})\cdot\mathbf{v}^{T} , where \mathbf{v}\in{\rm I\!R}^{M\times 1} is a fixed vector of unit length.

To summarize, during the forward pass, we compress the intermediate activation tensor \mathbf{Z} into a compact representation \mathbf{z_{p}} using \mathbf{v} . Then, in the backward pass when the original activation \mathbf{Z} is needed for gradient computation, we reconstruct an approximation \hat{\mathbf{Z}} by projecting \mathbf{z}_{p} back onto \mathbf{v} .
These steps are fundamentally different and complementary to recent works that leverage the low-rank property of gradients like GaLore [ 53 ] in two ways: Firstly, they store the uncompressed intermediate activations in memory for the gradient computation. In contrast, we compress these activations explicitly during the forward pass. Secondly, GaLore relies on periodically computing the principal components with SVD to optimally capture the underlying gradient subspace. Our compression method avoids such costly operations, making it much more efficient and scalable.