Learning without Forgetting: Task Aware Multitask Learning for Multi-Modality Tasks

1 INTRODUCTION . The process of Multi-Task Learning ( MTL ) on a set of related tasks is inspired by the patterns displayed by human learning . It involves a pretraining phase over all the tasks , followed by a finetuning phase . During pretraining , the model tries to grasp the shared knowledge of all the tasks involved , while in the finetuning phase , task-specific learning is performed to improve the performance . However , as a result of the finetuning phase , the model forgets the information about the other tasks that it learnt during pretraining . Humans , on the other hand , are less susceptible to forgetfulness and retain existing knowledge/skills while mastering a new task . For example , a polyglot who masters a new language learns to translate from this language without losing the ability to translate other languages . Moreover , the lack of task-based flexibility and having different finetuning/pretraining phases cause gaps in the learning process due to the following reasons : Role Mismatch : Consider the MTL system being trained to perform the Speech Translation ( ST ) , Automatic Speech Recognition ( ASR ) and Machine Translation ( MT ) tasks . The Encoder block has a very different role in the standalone ASR , MT and ST models and hence we can not expect a single encoder to perform well on all the tasks without any cues to identify/use task information . Moreover , there is a discrepancy between pretraining and finetuning hampering the MTL objective . Task Awareness : At each step in the MTL , the model tries to optimize over the task at hand . For tasks like ST and ASR with the same source language , it is impossible for the model to identify the task and alter its parameters accordingly , hence necessitating a finetuning phase . A few such examples have been provided in Table 1 . Humans , on the other hand , grasp the task they have to perform by means of context or explicit cues . Although MTL strategies help the finetuned models to perform better than the models directly trained on those tasks , their applicability is limited to finding a good initialization point for the finetuning phase . Moreover , having a separate model for each task increases the memory requirements , which is detrimental in low resource settings . In order to achieve the goal of jointly learning all the tasks , similar to humans , we need to perform shared learning in synergy with task-specific learning . Previous approaches such as Raffel et al . ( 2019 ) trained a joint model for a set of related text-to-text tasks by providing the task information along with the inputs during the joint learning phase . However , providing explicit task information is not always desirable , e.g. , consider the automatic multilingual speech translation task . In order to ensure seamless user experience , it is expected that the model extracts the task information implicitly . Thus , a holistic joint learning strategy requires a generic framework which learns task-specific information without any explicit supervision . In this work , we propose a generic framework which can be easily integrated into the MTL strategies which can extract task-based characteristics . The proposed approach helps align existing MTL approaches with human learning processes by incorporating task information into the learning process and getting rid of the issues related to forgetfulness . We design a modulation network for learning the task characteristics and modulating the parameters of the model during MTL . As discussed above , the task information may or may not be explicitly available during the training . Hence , we propose two different designs of task modulation network to learn the task characteristics ; one uses explicit task identities while the other uses the examples from the task as input . The model , coupled with the modulation network , jointly learns on all the tasks and at the same time , performs the task-specific learning . The proposed approach tackles issues related to forgetfulness by keeping a single model for all the tasks , and hence avoiding the expensive finetuning phase . Having a single model for all the tasks also reduces memory constraints , improving suitability for low resource devices . To evaluate the proposed framework , we conduct two sets of experiments . First , we include the task information during MTL on text-to-text tasks to show the effect of task information . Secondly , we train a model on tasks with different modalities and end goals , with highly confounding tasks . Our proposed framework allows the model to learn the task characteristics without any explicit supervision , and hence train a single model which performs well on all the tasks . The main contributions of this work are as follows : • We propose an approach to tackle the issue of forgetfulness which occurs during the finetuning phase of existing MTL strategies . • Our model , without any finetuning , achieves superior performance on all the tasks which alleviates the need to keep separate task-specific models . • Our proposed framework is generic enough to be used with any MTL strategy involving tasks with multiple modalities . 2 TASK-AWARE MULTITASK LEARNING . An overview of our proposed approach is shown in Figure 1 . 2.1 BASE MODEL . In general , the sequence-to-sequence architecture consists of two components : ( 1 ) an encoder which computes a set of representationsX = { x1 , · · · , xm } ∈ Rm×d corresponding to x , and a decoder coupled with attention mechanism ( Bahdanau et al. , 2015 ) dynamically reads encoder ’ s output and predicts target language sequence Y = { y1 , · · · , yn } ∈ Rn×d . It is trained on a dataset D to maximize the p ( Y |X ; θ ) , where θ are parameters of the model . We use the Transformer Vaswani et al . ( 2017 ) as our base model . Based on the task modalities , we choose the preprocessing layer in the Transformer , i.e. , speech or the text ( text-embedding ) preprocessing layer . The speech preprocessing layer consists of a stack of k CNN layers with stride 2 for both time and frequency dimensions . This layer compresses the speech sequence and produces the output sequence such that input sequences corresponding to all the tasks have similar dimensions , d. The overview of the base sequence-to-sequence model is shown in the rightmost part of Figure 1 . 2.2 TASK MODULATION NETWORK . The task modulation network performs two operations . In the first step , it computes the task characteristics ( te ) using the task characteristics layer . It then modulates the model parameters θ using te in the second step . 2.2.1 TASK CHARACTERISTICS NETWORK : . We propose two types of Task Characteristics Networks ( TCN ) to learn the task characteristics , where one uses explicit task identities while the other uses source-target sequences as input . Explicit Task Information : In this approach , the tasks involved are represented using different task identities and fed as input to this TCN as one hot vectors . This network consists of a feed-forward layer which produces the task embedding used for modulating the model parameters . te = FFN ( e ) , ( 1 ) where e ∈ Rs is a one-hot encoding of s tasks used during joint learning . Implicit Task Information : The Implicit TCN computes the task embeddings using example sequences from the tasks without any external supervision . It consists of four sub-layers : ( 1 ) Sequence Representation Layer , ( 2 ) Bi-directional Attention Layer , ( 3 ) Sequence Summary Layer , and ( 4 ) Task Embedding Layer . The sequence representation sub-layer consists of uni-directional Transformer Encoder ( TE ) blocks Vaswani et al . ( 2017 ) . It takes the source and target sequences from the tasks as input and produces self-attended source and target sequences . Xsa = TE ( X ) , Y sa = TE ( Y ) , ( 2 ) whereXsa ∈ RM×d , Y sa ∈ RN×d . This sub-layer computes the contextual representation of the sequences . The Bi-directional Attention ( BiA ) sub-layer takes the self-attended source and target sequences from the previous layer as input and computes the relation between them using Dot-Product Attention Luong et al . ( 2015 ) . As a result , we get target aware source ( Xat ∈ RM×d ) and source aware target ( Y asRN×d ) representations as outputs . Xat = BiA ( Xsa , Y sa ) , Y as = BiA ( Y sa , Xsa ) . ( 3 ) The sequence summary sub-layer is similar to the sequence representation sub layer and summarizes the sequences . The sequence summaries are given by : Xs = TEu ( X at ) , Y s = TEu ( Y as ) , ( 4 ) whereXs ∈ RM×d , Y s ∈ RN×d . The Equation 4 summarizes the sequencesXat and Y as which contain the contextual and attention information . We take the last tokens from both the xs ∈ Rd and ys ∈ Rd , since the last token can see the whole sequence and acts as a summary of the sequence . The task embedding layer computes te by taking the outputs of the sequence summary sub-layer and applying a feed-forward network : te = FFN ( [ x s : ys ] ) . ( 5 ) 2.2.2 MODULATING MODEL PARAMETERS . We modulate the parameters ( θ ) of the network ( Section 2.1 ) to account for the task-specific variation during MTL over a set of tasks . We achieve this by scaling ( γ ) and shifting ( β ) the outputs of each layer ( e.g. , transformer block ) including any preprocessing layers in the model adopted based on the Feature-wise Linear Modulation ( FiLM ; Perez et al . ( 2018 ) ) . The γ and β parameters are obtained from the task embedding te either by using equation 1 or 5. γ = te [ : d ] , β = te [ d : ] , ( 6 ) where te ∈ R2d , and d is the hidden dimension of the model . Once we have γ and β , we apply the feature-wise linear modulation ( Perez et al. , 2018 ) to compute the modulated output ( Ol ) for each block of the model . Ol = γ ∗ fl ( vl ; θl ) + β , l = 1 , · · · , L , ( 7 ) where L is the total number of blocks in the model and fl represents the lth block of the model with parameters θl ∈ θ and inputs vl . 2.3 TRAINING . MTL has been successfully applied across different applications of machine learning such as natural language processing ( Hashimoto et al. , 2016 ; Collobert & Weston , 2008 ) , speech recognition ( Liu et al. , 2019 ; Deng et al. , 2013 ) , computer vision ( Zhang et al. , 2014 ; Liu et al. , 2015 ; Girshick , 2015 ) , and drug discovery ( Ramsundar et al. , 2015 ) . It comes in many forms : joint learning , learning to learn , and learning with auxiliary tasks . We consider two MTL strategies : ( 1 ) joint learning and ( 2 ) learning to learn to train on set of S tasks , T = { τ1 , · · · , τS } with corresponding datasets D = { D1 , · · · , DS } . As our first training strategy , we use Joint Learning ( JL ) ( Caruana , 1997 ) , which is the most commonly used training strategy for MTL . In JL , the model parameters , including the output layer , are shared across all the tasks involved in the training . For the second training strategy under the learning-tolearn approach , we use a variant of meta-learning , Modality Agnostic Meta Learning ( MAML ) ( Finn et al. , 2017a ) . Even though MAML is mostly used in few-shot learning settings , we use it since it allows for task-specific learning during the meta-train step and it has also been shown to provide improvements in the field of speech translation ( Indurthi et al. , 2020 ) . We resolve the source-target vocabulary mismatch across different tasks in MTL by using a vocabulary of subwords ( Sennrich et al. , 2016 ) computed from all the tasks . We sample a batch of examples from Ds and use this as input to the TCN and the Transformer model . To ensure that each training example uses the task embedding computed using another example , we randomly shuffle this batch while using them as input to the TCN . This random shuffling improves the generalization performance by forcing the network to learn task-specific characteristics ( te ) in Equation 1 or 5 . We compute the task embedding in the meta-train step as well ; however , the parameters of the TCN are updated only during the meta-test step . During inference time , we use the precomputed task embeddings using a batch of examples randomly sampled from the training set .

Summary: This paper proposes a new framework that computes the task-specific representations to modulate the model parameters during the multi-task learning (MTL). This framework uses a single model with shared representations for learning multiple tasks together. Also, explicit task information may not be always available, in such cases the proposed framework is useful. The proposed framework is evaluated on various datasets spanning multiple modalities, where the MTL model even achieves state-of-the-art results on some datasets.