From Low Rank Gradient Subspace Stabilization to Low-Rank Weights: Observations, Theories, and Applications (2407.11239v2)

Abstract: LLMs' (LLMs) weight matrices can often be expressed in low-rank form with potential to relax memory and compute resource requirements. Unlike prior efforts that focus on developing novel matrix decompositions, in this work we study the non-uniform low-rank properties of weight matrices in LLMs through the lens of stabilizing gradient subspace. First, we provide a theoretical framework to understand the stabilization of gradient subspaces through Hessian analysis. Second, we empirically establish an important relationship between gradient dynamics and low-rank expressiveness of weight matrices. Our findings reveal that different LLM components exhibit varying levels of converged low-rank structures, necessitating variable rank reduction across them to minimize drop in performance due to compression. Drawing on this result, we present Weight Low-Rank Projection(WeLore) that unifies weight compression and memory-efficient fine-tuning into one, in a data-agnostic and one-shot manner. When used as a compression technique, WeLore categorizes weight matrices into Low-rank Components (LRCs) and Non-Low-rank Components (N-LRCs) and suitably encodes them for minimum performance loss. Our gradient dynamics perspective illustrates that LRCs tend to have better fine-tuning capabilities and their standalone fine-tuning can closely mimic and sometimes outperform the training loss trajectory and performance of full fine-tuning with notable memory and compute footprint reduction. Codes are available at https://github.com/VITA-Group/WeLore.

• The paper shows that gradient dynamics drive the emergence of non-uniform low-rank structures, enabling targeted compression in LLMs.
• It introduces the WeLore method that categorizes layers into low-rank and non-low-rank components to optimize fine-tuning and reduce memory usage.
• Experimental results demonstrate significant performance gains, with improved perplexity and reduced parameter counts in models like LLaMa-2.

From GaLore to WeLore: How Low-Rank Weights Non-uniformly Emerge from Low-Rank Gradients

Overview

The paper explores the emergence of low-rank structures within large matrices used in modern LLMs and introduces Weight Low-Rank Projection (WeLore), a technique that leverages these low-rank structures for effective model compression and memory-efficient fine-tuning. The authors diverge from traditional methods that uniformly apply low-rank approximations to all layers, revealing that different layers of LLMs exhibit varying degrees of low-rank expressiveness. They establish a consequential relationship between gradient dynamics and the emergence of low-rank structures, allowing for a non-uniform rank reduction across different layers to minimize performance degradation.

Key Contributions

1. Gradient Dynamics:
   • The paper begins by investigating gradient behaviors during back-propagation, identifying that gradients for some layers in LLMs (e.g., middle MLP layers) quickly saturate, while others (e.g., attention layers in terminal transformer blocks) continue to accumulate rich error signals, fostering low-rank gradient subspaces.
   • Consequentially, layers that consistently exhibit rich gradient dynamics tend to have stable low-rank structures in their weight matrices.
2. Layer Categorization:
   • Layers are categorized into Low-rank Components (LRCs) and Non-Low-rank Components (N-LRCs) based on their ability to express low-rank structures. LRCs show a heavy-tail distribution in their singular values, making them suitable for significant rank reduction without substantial loss of information.
3. WeLore Method:
   • WeLore introduces a non-uniform rank reduction strategy leveraging the heavy-tail property of singular values. By decomposing LRCs into low-rank matrices, WeLore achieves significant compression ratios while maintaining performance.
   • A novel proposal is made to back-propagate only through LRCs during fine-tuning, allowing for memory-efficient training by confining full parameter optimization to layers with rich gradient dynamics.

Results

The experimental evaluation of WeLore validates its efficacy through a series of empirical assessments:

1. Compression:
   • WeLore's adaptive rank reduction significantly outperforms uniform and outlier-weighted rank reduction strategies. For example, WeLore achieves a perplexity improvement of up to 47 times over a 40% uniform rank reduction in the LLaMa-2 13B model.
2. Memory Efficiency:
   • Inference memory requirements are substantially reduced. For instance, a 50% compressed LLaMa-2 7B model with WeLore requires 0.67 times the parameters and achieves memory reductions of up to 0.45 times for a sequence length of 4096.
3. Fine-Tuning:
   • Empirical results demonstrate that WeLore's fine-tuning strategy matches or even surpasses the performance of dense full-parameter finetuning. Fine-tuning LRCs and freezing N-LRCs achieves comparable performance with lower computational and memory expenses. For example, a 50% compressed LLaMa-2 7B model finetuned with WeLore achieves 3x throughput and 0.35x trainable parameters compared to full finetuning.

Theoretical and Practical Implications

WeLore introduces a scalable methodology for effectively compressing and fine-tuning LLMs by recognizing and exploiting the non-uniform emergence of low-rank structures in weight matrices. This has several implications:

• Theoretical:
  • Establishing a direct correlation between gradient dynamics and low-rank weight subspaces provides a new lens through which model compression can be viewed. This opens avenues for exploring other gradient-oriented optimization techniques for better model efficiency.
• Practical:
  • By utilizing WeLore, organizations can deploy high-performance LLMs on consumer-grade GPUs, making sophisticated AI technologies more accessible and reducing the dependency on large-scale high-performance computing infrastructures.

Future Directions

The paper paves the way for several future developments:

• Extending the gradient-driven low-rank decomposition strategy to other model architectures beyond transformers, further generalizing the approach.
• Combining WeLore with other compression techniques such as sparsity and quantization to explore synergistic effects and maximize compression benefits without significantly impacting performance.
• Investigating the scalability of WeLore for extremely large-scale LLMs, such as GPT-4, to validate its robustness and efficiency in ultra-large model environments.
• Developing a more sophisticated understanding of the relationship between gradient dynamics and weight matrix structure across a broader range of tasks and datasets to refine the methodology further.

In summary, WeLore stands out as a sophisticated, data-agnostic technique for LLM compression and fine-tuning, guided by a nuanced understanding of gradient dynamics and their impact on low-rank expressiveness.