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

TL;DR

Lillama introduces a low-rank activation distillation method for compressing large language models efficiently, achieving high compression ratios with minimal performance loss and faster convergence.

Contribution

The paper proposes a novel low-rank feature distillation approach that accelerates LLM compression, reducing memory and time requirements while maintaining high accuracy.

Findings

Compresses Mixtral-8x7B by 10 billion parameters in minutes.

Retains over 95% of original performance after compression.

Generalizes to non-transformer architectures like Mamba-3B.

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.