Conditionally Adaptive Multi-Task Learning: Improving Transfer Learning in NLP Using Fewer Parameters & Less Data

1 INTRODUCTION . The introduction of deep , contextualized Masked Language Models ( MLM ) 1 trained on massive amounts of unlabeled data has led to significant advances across many different Natural Language Processing ( NLP ) tasks ( Peters et al. , 2018 ; Liu et al. , 2019a ) . Much of these recent advances can be attributed to the now well-known BERT approach ( Devlin et al. , 2018 ) . Substantial improvements over previous state-of-the-art results on the GLUE benchmark ( Wang et al. , 2018 ) have been obtained by multiple groups using BERT models with task specific fine-tuning . The “ BERT-variant + fine-tuning ” formula has continued to improve over time with newer work constantly pushing the state-of-the-art forward on the GLUE benchmark . The use of a single neural architecture for multiple NLP tasks has shown promise long before the current wave of BERT inspired methods ( Collobert & Weston , 2008 ) and recent work has argued that autoregressive language models ( ARLMs ) trained on large-scale datasets – such as the GPT family of models ( Radford et al. , 2018 ) , are in practice multi-task learners ( Brown et al. , 2020 ) . However , even with MLMs and ARLMs trained for multi-tasking , single task fine-tuning is usually also employed to achieve state-of-the-art performance on specific tasks of interest . Typically this fine-tuning process may entail : creating a task-specific fine-tuned model ( Devlin et al. , 2018 ) , training specialized model components for task-specific predictions ( Houlsby et al. , 2019 ) or fine-tuning a single multi-task architecture ( Liu et al. , 2019b ) . ∗Joint first-authors 1For reader convenience , all acronyms in this paper are summarized in section A.1 of the Appendix . Single-task fine-tuning overall pretrained model parameters may have other issues . Recent analyses of such MLM have shed light on the linguistic knowledge that is captured in the hidden states and attention maps ( Clark et al. , 2019b ; Tenney et al. , 2019a ; Merchant et al. , 2020 ) . Particularly , BERT has middle Transformer ( Vaswani et al. , 2017 ) layers that are typically the most transferable to a downstream task ( Liu et al. , 2019a ) . The model proxies the steps of the traditional NLP pipeline in a localizable way ( Tenney et al. , 2019a ) — with basic syntactic information appearing earlier in the network , while high-level semantic information appearing in higher-level layers . Since pretraining is usually done on large-scale datasets , it may be useful , for a variety of downstream tasks , to conserve that knowledge . However , single task fine-tuning causes catastrophic forgetting of the knowledge learned during MLM ( Howard & Ruder , 2018 ) . To preserve knowledge , freezing part of a pretrained network and using Adapters for new tasks have shown promising results ( Houlsby et al. , 2019 ) . Inspired by the human ability to transfer learned knowledge from one task to another new task , Multi-Task Learning ( MTL ) in a general sense ( Caruana , 1997 ; Rajpurkar et al. , 2016b ; Ruder , 2017 ) has been applied in many fields outside of NLP . Caruana ( 1993 ) showed that a model trained in a multi-task manner can take advantage of the inductive transfer between tasks , achieving a better generalization performance . MTL has the advantage of computational/storage efficiency ( Zhang & Yang , 2017 ) , but training models in a multi-task setting is a balancing act ; particularly with datasets that have different : ( a ) dataset sizes , ( b ) task difficulty levels , and ( c ) different types of loss functions . In practice , learning multiple tasks at once is challenging since negative transfer ( Wang et al. , 2019a ) , task interference ( Wu et al. , 2020 ; Yu et al. , 2020 ) and catastrophic forgetting ( Serrà et al. , 2018 ) can lead to worse data efficiency , training stability and generalization compared to single task fine-tuning . Using Conditionally Adaptive Learning , we seek to improve pretraining knowledge retention and multi-task inductive knowledge transfer . Our contributions are the following : • A new task conditioned Transformer that adapts and modulates pretrained weights ( Section 2.1 ) . • A novel way to prioritize tasks with an uncertainty based multi-task data sampling method that helps balance the sampling of tasks to avoid catastrophic forgetting ( Section 2.2 ) . Our Conditionally Adaptive Multi-Task Learning ( CA-MTL ) approach is illustrated in Figure 1 . To the best of our knowledge , our work is the first to explore the use of a latent representation of tasks to modularize and adapt pretrained architectures . Further , we believe our work is also the first to examine uncertainty sampling for large-scale multi-task learning in NLP . We show the efficacy of CA-MTL by : ( a ) testing on 26 different tasks and ( b ) presenting state-of-the-art results on a number of test sets as well as superior performance against both single-task and MTL baselines . Moreover , we further demonstrate that our method has advantages over ( c ) other adapter networks , and ( d ) other MTL sampling methods . Finally , we provide ablations and separate analysis of the MT-Uncertainty Sampling technique in section 4.1 and of each component of the adapter in 4.2 . 2 METHODOLOGY . This section is organized according to the two main MTL problems that we will tackle : ( 1 ) How to modularize a pretrained network with latent task representations ? ( 2 ) How to balance different tasks in MTL ? We define each task as : Ti , { pi ( yi|xi , zi ) , Li , p̃i ( xi ) } , where zi is task i ’ s learnable shallow embedding , Li is the task loss , and p̃i ( xi ) is the empirical distribution of the training data pair { xi , yi } , for i ∈ { 1 , . . . , T } and T the number of supervised tasks . The MTL objective is : min φ ( z ) , θ1 , ... , θT T∑ i=1 Li ( fφ ( zi ) , θi ( xi ) , yi ) ( 1 ) where f is the predictor function ( includes encoder model and decoder heads ) , φ ( z ) are learnable generated weights conditioned on z , and θi are task-specific parameters for the output decoder heads . z is constructed using an embedding lookup table . 2.1 TASK CONDITIONED TRANSFORMER . Our task conditioned Transformer architecture is based on one simple concept . We either add conditional layers or modulate existing pretrained weights using a task representation by extending Feature Wise Linear Modulation ( Perez et al. , 2018 ) functions in several ways depending on the Transformer layer . We define our framework below . Definition 1 ( Conditional Weight Transformations ) . Given a neural network weight matrix W , we compute transformations of the form φ ( W|zi ) = γi ( zi ) W + βi ( zi ) , where γi and βi are learned functions that transform the weights based on a learned vector embedding zi , for task i . Definition 2 ( Conditionally Adaptive Learning ) . In our setting , Conditionally Adaptive Learning is the process of learning a set of φs for the conditionally adaptive modules presented below along with a set of task embedding vectors zi for T tasks , using a multi-task loss ( see equation 1 ) . In the subsections that follow : We introduce a new Transformer Attention Module using blockdiagonal Conditional Attention that allows the original query-key based attention to account for task-specific biases ( section 2.1.1 ) . We propose a new Conditional Alignment method that aligns the data of diverse tasks and that performs better than its unconditioned and higher capacity predecessor ( section 2.1.2 ) . We adapt layer normalization statistics to specific tasks using a new Conditional Layer Normalization module ( section 2.1.3 ) . We add a Conditional Bottleneck that facilitates weight sharing and task-specific information flow from lower layers ( section 2.1.4 ) . In our experiments we provide an ablation study of these components ( Table 1 ) examining performance in terms of GLUE scores . 2.1.1 CONDITIONAL ATTENTION Given d , the input dimensions , the query Q , the key K , and the value V as defined in Vaswani et al . ( 2017 ) , we redefine the attention operation : Attention ( Q , K , V , zi ) ) = softmax [ M ( zi ) + QKT√ d ] V M ( zi ) = N⊕ n=1 A′n ( zi ) , A ′ n ( zi ) = Anγi ( zi ) + βi ( zi ) where ⊕ is the direct sum operator ( see section A.6 ) , N is the number of block matrices An ∈ R ( L/N ) × ( L/N ) along the diagonal of the attention matrix , L is the input sequence , M ( zi ) = diag ( A′1 , . . . , A′N ) is a block diagonal conditional matrix . Note that An is constructed using L/N trainable and randomly initialized L/N dimensional vectors . While the original attention matrix depends on the hidden states h , M ( zi ) is a learnable weight matrix that only depends on the task embedding zi ∈ Rd . γi , βi : Rd 7→ RL 2/N2 are Feature Wise Linear Modulation ( Perez et al. , 2018 ) functions . We also experimented with full-block Conditional Attention ∈ RL×L . Not only did it have N2 more parameters compared to the block-diagonal variant , but it also performed significantly worse on the GLUE development set ( see FBA variant in Table 10 ) . It is possible that GLUE tasks derive a certain benefit from localized attention that is a consequence of M ( zi ) . With M ( zi ) , each element in a sequence can only attend to other elements in its subsequence of length L/N . In our experiments we used N = d/L . The full Conditional Attention mechanism used in our experiments is illustrated in Figure 2 . 2.1.2 CONDITIONAL ALIGNMENT . Wu et al . ( 2020 ) showed that in MTL having T separate alignment modules R1 , . . . , RT increases BERTLARGE avg . scores on five GLUE tasks ( CoLA , MRPC , QNLI , RTE , SST-2 ) by 2.35 % . Inspired by this work , we found that adding a task conditioned alignment layer between the input embedding layer and the first BERT Transformer layer improved multi-task model performance . However , instead of having T separate alignment matrices Ri for each T task , one alignment matrix R̂ is generated as a function of the task embedding zi . As in Wu et al . ( 2020 ) , we tested this module on the same five GLUE tasks and with BERTLARGE . Enabling task conditioned weight sharing across covariance alignment modules allows us to outperforms BERTLARGE by 3.61 % . This is 1.26 % higher than having T separate alignment matrices . Inserting R̂ into BERT , yields the following encoder function f̂ : f̂ = T∑ t=1 gθi ( E ( xi ) R̂ ( zi ) B ) , R̂ ( zi ) = Rγi ( zi ) + βi ( zi ) ( 2 ) where xi ∈ Rd is the layer input , gθi is the decoder head function for task i with weights θi , E the frozen BERT embedding layer , B the BERT Transformer layers and R the linear weight matrix of a single task conditioned alignment matrix . γi , βi : Rd 7→ Rd are Feature Wise Linear Modulation functions .

Summary: This paper proposed a new transformer model framework for multitask learning on NLP tasks. To deal with challenges in multitask learning/co-training, the authors proposed five improvements, including modifications on the transformer layers with task conditioning and uncertainty sampling. In the experiments, the authors showed that the proposed model can outperform full fine-tuned BERT large model with less parameters (adding some parameters on a single co-trained model), and with less training data.