A da M ix: Mixture-of-Adaptations for Parameter-efficient Model Tuning (2022)

Abstract: Standard fine-tuning of large pre-trained language models (PLMs) for downstream tasks requires updating hundreds of millions to billions of parameters, and storing a large copy of the PLM weights for every task resulting in increased cost for storing, sharing and serving the models. To address this, parameter-efficient fine-tuning (PEFT) techniques were introduced where small trainable components are injected in the PLM and updated during fine-tuning. We propose AdaMix as a general PEFT method that tunes a mixture of adaptation modules – given the underlying PEFT method of choice – introduced in each Transformer layer while keeping most of the PLM weights frozen. For instance, AdaMix can leverage a mixture of adapters like Houlsby or a mixture of low rank decomposition matrices like LoRA to improve downstream task performance over the corresponding PEFT methods for fully supervised and few-shot NLU and NLG tasks. Further, we design AdaMix such that it matches the same computational cost and the number of tunable parameters as the underlying PEFT method. By only tuning 0.1-0.2% of PLM parameters, we show that AdaMix outperforms SOTA parameter-efficient fine-tuning and full model fine-tuning for both NLU and NLG tasks.

## introduction standard fine-tuning of large pre-trained language models (plms) @xcite @xcite to downstream tasks requires updating all model parameters. given the ever-increasing size of plms (e.g., 175 billion parameters for gpt-3 @xcite and 530 billion parameters for mt-nlg @xcite ), even the fine-tuning step becomes expensive as it requires storing a full copy figure 1 : performance of different parameter-efficient fine-tuning methods on glue development set with roberta-large encoder following a setup similar to @xcite for fair comparison. we report the performance of pfeiffer @xcite , houlsby @xcite and @xcite with their default number of fine-tuned parameters as well as the number of fine-tuned parameters used in adamix with a mixture of adaptations . red dash shows the performance of full model fine-tuning. in this paper, we present adamix, a mixture of adaptation modules approach, and show that it outperforms sota peft methods and also full model fine-tuning while tuning only 0.1 -0.2% of plm parameters. in contrast to traditional peft methods that use a single adaptation module in every transformer layer, adamix uses several adaptation modules that learn multiple views of the given task. in order to design this mixture of adaptations, we take inspiration from sparsely-activated mixture-of-experts (moe) models. in traditional dense models (e.g., bert @xcite , gpt-3 @xcite )), all model weights are activated for every input example. moe models induce sparsity by activating only a subset of the model weights for each incoming input. consider adapters @xcite , one of the most popular peft techniques, to illustrate our method. a feedforward layer (ffn) is introduced to down-project the hidden representation to a low dimension d (also called the bottleneck dimension) followed by another up-project ffn to match the dimensionality of the next layer. instead of using a single adapter, we introduce multiple project-up and project-down ffns in each transformer layer. we route input examples to one of the project-up and one of the project-down ffn's resulting in the same amount of computational cost (flops) as that of using a single adapter. for methods like lora @xcite , that decomposes the gradient of pre-trained weights into low-rank matrices (a and b), we introduce multiple lowrank decompositions and route the input examples to them similar to adapters. we discuss different routing mechanism and show that stochastic routing yields good performance while eliminating the need for introducing any additional parameters for module selection. to alleviate training instability that may arise from the randomness in selecting different adaptation modules in different training steps, we leverage consistency regularization and the sharing of adaptation modules during stochastic routing. the introduction of multiple adaptation modules results in an increased number of adaptation parameters. this does not increase computational cost but increases storage cost. to address this, we develop a merging mechanism to combine weights from different adaptation modules to a single module in each transformer layer. this allows us to keep the number of adaptation parameters the same as that of a single adaptation module. our merging mechanism is inspired by model weight averaging model soups @xcite and multi berts @xcite . weight averaging of models with different random initialization has been shown to improve model performance in recent works @xcite @xcite ) that show the optimized models to lie in the same basin of error landscape. while the above works are geared towards fine-tuning independent models, we extend this idea to parameter-efficient fine-tuning with randomly initialized adaptation modules and a frozen language model. overall, our work makes the following contributions: (a) we develop a new method adamix as a mixture of adaptations for parameter-efficient fine-tuning (peft) of large language models. given any peft method of choice like adapters and low-rank decompositions, adamix improves downstream task performance over the underlying peft method. (b) adamix is trained with stochastic routing and adaptation module merging to retain the same computational cost (e.g., flops, #tunable adaptation parameters) and benefits of the underlying peft method. to better understand how adamix works, we demonstrate its strong connections to bayesian neural networks and model ensembling. we present background on mixture-of-experts (moe) and adapters in section a of appendix. ## mixture-of-adaptations consider a set of m adaptation modules injected in each transformer layer, where a ij : @xmath0, @xmath1. for illustration, we will consider adapters (houlsby ffn_d 1 + ffn_d 2 ffn_d 3 ffn_d 4 ffn_d 1 + ffn_d 2 ffn_d 3 ffn_d 4 input batch mixture-of-adapter regularization loss 1 p1 multi-head attention add & norm feedforwardintermediate embedding add & norm feedforwardoutput softmax adapter p2 multi-head attention add & norm feedforwardintermediate embedding add & norm feedforwardoutput softmax adapter mixture-of-adapter ffn_u 1 ffn_u 2 ffn_u 3 ffn_u 4 ffn_u 1 ffn_u 2 ffn_u 3 ffn_u 4 figure 2: mixture-of-adaptations (adamix) with adapters (houlsby et al., 2019) as the underlying peft mechanism. for illustration, we show @xmath2. the above block shown for one transformer layer is repeated across all the layers. adamix stochastically routes instances from an input batch via randomly selected adaptation modules resulting in flops match to a single module with consistency regularization and parameter sharing. adaptation merging (figure 3) collapses multiple modules to match single-module parameters in each layer. et al., 2019) as the underlying parameter-efficient fine-tuning (peft) mechanism as a running example. similar principles can be used for other peft mechanism like lora (hu et al., 2021) for lowrank decompositions as we show in experiments. we adopt the popularly used transformer architecture @xcite consisting of l repeated transformer blocks, where each block consists of a self-attention sub-layer, a fully connected feed-forward network (ffn) and residual connections around the sub-layers followed by layer normalization. each adaptation module a ij corresponding to the adapters @xcite consists of a feedforward up w up ij and a feedforward down w down ij projection matrices. ## conclusions we develop a new framework adamix for parameter-efficient fine-tuning (peft) of large pretrained language models (plm). adamix leverages a mixture of adaptation modules to improve downstream task performance without increasing the computational cost (e.g., flops, parameters) of the underlying adaptation method. we demonstrate adamix to work with and improve over different peft methods like adapters and low rank decompositions across nlu and nlg tasks. by tuning only 0.1 -0.2% of plm parameters, adamix outperforms full model fine-tuning that updates all the model parameters as well as other state-of-the-art peft methods.