Efficient One-shot Compression via Low-Rank Local Feature Distillation (2025)

Abstract: Current structured pruning approaches for large language models typically involve two steps: (1) compression using 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. Moreover, prior research suggests that pretrained Transformer weights are not necessarily low-rank, unlike their activations, making one-shot structured pruning challenging. 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.

## introduction compressing large language models is a crucial challenge that can be achieved with several complementary approaches: quantization, pruning, compression, and distillation. we propose a new oneshot compression method that locally distills lowrank weights and satisfies the three objectives: to address the first constraint, we compress the model with a local distillation objective, as opposed to a global objective that is more costly. we further reduce the cost by compressing and distilling only a subset of all layers. about the second constraint, we initialize the compressed layers with method no custom kernel no ampere-only memory gain is one-shot calib. data sparsegpt @xcite ✗ ✗ ✗ ✓ -wanda @xcite ✗ ✗ ✗ ✓ -teal @xcite ✗ ✗ ✗ ✓ -shearedllama @xcite ✓ ✓ ✓ ✗ 52b minitron @xcite singular value decomposition (svd), which enables reducing the required number of gradient steps as well as the calibration dataset size. we further improve convergence by combining teacher and student activations with a joint distillation loss. we experimentally show the robustness of our approach by compressing several state-of-the-art small and large transformers, mixture-of-experts, and mamba llms. figure 1 gives an intuitive overview of the proposed method, which is detailed in section 4 and validated in section 6. ## related works we discuss next previous compression methods, highlight their limitations, and explain how our approach addresses them. pruning methods remove unimportant weights in the pre-trained model @xcite . structured pruning removes entire groups of parameters, which results in a smaller and faster model @xcite . @xcite propose a new gradient-based criterion to eliminate substructures in llms, while @xcite use a joint loss combining a pruning mask loss with the language modeling loss. however, optimizing these criteria can be computationally intensive. for example, the pruning step of sheared-llama @xcite is 5x expensive compared to standard lm training, according to the authors. in contrast, thanks to the local gradient updates, our approach is computationally efficient, allowing us to compress a 47b model within minutes on a single a100 gpu. regarding unstructured pruning, these methods do not provide any gains in terms of memory or speedup, at least with current algorithmic implementations. semi-structured pruning (e.g., 2:4 and 4:8) @xcite @xcite does not lead to memory gain but can speed up processing on kernels optimized for such matrix structures. on the other hand, our method, which directly shrinks matrices, saves memory across all hardware and leads to speed up, as fewer computations are performed. low-rank decomposition compresses a model by approximating its pre-trained matrices with lowerdimension ones. based on the observation by @xcite that neural networks have lower intrinsic dimensions, @xcite show that transformer language models also require lower intrinsic dimensions depending on the language task. however, @xcite ; @xcite show that the weights in some pretrained transformer models are full-rank compared to their activations and propose an activation-aware compression method. our work is related to these approaches, as we also decompose the matrices of the models. however, we apply a lightweight layer-wise feature distillation objective to better recover from compression and propose different distillation strategies to accelerate convergence and improve performance. distillation. in knowledge distillation (kd), a small student model learns to reproduce the outputs of a larger pre-trained model. standard kd is dataintensive when the student model's parameters are randomly initialized. @xcite recently distilled gemma-2 9b from a larger model using 8 trillion tokens, which is as costly as a standard pretraining stage. another approach is layer-wise distillation, which can be more suitable in some cases, as intermediate layers may contain certain information, such as speaker identity or phonemes in pretrained speech models @xcite @xcite . for example, tinybert @xcite ) initializes the student model by removing layers from the teacher and then distills the student on task-specific knowledge, combining layer-wise and global distillation losses, which took 3.5 days to train. this is costly for llms. instead of removing layers, we leverage the observation that layer activations are low-rank and approximate them with fewer parameters. additionally, we compress the models to remain generalist rather than task-specific. ## proposed approach we identify and handle next three potential limitations of the approach presented in section 3.2. non-linear feature approximation. the previous analytical solution for eq-3 is limited to activations from linear layers. to generalize to non-linear modules, we first observe that eq-3 can be seen as a feature distillation objective that may be optimized numerically by gradient descent rather than analytically by eigendecomposition. given the input batch @xmath0, we denote the output activations of the i th teacher module t (i) as in the linear case, this should converge towards the eigendecomposition solution described in section 3.2. however, as we show later, we extend the distillation process to non-linear modules, which requires numerical optimization. we also show that initializing the student's low-rank parameters with svd, rather than randomly, improves convergence. ## fast and memory-efficient compression computing the optimal compression rank of each matrix individually may be costly: while search algorithms can be a solution for smaller models, they can be difficult to scale to llms with billions of parameters. assuming a target compression rate n , we propose to consider three simple strategies: uniform, top, and bottom first compression strategies, which are illustrated in appendix a.1. uniform. this approach removes n % parameters to each weight matrix in the model. this strategy involves training all llm weights. bottom. this approach removes n % parameters from the model by prioritizing lower layers as detailed in algorithm 1. the core of the algorithm is line 7, where the layers are ordered in bottom-top order and ranks are sorted in decreasing order. the argument k controls the spread of compression: the lower its value the more weights in the bottom layer will be compressed first. the algorithm is also illustrated in steps 1 and 2 in figure 1 . this strategy is efficient as only a subset of weights in the bottom layers will be trained. consequently, since there is no need to forward through the whole model, only a subset of the weights needs to be loaded into gpu memory, making this approach scalable to larger models. top. this approach removes n % parameters from the model by prioritizing top layers. this consists roughly of reversing the order of layers in the line 7 of algorithm 1. however, this strategy requires loading all the model's parameters into gpu memory. it also involves forwarding through the entire model, even if only the top layer weights are being trained. ## other architectures and modalities in this section, we show the generalizability of our approach by evaluating it on the mamba architecture for text @xcite and on the speech modality with whisper @xcite . we further evaluate it on a lower-resource african language in appendix a.11. ## analysis in this section, we give some ablation studies. in all experiments, the models are compressed for 20%. which distillation loss to use? how to choose the rank values? in our bottomfirst compression strategy, a small rank k in algorithm 1 compresses more severely the bottom layers and leaves more upper layers uncompressed, hence saving gpu memory. however, a small k may also destroy knowledge and the model will struggle to recover during the distillation phase. figure 4 compares 4 rank values for mistral-v0.1 7b and phi-3 14b and shows that higher ranks lead to better performance, but are also more costly as more layers are distilled. is svd initialization necessary? figure 5 shows the perplexity curves for phi-2 3b and phi-3 14b when the low-rank matrices are initialized with svd or randomly. svd initialization enables faster convergence, with 8 million tokens being sufficient, i.e., roughly 4k examples of 2048 tokens. cost. all models in section 6 are compressed in less than one hour on a single a100 gpu. the precise cost depends on the number of layers distilled determined by the compression strategy (bottom, top, or uniform) and the minimum rank k of algorithm 1: see appendix a.1 for details. comparisons. ## conclusion we propose a new efficient compression method that requires far less calibration data than most state-of-the-art approaches and provides competitive performance for various models (dense and moe transformers, mamba), and modalities (text, speech). we show that combining svd initializaphi-2 3b phi-3 14b 0 2 5 7 9 12 0 2 5 7 9 12 100 10000 1000000 10 100 1000 10000 100000 number of seen tokens (millions) wikitext perplexity with svd without svd figure 5: svd initialization improves the sample-efficiency. convergence when initializing low-rank weights randomly (without svd) or with svd.