AUXILIARY TASK UPDATE DECOMPOSITION: THE GOOD, THE BAD AND THE NEUTRAL

1 INTRODUCTION . Multitask learning ( Caruana , 1997 ) and pretraining ( Devlin et al. , 2018 ; Caron et al. , 2019 ) have transformed machine learning by allowing downstream tasks with small training sets to benefit from statistical regularities from data-rich related tasks ( Collobert & Weston , 2008 ; Zhang et al. , 2014 ; Liu et al. , 2019 ; Kornblith et al. , 2019 ) . Despite these advances , leveraging the mixing of tasks is still an art left to the practitioner . When one is interested in a primary task , it is unclear how to select helpful auxiliary tasks , an appropriate parameter sharing architecture and a good way to filter out auxiliary data which might be detrimental to the primary tasks . Without careful choices , pretraining might hurt end-task performance ( Gururangan et al. , 2020 ) or have limited impact ( Raghu et al. , 2019 ) . Prior work has examined these problems and proposed solutions , either to choose auxiliary tasks depending on their impact on the primary task ( Du et al. , 2018 ; Lin et al. , 2019 ) or to equalize the impact of updates across tasks ( Sener & Koltun , 2018 ; Chen et al. , 2018 ; Hessel et al. , 2019 ) . Recently , several approaches ( Sinha et al. , 2018 ; Suteu & Guo , 2019 ; Yu et al. , 2020 ) have been proposed that attempt to minimize interference between the updates across tasks . Our work builds on this direction , but unlike these previous approaches , we do not consider a symmetric view of multi-task learning in the sense that our goal is not to train a model performing well on all tasks . Instead , we focus on improving generalization for a single task , the primary task , and the other tasks , the auxiliary tasks are considered only through their impact on the problem of interest . For that purpose , we introduce a framework which decomposes the gradient updates from the auxiliary tasks according to their impact on the primary task . We analyze the auxiliary task gradients in the subspace spanned by the primary task per-example gradients . This allows us to decompose auxiliary gradients into into three components : components that help , interfere or have no impact on the primary task according to the Taylor expansion of the expected primary loss . This decomposition allows us to re-weight each component differently prior to the update . Our framework enables us to treat each auxiliary update differently depending on its impact on the task of interest and it encompasses prior methods such as classical multitask learning ( Caruana , 1997 ) or more novel gradient surgery techniques ( Yu et al. , 2020 ) . To achieve a tractable approach , we introduce an efficient , robust algorithm ( ATTITTUD , Auxiliary Task Training with Influence from Target Task Update Direction ) to estimate the subspace spanned by the primary task gradients in an online manner and decompose the auxiliary updates appropriately . As a result , we can integrate our approach with the stochastic training of large neural networks in various contexts . The contribution of our work is four-fold . To our knowledge , this paper proposes the first approach to adapt auxiliary gradients using a decomposition built from the span of the primary task Jacobian . In order to scale this approach to deep neural nets , we contribute a tractable and efficient algorithm called ATTITTUD that leverages insights from randomized linear algebra and automatic differentiation such as the R-operator ( Pearlmutter , 1994 ) . As our third contribution , we show that the fine-grained manipulation of the auxiliary task gradients under ATTITTUD , represents a unified framework that encompasses several previous approaches to asymmetrical task learning as special cases . Finally , we demonstrate the efficacy of our approach in both data-rich and data-starved primary tasks , over both images and textual data . 2 RELATED WORK . Methods to leverage data outside of the task of interest have been popular in machine learning since the inception of multitask learning ( Caruana , 1997 ; Ruder , 2017 ; Vandenhende et al. , 2020 ) . These methods address multiple task simultaneously and have been successful in various application domains ( Collobert & Weston , 2008 ; Zhang et al. , 2014 ; Misra et al. , 2016 ) . The optimization problem induced by multitask learning is difficult and solutions have been proposed for the various difficulties , including dealing with task gradients of different magnitude ( Sener & Koltun , 2018 ; Chen et al. , 2018 ; Hessel et al. , 2019 ) , or interfering with each others ( Sinha et al. , 2018 ; Suteu & Guo , 2019 ; Yu et al. , 2020 ) . The specific problem of interference has been studied extensively in the context of continual learning . Continual learning visits task in sequence and update interference is particularly problematic as it yields newer tasks to damage previously mastered tasks . In particular , a family of methods to project the gradient of the new tasks to be orthogonal to the gradient of the previous tasks has been proposed ( Lopez-Paz & Ranzato , 2017 ; Chaudhry et al. , 2018 ; Farajtabar et al. , 2019 ) . Different from many previous approaches , we are not interested in addressing multiple tasks per se . In our setting , only the primary task matters and the other auxiliary task have the sole role of improving generalization on the primary task . This is the setting considered by Du et al . ( 2018 ) ; Lin et al . ( 2019 ) , who favor auxiliary tasks whose gradient directions are helpful to the primary task . Unlike these works that use coarse properties like the cosine similarity between averaged gradients , our approach allows fine-grained gradient manipulation within a subspace . Also , in our case , we do not distinguish between the different auxiliary tasks . Instead , we aim at correcting every auxiliary gradient in the same manner to improve the loss on the primary task . This type of gradient correction is related to Yu et al . ( 2020 ) , which considers projecting multi-task gradients such that the directions of disagreement are removed . This method is actually a special case of our framework . Our work also shares some similarities with data selection and domain adaptation approaches . In this case , the training data comes from a single task but its distribution is different from the validation/test distribution ( Moore & Lewis , 2010 ; Axelrod et al. , 2011 ; Ngiam et al. , 2018 ) . This classical problem has recently been addressed by sampling training points whose gradient aligns well with the expected validation gradient ( Wang et al. , 2020b ; a ) . Instead of sampling individual points based on an estimated distribution of how helpful they will be to the primary task , our work avoids the use ( and inherent challenges ) of this reinforcement learning approach by operating on batch gradients of groups of points . Our primary task/auxiliary task setting is also related to the pre-training then fine-tuning paradigm in which the auxiliary tasks are visited first ( pre-training ) to give an initialization for training on the primary task ( fine-tuning ) . These methods have been very successful in settings where primary task data are rare . In particular , it is common to first rely on an unsupervised task over very large datasets prior to fine tuning over a supervised task ( Devlin et al. , 2018 ; Liu et al. , 2019 ; Kornblith et al. , 2019 ; Yang et al. , 2019 ; Song et al. , 2019 ; Caron et al. , 2018 ) . 3 AUXILIARY TASK UPDATE DECOMPOSITION . This section introduces a new method to improve generalization on a primary task T ∗ using training data from auxiliary tasks T = { T1 , . . . , Tn } , where θ ∈ RD denote the parameters shared by all tasks . Our approach leverages gradient updates from the auxiliary tasks , but unlike the traditional approach , we decompose these gradients to maximize their usefulness to T ∗ . Precisely , we decompose the auxiliary task gradients into directions which decrease a first-order approximation of the primary task loss , increase it or have no effect . This decomposition allows weighting these three directions differently when learning from the auxiliary tasks . In order to decompose the auxiliary gradient , we must collect more fine-grained statistics about the primary task . At each training step , we collect the gradient of the loss with respect to θ for individual examples from the primary task , { ∇θLprimi , ∀i } . The span of these vectors , S = Span { ∇θLprimi , ∀i } defines a subspace in which any linear combination of primary task gradients lies , including the gradient of the expected primary task loss , i.e . gprim = E ( ∇θLprimi ) ∈ S. We denote the size of the subspace , |S| = K. This is upper-bounded by the number of examples m , used to construct S. If we define the orthogonal complement of S as S⊥ , any vector v ∈ S⊥ , is therefore orthogonal to gprim , i.e . v · gprim = 0 . This means that adding such a vector to the parameters has no impact on the expected primary task loss , according the order-1 Taylor expansion of Lprim , i.e . Lprim ( θ + v ) ' Lprim ( θ ) + v · gprim = Lprim ( θ ) . We propose to project auxiliary task gradients onto S and S⊥ . This allow to distinguish between the directions of the auxiliary task updates which impact the primary task loss and those which do not . If we denote the averaged auxiliary task gradient as gaux = E ( ∇θLauxi ) , we can decompose this gradient as gaux = g±aux + g ⊥ aux . where g ± aux ∈ S is the portion of the gradient that lies in the span of the primary task example gradients and g⊥aux ∈ S⊥ is the portion that lies outside of it . Since g⊥aux ∈ S⊥ , it is orthogonal to the average primary task gradient and parameter updates along the direction of g⊥aux are expected to have limited impact on the primary task loss . On the other hand , updates along the direction of g±aux can potentially improve or damage the averaged primary task loss . This component deserves a more careful treatment . For that purpose , we introduce { ui , i = 1 , . . . , K } an orthonormal basis of S. In this basis , we can measure if the components of g±aux agree or disagree with gprim . We say that the two gradients agree along ui iif sign ( g±aux · ui ) = sign ( gprim · ui ) . This means that we can decompose g±aux = g+aux + g−aux where g+aux refers to the projection of g ± aux onto the basis vectors where g ± aux and gprim agree . By this decomposition , g+aux helps the primary task , g + aux · gprim > 0 , while g−aux interfere with the primary task , g−aux · gprim < 0 . Guided by the primary task , we can therefore decompose the auxiliary task gradient as gaux = g ⊥ aux + g + aux + g − aux ( 1 ) which is described on Fig 1 . Our approach proposes to re-weight differently the components of gaux , i.e . g̃aux = η⊥g ⊥ aux + η+g + aux + η−g − aux ( 2 ) where ηaux = ( η⊥ , η+ , η− ) are hyper-parameters adjusting the auxiliary gradient according to the impact on the main task . If we also wish to include the primary task gradient in descent , as with multitasking , we can introduce ηprim as a scalar control variable to modulate its weighting . A consequence of introducing ηaux is that specific configurations lead us to gradient updates that are guaranteed to do no harm to both tasks . This is captured by Theorem 1 below . Theorem 1 . Let Laux ( θt ) and Lprim ( θt ) represent the full batch losses of the auxiliary tasks and primary task respectively at step t. We assume the gradients of Laux and Lprim are Lipschitz continuous with constant L > 0 . Following the update rule : θt+1 = θt − α · g̃aux , where α ≤ 1L is the learning rate , we are guaranteed : Laux ( θt+1 ) ≤ Laux ( θt ) Lprim ( θt+1 ) ≤ Lprim ( θt )

Summary: The work studies the auxiliary task selection in deep learning to resolve the burden of selecting relevant tasks for pre-training or the multitask learning. By decomposing the auxiliary updates, one can reweight separately the beneficial and harmful directions so that the net contribution to the update of the primary task is always positive. The efficient implementation is experimented in text classification, image classification, and medical imaging transfer tasks.