F lexi GPT : Pruning and Extending Large Language Models with Low-Rank Weight Sharing (2025)

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.

## introduction the widespread adoption of llms has revolutionized nlp applications, driving significant advancements in areas such as virtual assistants, automated customer support, and real-time language translation @xcite . however, deploying these models on memoryconstrained devices, such as smartphones and edge devices, remains a formidable challenge due to their substantial parameter sizes and computational demands @xcite . this paper addresses this challenge by presenting a novel approach that targets parameter efficiency to make llms more suitable for on-device applications with minimal performance compromises. (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. parameter efficiency is particularly critical as it directly impacts the feasibility of deploying llms on devices with limited memory and storage resources. recent model pruning techniques, such as slicegpt @xcite , llm surgeon @xcite , llm-pruner @xcite , laco @xcite , and shortgpt @xcite , reduce the number of parameters but often result in significant performance degradation with minimal recovery after pruning. this gap in existing techniques underscores the need for an end-to-end pruning method that not only reduces the model size but also facilitates performance recovery. in this work, we propose to recover performance by utilizing existing weights within the model. specifically, we introduce a comprehensive pruning strategy combined with an innovative weight sharing technique and low-rank adapters (lora) @xcite , facilitating efficient parameter usage while preserving the model's performance. we begin by pruning model blocks based on shortgpt's block influence (bi) score @xcite . to replace the pruned blocks, we introduce a low-parameter weight-sharing mechanism that leverages existing block modules within the model and incorporates block-specific lora parameters, ensuring the selected replacement blocks have high similarity to the pruned blocks while maintaining block diversity. furthermore, we introduce a novel method to initialize the lora adapters in weight-sharing blocks, setting them to be the low-rank difference between the pruned block and the weight-shared replacement block. this initialization minimizes initial disruptions and facilitates smoother model adaptation. finally, we incorporate output feature normalization for pruned blocks to ensure a smooth transition and adaptation, allowing the model to gradually learn and stabilize its performance over time. empirical evaluations of our method, which we refer to as flexigpt, demonstrate substantial performance gains over existing methods. specifically, we achieve 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% for the popular llama-2 7b model @xcite . as visualized in figure 1 , our proposed technique not only effectively prunes large models for ondevice deployment but also extends smaller models, improving their performance at minimal additional parameter costs. specifically, our method shows that a 22-layer tinyllama @xcite model can be extended with repeated blocks, boosting performance on 6/6 benchmarks using only ≈0.3% tokens of extended training with minimal additional parameter costs. in summary, we make the following contributions: pruning -pruning selectively removes less important parameters, reducing model size and computational complexity while maintaining performance. its greatest benefit lies in optimizing llms for deployment in resource-constrained environments such as mobile devices, facilitating faster inference. several works have been proposed to reduce the size of llms by pruning model structures. llm-pruner (ma et al., 2023) removes unimportant coupled structures and the importance is calculated from taylor expansions. slicegpt (ashkboos et al., 2024) applies orthogonal projections to the feature maps and then it performs pruning in the projected space. llm surgeon (van der ouderaa et al., 2024) periodically updates model weights and structures, resulting in a higher cost compared to other methods. besides reducing the width of the model, shortgpt (men et al., 2024) is proposed to remove blocks by using block influence scores, and laco (yang et al., 2024) is proposed to collapse layers. existing pruning methods primarily focus on removing redundant model weights, often neglecting the loss of model capacity. our approach addresses this limitation by sharing weights from the pruned model to restore its capacity. peft -parameter-efficient fine-tuning (peft) methods aim to mitigate the extensive computational and memory demands of fine-tuning large models by focusing on a smaller subset of parameters. one prominent category of peft methods is adapters, which involves adding trainable modules to the existing frozen layers of the model @xcite . another significant category is prompt methods, which augment the initial input sequence with additional trainable vectors known as prompts. this technique focuses on fine-tuning these added tokens rather than the entire model, as demonstrated in works such as @xcite ; @xcite . recently, lora @xcite has emerged as the most efficient and highest-performing peft approach. lora introduces the use of low-rank matrices to adjust model weights efficiently, merging with pre-trained weights before inference to maintain the model's operational speed. building on this, dora @xcite decomposes pretrained weights into magnitude and direction components for fine-tuning, focusing on fine-tuning the directional components using lora. our research extends these low-rank peft methods by incorpo-rating lora and normalization for efficient weight sharing. similar to dora, our method involves a normalization stage; however, we normalize in the feature-space instead of the weight-space. weight-sharing -for gpt llms, @xcite proposed to share all the layers with a dynamic halting mechanism to improve accuracy on the downstream tasks. however, it requires the number of parameters of the base layer (unshared) to match the number of parameters of all layers of vanilla transformers @xcite . subformer @xcite applies a sandwich-like method of parameter sharing where only middle layers are shared but it does not use any adapter. @xcite developed efficient cyclic sharing patterns to increase the accuracy, however their sharing patterns are mainly based on ablation studies. mobilellm @xcite proposed sub-billion parameter architectures for mobile devices and adopted immediate block-wise weight sharing for further accuracy improvement. however, they do not use any pruning and adapters. @xcite introduced matching functions to develop head-wise shareable attention in a principled fashion. although they use pretrained weights for faster convergence, their matching functions are only applied to share weights among multiple heads of the same layer. svd -the singular value decomposition (svd) of a matrix @xmath0, @xmath1, where @xmath2, and σ ∈ r m×n is a diagonal matrix of singular values. svd is widely used to obtain a low-rank representation of w by selecting the k most significant singular values and their corresponding singular vectors, where @xmath3, n). hence, the low-rank representation of w is given as: ## flexigpt in this section, we detail our approach, flexigpt (figure 3 ), which prunes and extends llms using lora adapters, weight sharing techniques, and output feature normalization. our method focuses on achieving parameter efficiency while minimizing performance degredation, particularly for memory-constrained devices. our method is based on the transformer architecture @xcite , which consists of multi-head self-attention (mhsa) and multi-layer perceptron (mlp) layers. however, our approach is not constrained to this architecture. in general, we refer to blocks (mhsa+mlp), layers (mhsa or mlp), and weights (denoted as w ), as our method affects the weights in a uniform manner. ## conclusion in this paper, we presented an approach to pruning and extending llms using lora and weightsharing techniques. our method targets memoryconstrained devices by selectively pruning model blocks based on an importance score and replacing them with a low-parameter replacement strategy. empirical evaluations show substantial performance gains over existing methods, highlighting our technique's effectiveness. furthermore, our approach can extend smaller models, achieving significant performance improvements with minimal additional parameters. this work paves the way for more accessible and efficient on-device nlp applications, leveraging our novel combination of pruning, weight-sharing, and parameter-efficient adapters, thereby bringing the power of llms to a broader range of memory-constrained devices and use cases.