Reproduce “PRAC: Principal-Random Subspace for LLM Activation Compression and Memory-Efficient Training” on OpenResearch

About this paper

Yanyi Li, Yimu Zhang, Cong Fang

Researchers from Peking University developed Principal-Random Subspace for LLM Activation Compression (PRAC), a framework that leverages the spectral structure of activations to reduce memory usage during large language model training. PRAC achieved up to 36% total memory reduction and enabled a 50% increase in batch size, leading to a 25% reduction in training time for LLaMA-1B while maintaining competitive model performance.

Key insights

• For effective activation compression, a gradient estimator must be both unbiased for accurate convergence and low-variance for a fast convergence rate.
• Activations in LLMs exhibit a 'degenerate condition,' characterized by a few dominant singular values and a long tail of smaller ones, which can be strategically exploited for compression.
• A hybrid principal-random subspace projection strategy, employing a specific scaling factor for the random component, can achieve the theoretical lower bound for variance while maintaining unbiasedness.

Summary by alphaXiv.