Independent Component Alignment for Multi-task Learning

1 INTRODUCTION . In multi-task learning ( MTL , e.g. , Caruana , 1993 ; Doersch & Zisserman , 2017 ) , several tasks are optimized simultaneously using a single model and by leveraging shared information across tasks to improve generalization and boost performance for all tasks . MTL is a key component in real-world applications , such as natural language processing ( Liu et al. , 2019c ) , computer vision ( Kendall et al. , 2015 ; Kokkinos , 2017 ) , and reinforcement learning ( Parisotto et al. , 2016 ; Teh et al. , 2017 ) . However , these systems are difficult to train since different tasks need to be properly balanced , which is challenging when the multi-task gradient is dominated by the gradient of one of the tasks . MTL approaches have demonstrated rapid progress , and in general , they can be divided into two categories . The first group of methods aims at designing multi-task neural network architectures which should learn both task-shared and task-specific features by using hard and soft parameter sharing ( Liu et al. , 2019b ; Lu et al . ; Houlsby et al. , 2019 ; Pfeiffer et al. , 2021 ) . The methods in the second group focus on multi-task optimization strategies . To learn all tasks with equal importance , some methods use gradient replacement ( Yu et al. , 2020 ; Lopez-Paz & Ranzato , 2017 ) or balance the individual loss functions with weights sampled from a uniform distribution , learned by using task-dependant uncertainty ( Kendall et al. , 2018 ) , determined by utilising the learning speed for the task ( Liu et al. , 2019a ; b ; Lee & Son , 2020 ) , or obtained by minimizing an auxiliary loss ( Chen et al. , 2018 ) . However , combining individual loss functions into one using a weighted average can be hard and prone to errors if some task gradients are in conflict with each other . To alleviate this limitation , a variety of multi-objective optimization methods have been proposed ( Désidéri , 2012 ; Sener & Koltun , 2018 ) . They aim to find Pareto-stationary solutions for which the performance of any of the objectives can only be improved by deteriorating performance of some other objective . Although these methods demonstrate good performance , their main limitation is local convergence . Specifically , if the objective function is non-convex , which is typical in deep learning , the methods converge to a poor Pareto-stationary solution . Moreover , there is a gap between optimization and generalization performance as proven by Lin et al . ( 2019 ) . This work hinges on the realization that improving upon the overall performance in multi-task learning can be formulated by conservatism in the individual learning steps of uncorrelated losses . We formulate an approach that aims to guarantee consistent improvement over dynamically constructed linearly independent tasks at every step . This results in a multi-objective optimization algorithm that automatically balances training in multi-task systems by aligning orthogonal components of the objective . We address the problem of a large difference in the task gradient magnitudes through the condition number of the gradient matrix . This number defines the dominance rate in the multi- objective task and is obtained as the ratio of maximum and minimum singular values . Fig . 1 illustrates the proposed approach and shows its performance compared to other gradient-based MTL methods when the difference in the task gradient magnitudes is large . The contributions of this paper are summarized as follows . ( i ) We present a novel gradient-based MTL approach that dynamically constructs consistent losses and seamlessly handles the problem of a large difference in task gradients . ( ii ) We provide a computationally efficient approximation for the approach with constant time complexity and which is directly applicable to large-scale problems . ( iii ) We formulate a new metric , the dominance rate , as an interpretable way of establishing the maximum possible dominance in the gradient system . Additionally , we perform an extensive evaluation of the proposed approach on various MTL benchmarks and demonstrate that our method outperforms prior state-of-the-art MTL methods . 2 RELATED WORK . A multi-task formulation is well motivated and ubiquitous to many application domains such as computer vision ( Bilen & Vedaldi , 2016 ; Kokkinos , 2017 ; Zamir et al. , 2018 ; Nekrasov et al. , 2019 ; Kendall et al. , 2018 ) , natural language processing ( Collobert & Weston , 2008 ; Luong et al. , 2015 ; Dong et al. , 2015 ) , speech processing ( Seltzer & Droppo , 2013 ) , and robotics ( Wulfmeier et al. , 2020 ; Levine et al. , 2016 ) . In this work we focus on gradient-based malti-task learning strategies and refer the interested reader to comprehensive surveys by Caruana ( 1993 ) , Ruder ( 2017 ) , and Crawshaw ( 2020 ) for a more detailed overview . One way to ease multi-task optimization is to balance the individual loss functions for different tasks . Prior methods to simultaneously learning multiple tasks parameterize the total objective as a weighted sum of task-specific loss functions whose weighting coefficients are manually tuned ( Kendall et al. , 2015 ; Laskar et al. , 2017 ) . However , using grid search for an optimal weighting is computationally inefficient especially when the number of tasks is large . To address this limitation , Kendall et al . ( 2018 ) utilize homoscedastic uncertainty of each task to derive a weighted multi-task loss . Chen et al . ( 2018 ) propose a heuristic to tune loss weights based on task gradient magnitudes . One of the main challenges in MTL is the presence of conflicting gradients when tasks have gradients pointing in opposing directions , i.e . negative transfer ( Lee et al. , 2018 ) . Among other approaches to mitigate the problem of negative transfer , the methods based on explicit gradient modulation have shown great potential ( Lopez-Paz & Ranzato , 2017 ; Yu et al. , 2020 ; Chaudhry et al. , 2019 ) . The core idea is to replace a task gradient which conflicts with some other task by a modified , non-conflicting , gradient . Yu et al . ( 2020 ) propose a ‘ gradient surgery ’ technique which makes the new gradient non-conflicting with the average of task gradients at the previous step . However , it requires checking multiple gradients for conflicts at each update step leading to a large computational overhead . On the other hand , multi-objective optimization ( MOO ) is the process of optimizing a set of possibly contrasting objectives simultaneously ( Fliege & Svaiter , 2000 ; Schäffler et al. , 2002 ; Désidéri , 2012 ; Peitz & Dellnitz , 2018 ; Poirion et al. , 2017 ) . These methods find a direction that decreases all objectives relying on multi-objective Karush–Kuhn–Tucker ( KKT ) conditions ( Kuhn & Tucker , 1951 ) . The work by Sener & Koltun ( 2018 ) extends the classical MGDA ( Désidéri , 2012 ) method to a form that scales well to high-dimensional problems . However , it converges to a non-optimal Pareto-stationary solution if the objective function is non-convex . We empirically compare our approach to a number of MOO methods and demonstrate its superior performance . 3 PROBLEM FORMULATION AND NOTATION . Multi-task learning is concerned with finding the set of parameters for a parametric model that correspond to a high average performance across all tasks . More formally , we consider a MTL problem defined over the input space X and a series of targets { Y1 , . . . , YT } , where T is the number of tasks . We consider a parametric model f ( x ; θsh , θ1 , . . . , θT ) with shared , θsh , and task-specific parameters , θi . Each task has a corresponding differentiable loss function Li ( ŷi , yi ) : Yi × Yi → R. Moreover , we can include prior knowledge on user preferences p ( Y ) in the model . The MTL problem can be formulated as an empirical risk minimisation ( ERM ) problem on a data set { ( xj , y1j , . . . , yTj ) } Nj=1 of i.i.d . data points : min θsh θ1 , ... , θT EYi∼p ( Y ) [ Li ( θsh , θi ) ] = min θsh θ1 , ... , θT T∑ i=1 wiLi ( θsh , θi ) = min θsh θ1 , ... , θT T∑ i=1 wiLi ( θ = θsh ) . ( 1 ) The vector of preferences w defines the importance of each task : w = 1 for MTL ; w 6= 1 for auxiliary task learning . We consider shared parameters , unless otherwise specified . Since the gradients of θi and θsh are linearly independent pairwise , task-specific parameters have a single gradient , and thus can be omitted without loss of generality . 4 INDEPENDENT COMPONENT ALIGNMENT . One of the key challenges in MTL is overfitting to one of the tasks caused by the dominance of the gradient of some linear component in the loss function . In practice , the overfitting might be defined by a dominating direction which has the fastest learning rate compared to others . This dominance can be caused by , for example , gradient interference between the individual loss functions or a significant difference in their magnitudes . Such a situation is illustrated in Fig . 2a , where the gradient for task 2 dominates task 1 . Recent works on the interactions of individual tasks ( Zamir et al. , 2018 ; 2020 ) demonstrate that tasks may have a common component . This means that the uniqueness of a particular task is determined by another , unique , component independent to the common one . In this regard , the interference of gradients leads to a faster training rate of the joint subtask reducing the efficiency of learning unique components for each task . Therefore , we need to analyze the independent components of the gradient system , rather than the individual gradients themselves . We define dominance in the system of gradients as follows : Definition 1 ( Dominance rate ) . The rate of dominance of the system of loss functions is the condition number of the matrix composed of the gradients of these functions over the shared parameters : D ( L1 ( θ ) , . . . , LT ( θ ) ) = D ( Gθ ) = σmax ( Gθ ) σmin ( Gθ ) , where Gθ = ( ∇θL1 ( θ ) , . . . , ∇θLT ( θ ) ) . ( 2 ) This definition is based on the singular value decomposition ( SVD ) and has a geometric interpretation . The dominance coefficient defines the maximum possible dominance in the gradient system and is obtained as the ratio of the maximum and minimum learning rates of the orthonormal components of the gradient system ( as in Fig . 2a ) . Recently , Wang et al . ( 2020 ) propose a metric , OGR to quantitatively estimate the significance of overfitting of multi-modal networks for classification . Although related , the metric is obtained by using an assumption that target distribution is well approximated by a validation data set . In contrast , our measure does not require such a strong prior and solely based on gradients with respect to the training set . 4.1 MODEL-AGNOSTIC ALIGNMENT . For effective training over all tasks , the system of gradients of loss functions should be consistent and have a small dominance rate . We take this as a design principle , and aim for avoiding domination and inconsistency altogether . We propose constructing a new ( consistent ) gradient system that best ( under conditions that follow ) describes the given directions , but whose dominance coefficient is equal to minimum dominance over the task gradients , i.e . unit dominance . More formally , we define a new aligned gradient as follows . Definition 2 ( θ-aligned ) . We define a θ-aligned gradient as ∇θaθ L ( θ ) : = σmin ( Gθ ) Ĝθw , where vector w denotes a pre-selected task preference and Ĝθ is the solution to the Procrustes problem : Ĝθ = arg min Q ‖Gθ −Q‖F , s.t . Q > Q = I . ( 3 ) In order to calculate the θ-aligned gradient , the solution to the general Procrustes problem is needed . Schönemann ( 1966 ) shows that the problem has an analytical solution , namely : Ĝθ = UV > , where U , V : Gθ = UΣV > ( given by a SVD ) . ( 4 ) The analytical solution allows us to avoid the costly procedure for finding the minimum , which reduces the computational cost of our approach . Additionally , this leads to the dominance rate of the resulting matrix being equal to unit dominance : D ( σmin ( Gθ ) Ĝθ ) = σminσmin = 1 , i.e . we have eliminated the possibility of uneven learning of different tasks . Fig . 2c shows a sketch of the geometric intuition of the aligned gradients after the alignment step . At this stage we aim to verify that standard gradient descent with θ-aligned gradients corresponds to a sensible optimization procedure . We provide the following analysis , which states the convergence for general non-convex loss functions with standard assumptions on regularity . For a proof and more theoretical results , see App . B. Theorem 1 . Assume L1 ( θ ) , . . . , LT ( θ ) are lower bounded continuously differentiable functions with Lipschitz continuous gradients with Λ > 0 . Assume that the total loss function L ( θ ) is a weighted sum of these functions with a pre-selected vector of preferences w. Gradient descent with step size α ≤ 1Λ and θ-aligned gradient converges if the initial rate of dominance is finite , D ( Gθ ) < C , and task preferences vector w is consistent with the gradient system , ‖V w‖ ≥ ε .

Summary: This paper proposes a gradient-based multi-task learning method to balance multi-task training by aligning the independent components of the training objective. Experiments conducted on several MTL problems are given. The method is claimed to be scalable, robust to overfitting, and able to seamlessly handle multi-task objectives with a large difference in gradient magnitudes.