LEMON : Reviving Stronger and Smaller LM s from Larger LM s with Linear Parameter Fusion (2024)

Abstract: In the new era of language models, small models (with billions of parameter sizes) are receiving increasing attention due to their flexibility and cost-effectiveness in deployment. However, limited by the model size, the performance of small models trained from scratch may often be unsatisfactory. Learning a stronger and smaller model with the help of larger models is an intuitive idea. Inspired by the observing modular structures in preliminary analysis, we propose LEMON to learn competent initial points for smaller models by fusing parameters from larger models, thereby laying a solid foundation for subsequent training. Specifically, the parameter fusion process involves two operators for layer and dimension, respectively, and we also introduce controllable receptive fields to model the prior parameter characteristics. In this way, the larger model could be transformed into any specific smaller scale and architecture. Starting from LLaMA 2-7B, we revive two stronger and smaller models with 1.3B and 2.7B. Experimental results demonstrate that the fusion-based method exhibits flexibility and outperforms a series of competitive baselines in terms of both effectiveness and efficiency.

## introduction transformer-based language models (lms) shine in various natural language tasks due to their powerful understanding and generation capabilities @xcite @xcite . considering that the training and deployment of large-scale lms require a large amount of computing resources, small lms are more cost-effective in actual production environments @xcite . however, in the pre-training stage, it has been proven that there is a certain scaling law between model size and performance @xcite . how to † corresponding author. ‡ project lead. break through this limitation to train a stronger and smaller model with the help of large models has recently attracted the attention of researchers @xcite @xcite . inspired by the idea of function-reserving transformation in efficient training @xcite , there are some model compression efforts that utilize layer dropping, structured pruning, or weight selection to preserve the parameter functionality of large models as much as possible, so as to accelerate pre-training and improve the final performance of small models @xcite @xcite . however, since there are tens or even hundreds of billions of parameters in modern lms, searching for the essential parameter combinations requires significant computing resources, while pruning inappropriate parameters inevitably leads to performance degradation in downstream tasks @xcite . in this case, how to construct an appropriate initial point with capabilities close to large lms at the lowest possible cost remains an open problem. given the limited understanding of parameter characteristics and roles in modern large lms, we perform an empirical analysis from the perspectives of representation similarity, and observe that there are some modular structures with high internal similarity. in particular, there is a high similarity in the parameters of adjacent layers, while the similarity of distant layers to a lesser extent, which coincides with the sparsity found in previous studies @xcite . nonetheless, for a parameter module, each part has its own unique role, and thus no model parameter should be arbitrarily discarded. based on the observation, we believe that the functionality of a module should be reconstructed by selectively fusing corresponding parameters during the initial point construction phase, as illustrated in figure 1 . in this paper, we propose lemon to construct compact initial points for small models through linear fusion of large model parameters. specifically, we first identify the importance of parameters in larger models with fusion operators, and then compress them into new parameters for the smaller model, maintaining the same functionality of related parameters before and after transformation. to support mapping a larger model into any specific architecture, we decompose the fusion process into layer and dimension operations. the layer operator establishes layer-to-layer correspondence between larger and smaller models, while the dimension operator aims to learn the fusion of different dimensions within each parameter matrix. recalling the modular structures among layers, we further factorize the fusion operators into kronecker products @xcite with controllable receptive fields, instead of directly learning it on a huge number of parameters. eventually, the fusion operators are initiated by the prior features of modular parameters and optimized on unlabeled data. based on well-learned operators, it is convenient to efficiently achieve compact initial points for smaller models and facilitate subsequent training. to verify the effectiveness, we conduct experiments on the llama 2-7b model @xcite and compress it into two smaller models, 1.3b and 2.7b, respectively. within only 0.26 billion tokens, it is sufficient to learn competent initial points for smaller models, and with an additional 50 billion tokens of post-training, lemon outperforms a series of powerful baselines on 11 representative downstream tasks. meanwhile, compared to traditional models trained from scratch @xcite @xcite , lemon only requires 5.26% of pre-training token budget. when achieving comparable performance to recent pruning-based method (xia et al., 2022), the learning process of our fusion-based method saves 62.5% of training budget. overall, lemon exhibits considerable flexibility with controllable layer and dimension operation, making it a unified view of existing model compression techniques, which explore a new way for model compression using in-model parameters fusion. ## parameters characteristics in large language models to figure out the inherent characteristics of the parameters in the large language model, we calculate the layer-by-layer parameter similarity of individual components with centered kernel alignment (stated in appendix a.6). from figure 2 (a) (here we provide the parameter similarity of multi-head attention in llama 2-7b, see appendix a.5 for the similarity of feed forward layer), one can observe that the parameters of large lms tend to form cluster-like modular during pre-training, in which the parameter within the cluster, especially in neighboring layers, shows a high degree of similarity, while there are significant differences in the layers outside the cluster. it is intuitive to speculate that these modules naturally distinguish the functional areas of large models, allowing for the processing of language features at different levels. specifically, the cluster-like similarity structure primarily emerges in the middle to deeper layers, which indicates a shift from explicit textual information processing in shallower layers to more abstract and progressive feature processing in deeper layers. on the contrary, the start and final layers of these models show distinct characteristics, highlighting their unique functions. this enlightens that, on the one hand, the observation of similar parameters shows that there is a large amount of redundancy in the pre-trained large model, which opens up the exploration of merging modules with high parameter similarity to reduce the model size. on the other hand, the highest parameter similarity is still less than 80% (see figure 2 (b) ), indicating that each part in the module has its unique function behavior and should not be directly discarded. further motivated by findings that deep neural networks' parameter optima can be connected via linear low-loss paths @xcite , is imperative to explore the linear fusion of parameter modules in large models to achieve model compression. ## method building a stronger small model typically involves two steps. the first step is to construct a compact initial point from a larger model while preserving its functional behavior, and the second step is to perform (restorative) pre-training to further improve performance. the second step is technically no different from training from scratch, we focus on the first step in this paper. formally, we define a model as a set of layer weights θ l,@xmath0, w 2 , • • • , w l ], where w @xmath1. in the initialization stage, our goal is to transfer knowledge from a larger model θ (large) l 1 ,d 1 to a smaller one θ (small) l 2 ,d 2 , where total layers l 1 > l 2 , hidden dimension d 1 > d 2 , by learning a fusion operation that minimizes the functional difference between original and fused parameters, ## experiments 4.1 setup data to learn lemon operators and continue pre-training, we use the redpajama (togetherai, 2023) dataset, mirroring the training data in llama across seven domains (commoncrawl, c4, github, wikipedia, books, arxiv, stack-exchange). the dataset includes a validation set of 2 million tokens, a training set of 4 billion tokens, and a continued pre-training set of 50 billion tokens. training our setup uses sheared-llama @xcite code base on composer package (team, 2021), tested on 8 nvidia a100 gpus (80gb). the models were trained at the sequence length of 4096 using a batch size of 32 during the pruning phase and continued pre-training phase. following the setup in sheared-llama @xcite , baseline models were trained for 6400 steps (22 hours) with dynamic batch gradient loading (doremi @xcite ) during the pruning phase. in contrast, lemon were trained for 1200 steps (8 hours) without doremi. for continued pre-training, we employed 50 billion tokens with doremi with 256 batch size (387 hours). learning rates were set to 1e-4 for language modeling loss, and for layer and dimension operators, 5e-4 and 5e-5, respectively. for more information, please refer to appendix a.3. evaluation following the multiple assessment metrics reported in the baseline studies @xcite , we utilized lm-evaluation-harness @xcite to evaluate our models. in common sense and reading comprehension evaluation, we report 0-shot accuracy of sciq @xcite , piqa @xcite , winogrande (abbreviated as wino.) @xcite , arc easy @xcite , hellaswag @xcite , 25-shot accuracy of arc challenge @xcite . in world knowledge evaluation, we report accuracy of 32-shot nq @xcite and 5-shot mmlu @xcite . in continued qa and text understanding evaluation, we report accuracy of 0-shot logiqa @xcite , 32-shot boolq @xcite and 0-shot lambada @xcite . ## conclusion in this paper, we propose a new paradigm for building smaller lms based on larger lms. by learning a parameter fusion operator with controllable receptive fields from a larger model to smaller models, we can obtain a good starting point to facilitate subsequent training. this fusion mapping operator consists of a layer operator and a dimension operator. experimental results demonstrate that our method can compress larger lms into smaller lms of arbitrary architectures and better preserve the knowledge in large models. by comprising the requirements of pre-training data, we demonstrate the effectiveness of the method against several baseline approaches in terms of training speedup and computational expenditure savings. the paradigm has greater research value and has the potential to reach smaller performance losses and lower computational effort in the future.