FlexiGPT: Pruning and Extending Large Language Models with Low-Rank Weight Sharing

TL;DR

FlexiGPT addresses memory-constrained deployment of large language models by pruning blocks with a Block Influence score and replacing them via a low-parameter weight-sharing scheme that leverages block-specific LoRA adapters and output normalization. A low-rank SVD-based selection of sharing bases plus careful adapter initialization minimizes initial disruption and facilitates rapid adaptation, while also supporting model extension by repeating blocks with unique adapters. Empirical results show state-of-the-art performance on multiple benchmarks at 30% and full coverage at 40% pruning on LLaMA-2 7B, with consistent improvements across LLaMA-3 and OPT models, and effective extension of TinyLLaMA using only about 0.3% additional training tokens. Overall, FlexiGPT enables on-device NLP with substantial parameter efficiency and flexible model extension, offering practical gains for memory-constrained applications while preserving high task performance.

Abstract

The rapid proliferation of large language models (LLMs) in natural language processing (NLP) has created a critical need for techniques that enable efficient deployment on memory-constrained devices without compromising performance. We present a method to prune LLMs that selectively prunes model blocks based on an importance score and replaces them with a low-parameter replacement strategy. Specifically, we propose a principled metric to replace each pruned block using a weight-sharing mechanism that leverages unpruned counterparts from the model and block-specific low-rank adapters. Furthermore, we facilitate the learning of these replacement blocks with output feature normalization and an adapter initialization scheme built on low-rank SVD reconstructions. Empirical evaluations demonstrate substantial performance gains over existing methods, achieving state-of-the-art performance on 5/6 benchmarks for a compression rate of 30% and 6/6 benchmarks for a compression rate of 40%. We also demonstrate that our approach can extend smaller models, boosting performance on 6/6 benchmarks using only ~0.3% tokens of extended training with minimal additional parameter costs.

Figures (4)

• Figure 1: FlexiGPT is used for two settings: (1) pruning a model to reduce parameters with minimal performance cost or (2) extending a model to increase performance with minimal parameter cost. Left: For pruning models (setting 1), we prune entire blocks and replace them using weight sharing and learned adapters. Right: For extending models (setting 2), we repeat block patterns in the model using weight sharing and learned adapters.
• Figure 2: Comparison of block distance score versus block index distance ($i-j$) for different metrics. (a) Using the proposed metric in Eq. \ref{['eq:ws-score']} with high-rank pruning, showing that closer blocks score lower (better), matching our intuition that weights close in the model have similar function. (b) Ablation of high-rank pruning, where there is no clear trend except that blocks closer to 0 are lower and those closer to 31 are higher. (c) Simple Frobenius norm, showing a similar lack of clear trend as in (b). We found that using the score in (a) as the weight-sharing selection metric results in a much higher performing model compared to using the scores in (b) and (c).
• Figure 3: Overview of the FlexiGPT pruning process. Left: We prune model blocks with the lowest scores based on \ref{['eq:prune-score']}. Center: We select replacement blocks with high similarity using \ref{['eq:select-base-score']}. Right: We add feature normalization and learn adapters to recover performance.
• Figure 4: Comparison of block distance score versus candidate block index $j$ for different metrics. Dotted lines represent where candidate block index $j$ is equal to pruning block index $i$, which is not a valid candidate. (a) Using the proposed metric in Eq. \ref{['eq:ws-score']} with high-rank pruning, showing that closer blocks score lower (better), matching our intuition that weights close in the model have similar function. (b) Ablation of high-rank pruning, where there is no clear trend except that blocks closer to 0 are lower and those closer to 31 are higher. (c) Simple Frobenius norm, showing a similar lack of clear trend as in (b). We found that using the score in (a) as the weight-sharing selection metric results in a much higher performing model compared to using the scores in (b) and (c).