Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Lillama: Large Language Models Compression via Low-Rank Feature Distillation

Yaya Sy, Christophe Cerisara, Irina Illina

TL;DR

Lillama presents a data-efficient, one-shot compression framework that blends SVD-based initialization with a local, activation-aware distillation objective (Teacher+Student) to compress large language models. By distilling low-rank weight factors per layer and employing a bottom-first strategy, it achieves substantial parameter reductions with minimal zero-shot performance loss across diverse architectures, including Mixtral, Mistral, Phi, and Mamba, on modest calibration data. The method yields fast convergence, reduced memory, and up to ~20% speedups on A100 GPUs, while enabling larger context lengths and broad applicability beyond Transformers. These findings underscore a practical, hardware-friendly path to deploying large models at smaller scales without extensive retraining or data requirements.

Abstract

Current LLM structured pruning methods typically involve two steps: (1) compression with calibration data and (2) costly continued pretraining on billions of tokens to recover lost performance. This second step is necessary as the first significantly impacts model accuracy. Prior research suggests pretrained Transformer weights aren't inherently low-rank, unlike their activations, which may explain this drop. Based on this observation, we propose Lillama, a compression method that locally distills activations with low-rank weights. Using SVD for initialization and a joint loss combining teacher and student activations, we accelerate convergence and reduce memory use with local gradient updates. Lillama compresses Mixtral-8x7B within minutes on a single A100 GPU, removing 10 billion parameters while retaining over 95% of its original performance. Phi-2 3B can be compressed by 40% with just 13 million calibration tokens, resulting in a small model that competes with recent models of similar size. The method generalizes well to non-transformer architectures, compressing Mamba-3B by 20% while maintaining 99% performance.

Lillama: Large Language Models Compression via Low-Rank Feature Distillation

TL;DR

Lillama presents a data-efficient, one-shot compression framework that blends SVD-based initialization with a local, activation-aware distillation objective (Teacher+Student) to compress large language models. By distilling low-rank weight factors per layer and employing a bottom-first strategy, it achieves substantial parameter reductions with minimal zero-shot performance loss across diverse architectures, including Mixtral, Mistral, Phi, and Mamba, on modest calibration data. The method yields fast convergence, reduced memory, and up to ~20% speedups on A100 GPUs, while enabling larger context lengths and broad applicability beyond Transformers. These findings underscore a practical, hardware-friendly path to deploying large models at smaller scales without extensive retraining or data requirements.

Abstract

Current LLM structured pruning methods typically involve two steps: (1) compression with calibration data and (2) costly continued pretraining on billions of tokens to recover lost performance. This second step is necessary as the first significantly impacts model accuracy. Prior research suggests pretrained Transformer weights aren't inherently low-rank, unlike their activations, which may explain this drop. Based on this observation, we propose Lillama, a compression method that locally distills activations with low-rank weights. Using SVD for initialization and a joint loss combining teacher and student activations, we accelerate convergence and reduce memory use with local gradient updates. Lillama compresses Mixtral-8x7B within minutes on a single A100 GPU, removing 10 billion parameters while retaining over 95% of its original performance. Phi-2 3B can be compressed by 40% with just 13 million calibration tokens, resulting in a small model that competes with recent models of similar size. The method generalizes well to non-transformer architectures, compressing Mamba-3B by 20% while maintaining 99% performance.

Paper Structure

This paper contains 27 sections, 12 equations, 9 figures, 18 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Background: Low-Rank Approximation of Transformer Models
  4. Low-Rank Approximation from Weights.
  5. Low-Rank Approximation from Feature
  6. Proposed Approach
  7. Fast and Memory-Efficient Compression
  8. Experiments on Transformers
  9. Setup
  10. Results
  11. Experiments on Mamba Architecture
  12. Analysis
  13. Conclusion
  14. Limitations
  15. Appendix
  16. ...and 12 more sections

Figures (9)

  • Figure 1: Lillama approach: STEP 1 selects layers to compress for a target compression ratio (e.g., N%) using various strategies (see Section \ref{['sec:algo']}). STEP 2 compresses and initializes the chosen parameters via SVD. STEP 3 distills the low-rank weights with a small calibration dataset.
  • Figure 2: Accelerating convergence by learning from Teacher and Student activations through a joint loss. We propose to study the effect of three distillation strategies: (a) Teacher: the input to the compressed student layer comes from the output of the previous teacher layer; (b) Student: the input to the compressed student layer comes from the output of the previous student layer; (c) Teacher+Student: the compressed student layer receives both the output of the previous teacher layer and the output of the previous student layer, and the loss is the sum of two losses.
  • Figure 3: The joint loss converges generally better. Convergence of the three losses illustrated in Figure \ref{['fig:distillaton']}, evaluated on the Wikitext2 test corpus perplexity during distillation.
  • Figure 4: Mean accuracy over our test benchmarks as a function of the minimum rank $k$ in Algorithm \ref{['alg:two']}
  • Figure 5: Convergence when initializing low-rank weight randomly (without SVD) or with SVD.
  • ...and 4 more figures