Efficient Learning of Multiple NLP Tasks via Collective Weight Factorization on BERT (2022)

Abstract: The Transformer architecture continues to show remarkable performance gains in many Natural Language Processing tasks. However, obtaining such state-of-the-art performance in different tasks requires fine-tuning the same model separately for each task. Clearly, such an approach is demanding in terms of both memory requirements and computing power. In this paper, aiming to improve training efficiency across multiple tasks, we propose to collectively factorize the weighs of the multi-head attention module of a pre-trained Transformer. We test our proposed method on finetuning multiple natural language understanding tasks by employing BERT-Large as an instantiation of the Transformer and the GLUE as the evaluation benchmark. Experimental results show that our method requires training and storing only 1% of the initial model parameters for each task and matches or improves the original fine-tuned model’s performance for each task while effectively decreasing the parameter requirements by two orders of magnitude. Furthermore, compared to well-known adapter-based alternatives on the GLUE benchmark, our method consistently reaches the same levels of performance while requiring approximately four times fewer total and trainable parameters per task.

## introduction transformer-based language models such as bert @xcite , roberta @xcite and t5 @xcite have shown remarkable performance in many natural language processing (nlp) tasks, including language understanding @xcite , machine translation @xcite ) and text generation @xcite , to mention but a few examples. such deep language models are first pretrained on a large-scale unlabelled text dataset, and are then fine-tuned on various downstream tasks by employing labeled data. to this end, sequential transfer learning is adopted, where all the pretrained model parameters are optimized on the downstream taskspecific loss. despite the simplicity and the practical success of such a training scheme, the exponential increase in the size of the models and training data can make it difficult and costly for researchers and practitioners with limited computational resources to benefit from these models in domain or task-specific applications. to mitigate this issue, many recent papers have focused on learning multiple tasks simultaneously by sharing most of the parameters of a given model between tasks, and introducing a task-specific classifier on top of the shared network @xcite . along with cutting down on parameter costs through the aforementioned parameter sharing scheme, multi-task learning also improves the overall performance of the pretrained model across all tasks. however, training on multiple tasks requires access to all task-specific data simultaneously, hindering the extension of such models to new domains incrementally. an alternative approach to parameter sharing is the introduction of task specific non-shared parameters, also known as adapters, to the initial model architecture. adapters @xcite are typically small non-shared feedforward networks (ffn) inserted at each layer of a main network (e.g., bert). for each task, a separate set of adapters needs to be trained. in particular, the standard practice is to freeze the main model during task training and update only the adapter parameters, introducing a huge parameter cost reduction and the ability to incrementally train a model for an arbitrary number of tasks without having access to each dataset at the same time. however, the effectiveness of adapters is subject to extensive architecture search, that includes not only the adapter structure, but also their position in the overall network architecture, as well as the evaluation of parameter reduction and performance trade-offs @xcite . motivated by the above mentioned shortcomings, we introduce a tensor-based method for adapting a pre-trained transformer to new unseen nlp tasks by collective weight factorization of network weights. in particular, our method relies on collecting the multi head attention (mha) module weight matrices into a high-order tensor and represent it in tucker format @xcite , by applying a full-rank tucker decomposition to it (see fig. 1 ). this allows training only the decomposition factor matrices while keeping the rest of the network frozen, leading to significant parameter cost savings without any architecture search. in short, our contributions are the following: the organization of the rest of the paper is as follows. section 2 discusses closely related work. in section 3, we introduce the proposed methodology. in section 4, we provide implementation details of our method and present experimental results. conclusions are drawn in section 5. ## related work this work lies at the intersection of multi-source domain adaption, while further incorporating tensor decompositions for the parametrization of the transformer architecture. a primary approach for fine-tuning deep learning models in the context of multi-source domain adaptation is adapter learning. adapters were introduced in computer vision as a way to efficiently fine-tune a model by adding extra adapter modules to the network @xcite and have been extended to the nlp domain as well. specifically, involving the transformer architecture @xcite , stickland and @xcite fine-tune the pretrained transformer bert @xcite with a multi-task strategy by training task-specific adapters along with the rest of the model. @xcite propose using adapters to efficiently learn multiple tasks in an incremental manner by fine-tuning only the task specific model adapters each time. likewise, @xcite adapt the pretrained multilingual model xlm-r @xcite to other languages through training language-specific adapters. in @xcite , a simpler architecture is presented where adapters are only added after the feed-forward part of the attention module. the placement and internal structure of adapters is non-trivial and may impact model efficiency, thus adapters require extensive experimentation before they are deployed @xcite . in contrast, the method proposed in this paper is straightforward to implement, without requiring any architecture search. additionally, the number of trainable parameters is even less than adapter-based counterparts, while performance is comparable or better. instead of adding extra task-specific parameters to the model as in the adapter case, a number of works aim to directly modulate the layer weights for each task. these methods are mostly based on the hypernetwork scheme @xcite in which the weights of a network are generated dynamically by using another network. an alternative strategy is modulating the activations of the network layers as seen in @xcite where question encodings are used to modulate the activation maps of a convolutional network. @xcite use a hypernetwork scheme based on decomposable projections to create efficient multi-task transformers utilising training strategies similar to @xcite . in the context of deep learning, tensor methods are often used to compress a model or make its operations faster @xcite . a prominent idea in this area is reshaping the weight matri-ces of network layers into high order tensors and decomposing them to achieve high compression rates. similar to the decomposition of convolutional @xcite and fully connected layers @xcite , @xcite propose to compress the embedding matrices of nlp architectures by reshaping them into tensors and then use the tensor-train (tt) decomposition (oseledets, 2011) on them. the same strategy is used in @xcite where all the weight matrices of an rnn or an lstm unit are again decomposed into a tt format. the main novelty behind our work lies in utilizing the tucker decomposition in a collective weight factorization context, in order to efficiently extend pretrained transformer models to new tasks within the nlp domain. we show that selectively grouping and decomposing specific layers (such as the mha modules) can lead to efficient nlp task adaptation with little to no architecture search, while retaining strong performance. ## collective weight factorization for in this section, we present our method which is based on the collection of a pretrained network's weight matrices into a single high-order tensor. we then introduce task-specific factor matrices by decomposing the said tensor using a tucker decomposition on the collected weight matrices and training the factor components of the decomposition via backpropagation. a visual overview of the proposed method is presented in fig. 1 . ## experimental evaluation in this section, we provide implementation details and discuss our experiments. our models are based on the existing open source pytorch implementation of a well-known transformer model foot_0 , while tensorly @xcite is used for the tensor operations. ## conclusions and future work we introduced a collective weight factorization method for adapting a pre-trained transformer to multiple nlp tasks in an efficient manner. the proposed method is more efficient in terms of total parameters needed and percentage of model parameters trained and stored for each new task. furthermore, we investigated the effect of our method when applied to different components of the transformer. the results indicate that the multi-head attention mechanism can effectively adapt to different tasks and achieve competitive results when it is the only trained component, leading to the conclusion that more research needs to be done on the contribution of the different parts of the transformer to its performance, and the correlation of the performance to the data used for training. a limitation of out work is that the proposed factorization scheme operates on a single type of module, i.e. requiring same number of weight matrices and matrix dimensions. we plan to extend our method in order to be able to incorporate different types of modules (such as the ffn and mha) in a joint factorization scheme. taking into account recent advances in the extension of pretrained transformers to different modalities @xcite , one can also investigate the extension of the proposed method to adapt pretrained transformer networks to new modalities, incorporating diverse tasks while maintaining a minimum overhead in terms of additional trainable and total parameters.