Surgical Feature-Space Decomposition of LLM s: Why, When and How? (2024)

Abstract: Low-rank approximations, of the weight and feature space can enhance the performance of deep learning models, whether in terms of improving generalization or reducing the latency of inference. However, there is no clear consensus yet on how, when and why these approximations are helpful for large language models (LLMs). In this work, we empirically study the efficacy of weight and feature space decomposition in transformer-based LLMs. We demonstrate that surgical decomposition not only provides critical insights into the trade-off between compression and language modelling performance, but also sometimes enhances commonsense reasoning performance of LLMs. Our empirical analysis identifies specific network segments that intrinsically exhibit a low-rank structure. Furthermore, we extend our investigation to the implications of low-rank approximations on model bias. Overall, our findings offer a novel perspective on optimizing LLMs, presenting the low-rank approximation not only as a tool for performance enhancements, but also as a means to potentially rectify biases within these models.

## introduction the mass adoption of large language models (llms) in various domains has brought to the forefront the critical need to understand and optimize their structural and functional properties. this understanding is not only crucial from a theoretical point of view, but also essential for practical, realworld applications. llms, characterized by their vast size and complexity, present significant challenges in terms of resource utilization and operational efficiency. in this context, model compression has become a crucial topic of research. in this work, we delve into the low-rank decomposition properties of llms. we meticulously examine and surgically decompose the individual layers of these models with the aim of achieving training-free compression while maintaining or, in some cases, enhancing performance. this layerwise decomposition approach allows for granular analysis of the model's components, leading to a more tailored and effective compression strategy. furthermore, we extend our study to investigate the effects of low-rank decomposition on the intrinsic biases inherent in llms. bias in language models is a topic of growing concern given the widespread application of these models in society. by examining how decomposition influences these biases, we contribute to the broader conversation on ethical ai and responsible model deployment. in summary, this work presents a comprehensive study of low-rank decomposition in llms, exploring its potential to achieve considerable compression levels and its impact on model performance and biases. through our empirical analyses, we aim to provide insights that can guide future de-1 2389 velopments in the field of language model optimization. the contributions of this paper can be summarized as follows: ## related work this work focuses on surgically creating low-rank approximations of the feature space for achieving efficient and effective compression of llms. further, we also study the behavior of the resultant models from a fairness perspective. in this regard, we present below an overview of the existing works related to these areas of research. surgical fine-tuning. recent works have shown that fine-tuning only specific layers of a deep neural network improves adaptations to distribution shifts @xcite . @xcite provided strong empirical results on surgical fine-tuning of the bert @xcite model. their results show that fine-tuning only specific layers often surpass downstream task performance on glue and superglue @xcite datasets. @xcite adopt an evolutionary search mechanism to identify layer-specific learning rates and freeze some layers in the process for better performance on few-shot learning tasks. @xcite perform layer wise auto-weighting using the fischer importance matrix for test-time adaptation. in this process, their method even makes certain layers nearly frozen to mitigate outliers. these works have shown that for any given downstream task, fine-tuning and/or freezing only specific layers lead to better performance as compared to the full model. we argue that not layers, but sub-components in each layer identified through traditional rank decomposition can be surgically eliminated form deep neural networks with minimal loss in performance. ## low-rank decomposition of llms before describing low-rank decomposition of llms, we first present here a brief understanding of the llm structure. any llm typically involves a deep neural network architecture based on the transformer model. without loss of generality, any transformer model has multi-head self attention (mhsa) and multi-layer perceptron (mlp) blocks repeated across the model depth. we refer to a combination of mhsa and mlp blocks as a single module. each mhsa block consists of three linear transformation layers corresponding to query, key and value. we denote the weight matrices corresponding to them as w q , w k and w v respectively. additionally, after the attention operation, an output linear transformation layer is present with weight denoted as w o . similarly, mlp block typically consists of a gate projection layer, an up projection layer and a down projection layer. we denote the weight matrices corresponding to them as w g , w u and w d respectively. combining the weights from the mhsa and mlp blocks, the weight space of a llm with l such repeated blocks can be stated as ## experimental analysis here, we offer various insights into the analysis of llms, focusing specifically on experiments conducted with llama-7b @xcite and mistral-7b models @xcite . our experiments are centered on downstream common-sense reasoning tasks, with detailed task descriptions available in the appendix a. throughout our experiments, we utilize a calibration dataset comprising 512 samples, each with a maximum sequence length of 128. importantly, this data is randomly selected from the train splits of the downstream datasets, ensuring no data leakage. it is worth not-ing that the results demonstrated at the ∼7b scale tend to extrapolate more effectively to larger-scale models @xcite . decomposition is carried out on a cpu machine and rank search is done on a single l4 gpu with 24gb of vram for faster evaluations during search. it is a possibility that a larger size of calibration data with longer sequence lengths can further provide better low-rank approximations. ## limitations while we have sufficiently demonstrated the efficacy of sfsd in compressing and building efficient llms, there are still a few limitations that need to be overcome before sfsd can be looked at the preferred full-blown solution for improving llms from the efficiency perspective. first is the scale of experimentation. while we have observed ourselves from very early experiments that sfsd works well on llms with parameters above 20b, the extensive experimentation presented in the paper is on 7b models, and for unhindered generalization, additional experiments with more parameters might be needed. next, our rank search strategy is currently surgical, and it is not clear yet whether it exploits the best out of the proposed method. this might be clearly limiting the extent of compression that can be achieved with minimal performance drop, and a better choice in terms of a first-principled approach might be able to develop more efficient llms.