PRAC: Principal-Random Subspace for LLM Activation Compression and Memory-Efficient Training

Activations have become the primary memory bottleneck in large-batch LLM training. However, existing compression methods fail to exploit the spectral structure of activations, resulting in slow convergence or limited compression. To address this, we bridge the relationship between the algorithm’s fast convergence and the requirements for subspace projection, and show that an effective compression should yield an unbiased estimate of the original activation with low variance. We propose Principal-Random Subspace for LLM Activation Compression (PRAC), which novelly decomposes activations into two components: a principal subspace captured via SVD to retain dominant information, and a random subspace sampled from the orthogonal complement to approximate the tail. By introducing a precise scaling factor, we prove that PRAC yields an unbiased gradient estimator with minimum variance under certain conditions. Extensive experiments on pre-training and fine-tuning tasks demonstrate that PRAC achieves up to 36\% total memory reduction with negligible performance degradation and minimal computational cost.

Training large language models like ChatGPT requires enormous computer memory. A big part of that memory is used to store intermediate data – called “activations” – that the model computes between layers. When using larger batch sizes to speed up training, these activations can quickly exceed the memory capacity of a single GPU. We introduce PRAC, a method that compresses activations during training without hurting the model’s learning. The key idea is simple: split the activation data into two parts. One part keeps the most important information (the “principal” directions). The other part uses random sampling to capture the remaining details, which would otherwise be too costly to store. A special scaling factor then combines them to ensure that the compressed signal remains unbiased. Experiments show that PRAC reduces total memory usage by up to 36% while maintaining the same training performance as the standard method.