Adaptive Rank Selections for Low-Rank Approximation of Language Models (2024)

Abstract: Singular Value Decomposition (SVD) or its weighted variants has significantly progressed in compressing language models. Previous works assume the same importance for all operations and assign the same number of ranks for different layers in a language model. However, such a uniform rank selection is sub-optimal since different operations (layers) have non-uniform demand in capacity. In other words, a desired SVD strategy should allocate more ranks for important operations and vice versa. However, a globally-optimized selection of ranks for neural networks is still an open problem, and this is a non-trivial challenge since the selection is discrete. In this work, we propose a novel binary masking mechanism for optimizing the number of ranks in a differentiable framework. Our strategy uses a novel regularization to enable the masking to comply with the SVD property where the ranks have sorted singular values. The experiments examined both types of language models, encoder-only and decoder-only models, including large language models like LLaMA. Our compressed model achieves much better accuracy than previous SVD and their SOTA variants. More interestingly, our method retains significantly better accuracy with zero or limited fine-tuning, proving the substantial advantage of adaptive rank selection.

## introduction transformer-based models @xcite have been very popular across different natural language processing tasks, such as text classification @xcite , question answering @xcite , and summarization @xcite . despite its success on these tasks, the size of these models often scales up to millions or billions of parameters, especially for recently proposed large language models @xcite . such a huge number of parameters makes these models very hard to be deployed on resource-limited devices, such as mobile phones or edge devices. as a result, the compression of transformer-based language models has drawn much attention. transformers-based models have two core operations: self-attention layers and feed-forward layers. these operations are built on linear layers, making them straightforward to compression techniques like low-rank weight factorization @xcite with svd or its variants. low-rank weight factorization decomposes a large linear layer into two small linear layers without changing other model parts, providing a friendly property for deployment. in addition, it is orthogonal to other compression techniques, such as structural pruning @xcite , quantization @xcite , and knowledge distillation @xcite . previous work @xcite shows that using vanilla svd for compression can result in a significant performance drop. they argue that low reconstruction error is not equivalent to high accuracy. as a result, hsu et al. @xcite proposed to apply the fisher information @xcite ) matrix to re-weight the weight matrix so that the factorization results can capture information from both the task and the reconstruction error. empirically, fisher information weighted svd performs much better than the original svd. despite using the fisher information matrix, other importance scores, like first-order taylor expansion @xcite , can also be used to re-weight the weight matrix. although the mentioned weighted svd methods above achieved promising results, they treat all layers uniformly and use the same number of ranks for all weight matrices. on the other hand, some prior works suggest that the compression rate for different layers should be different in the cases of vision @xcite and language (la-gunas et al., 2021) models. these observations provide clues to improve the performance of existing weight factorization by selecting the proper number of ranks for each layer. inspired by the above observations, the target of our problem setting is to find the optimal number of ranks for all the layers in a neural network. however, this optimization is not trivial since it is a discrete, nonsmooth, and non-convex problem. reinforcement learning @xcite and evolutionary algorithms @xcite ) may find a solution for this problem, but they introduce substantial optimization costs that are not affordable for larger models. to address the above challenge, we propose to use regularized differentiable binary masks to learn the number of ranks for each operation. the entire learning pipeline is built upon an end-to-end differentiable learning framework. we use the sum of a binary mask to capture the number of ranks for each layer. the proposed binary mask is properly regularized to be aligned with the sorted singular values of svd. moreover, we use a hypernetwork to improve the effectiveness of our method, which further accelerates the learning process. with all these designs, our method can efficiently find the number of ranks of different operations. the contribution of our work can be summarized as the following points: ## related works the benefit of low-rank factorization is that it can be applied to any linear layer. an early work @xcite applies svd for the lstm cell and explores the effectiveness on different nlp tasks @xcite @xcite and model components. (noach and goldberg, 2020) propose a two-stage approach to compress a pre-trained language model. the first stage decomposes the weight matrix with svd in the pre-trained language model. then, they fine-tune weights with knowledge distillation to regain performance. the standard svd can not capture all the information from tasks. the fisher information is introduced to reweight the weight matrix, and svd is applied to the reweighted matrix @xcite . on top of @xcite , several numeric optimization methods are used to find the optimal solution to the weighted svd problem @xcite when the weighting matrix is not diagonal. besides model weights, svd can also be applied to embedding layers. the albert model @xcite addresses the issue of redundant parameters in the embedding layer by employing factorization. this layer tends to have high input and output dimensions, leading to inefficiencies. in their work, reid et al. @xcite introduce a novel approach called self-attentive factorized embeddings (safe). this method enhances performance by incorporating a small self-attention layer built upon linear projection. a crucial point omitted by previous works is that not all operations are created equally. some operations require more capacity than others. our method tackles this problem by automatically learning the number of ranks for each operation. our method is also related to network pruning methods, especially structural pruning. block pruning @xcite integrates structures of any size into the movement pruning paradigm for fine-tuning, and it prunes the model globally. in addition to nlp tasks, deciding the width of a convolution layer has also been studied extensively using reinforcement learning @xcite , evolutionary algorithm @xcite , etc. differentiable pruning @xcite @xcite @xcite ) is also a popular direction since the cost is often not high. however, they can not be directly applied to select the number of ranks due to the cost or difficulty of fine-tuning resulting from using binary masks. ## conclusion in this paper, we proposed a new algorithm that adaptively selects the number of ranks for low-rank approximation of language models. we proposed to use a hypernetwork to predict the number of ranks for each operation. the predicted number of ranks is regularized using the svd property and is encouraged to produce top-k-like binary masks. our method resolved the issue with the ordinary masking that resulted in element-wise rank selec-tions, delivering stable performance gain in a comprehensive collection of experiments. the extensive results also show our advantage over previous low-rank methods with uniform rank selections.