A Closer Look at Loss Weighting in Multi-Task Learning

1 INTRODUCTION . Multi-Task Learning ( MTL ) ( Zhang & Yang , 2021 ; Vandenhende et al. , 2021 ) aims to jointly train several related tasks to improve their generalization performance by leveraging common knowledge among them . Since this learning paradigm can not only significantly reduce the model size and increase the inference speed but also improve the performance , it has been successfully applied in various fields of deep learning , such as Computer Vision ( CV ) ( Vandenhende et al. , 2021 ) , Natural Language Processing ( NLP ) ( Chen et al. , 2021 ) , reinforcement learning ( Zhang & Yang , 2021 ) and so on . However , when all the tasks are not related enough , which may be reflected via conflicting gradients or dominating gradients ( Yu et al. , 2020 ) , it is more difficult to train a multi-task model than training them separately because some tasks dominantly influence model parameters , leading to unsatisfactory performance for other tasks . This phenomenon is related to the task balancing problem ( Vandenhende et al. , 2021 ) in MTL . Recently , several works focus on tackling this issue from an optimization perspective via balancing task losses or gradients . In this paper , we investigate eight State-Of-The-Art ( SOTA ) task balancing approaches and unify them as loss weighting strategies . According to the way of generating loss weights , those methods can be divided into three types , including the solving approach such as directly solving a quadratic optimization problem in a multi-objective formulation as weights ( Sener & Koltun , 2018 ) , the calculating approach such as projecting conflict gradients ( Yu et al. , 2020 ) , and the learning approach such as learning weights in a gradient descent manner ( Chen et al. , 2018b ) . On the other hand , since there are some discrepancies of the implementation details such as using different backbone networks for training or different metrics for the evaluation among those SOTA weighting methods , leading to inconsistent comparisons , we provide a unified testbed on six CV datasets and four multilingual problems from the XTREME benchmark ( Hu et al. , 2020 ) for those SOTA weighting strategies to show a fair comparison . In addition , inspired by dynamic weighting processes in those SOTA strategies where loss weights vary over training iterations or epochs , we have a sudden whim : what will happen if a MTL model is trained with random loss weights ? Specifically , in each training iteration , we first sample the loss weights from a distribution with some normalization and then minimize the aggregated loss weighted by the normalized random weights . Surprisingly , this seemingly unreliable method can not only converge but also achieve comparable performance with the SOTA weighting strategies . Based on this observations , we proposed a simple yet effective weighting strategy for MTL , called Random Loss Weighting ( RLW ) . It is very easy to implement RLW by adding only one line of code and this strategy does not incur any additionally computational cost . An implementation example of RLW in PyTorch ( Paszke et al. , 2019 ) is shown below . 1 outputs = model ( inputs ) 2 loss = criterion ( outputs , labels ) # [ 1 , task_num ] vector 3 weight = F.softmax ( torch.randn ( task_num ) , dim=-1 ) # RLW is only this ! 4 loss = torch.sum ( loss * weight ) 5 optimizer.zero_grad ( ) 6 loss.backward ( ) 7 optimizer.step ( ) To show the effectiveness of RLW , we provide both theoretical analyses and empirical evaluations . Firstly , the objective function of RLW can be considered as a doubly stochastic optimization problem when optimizing by stochastic gradient descent or its variants , where the randomness is from both the mini-batch sampling of the data for each task and the random sampling of loss weights . From this perspective , we give a convergence analysis for RLW . Besides , we can show that RLW has a higher probability to escape local minima when compared with fixing loss weights , resulting in a better generalization performance . Empirically , as described before , we compare RLW with SOTA weighting approaches on six CV datasets and four multilingual problems to show its competitive performance . In summary , the main contributions of this paper are four-fold . • We provide a unified testbed on six multi-task computer vision datasets and four multilingual problems from the XTREME benchmark for a fair comparison among eight SOTA weighting methods and the proposed RLW method . • We propose a simple yet effective RLW strategy , which we think is an ignored baseline in MTL . • We provide the convergence guarantee and effectiveness analysis for RLW . • Experiments show that RLW can achieve comparable performance with SOTA weighting methods without bringing any additionally computational cost . 2 PRELIMINARY . Suppose there are T tasks and task t has its corresponding dataset Dt . An MTL model usually contains two parts of parameters : task-sharing parameters θ and task-specific parameters { ψt } Tt=1 . For example , in CV , θ usually denotes parameters in the feature extractor shared by all tasks and ψt represents the task-specific output module for task t. Let ` t ( · ; θ , ψt ) denotes a task-specific loss function for task t. Then the objective function of a MTL model can be formulated as LMTL = T∑ t=1 λt ` t ( Dt ; θ , ψt ) , ( 1 ) where ` t ( Dt ; θ , ψt ) denotes the average loss on Dt for task t and { λt } Tt=1 are task-specific loss weights with a constraint that λt ≥ 0 for all t. When minimizing Eq . ( 1 ) by Stochastic Gradient Descent ( SGD ) or its variants , the task-specific parameters { ψt } Tt=1 are simply updated based on the corresponding task gradient ∇ψt ` t ( Dt ; θ , ψt ) , while the task-sharing parameters θ should be updated by all the task losses jointly as θ = θ − η T∑ t=1 λt∇θ ` t ( Dt ; θ , ψt ) , ( 2 ) where η is a learning rate . Obviously , for the update of task-sharing parameters θ , the loss weighting ( i.e. , { λt } Tt=1 in Eq . ( 1 ) ) influences θ via the aggregated gradient essentially and the gradient weighting in Eq . ( 2 ) during the backward process has the same effect as the loss weighting when they are using the same weights . Therefore , we can ignore the level on which the weights act and focus on the generation of weights . For simplicity , these two types of weighting are all referred to as loss weighting in the following sections . Apparently , the most simple method for loss weighting is to set a same weight for every tasks , i.e. , without loss of generality , λt = 1T for all t. This approach is a common baseline in MTL and it is called Equally Weighting ( EW ) in this paper . To tackle the task balancing problem and improve the performance of MTL model , there are several works to study how to generate appropriate weights . In this paper , we investigate eight SOTA weighting strategies , i.e . Gradient Normalization ( GradNorm ) ( Chen et al. , 2018b ) , Uncertainty Weights ( UW ) ( Kendall et al. , 2018 ) , MGDA ( Sener & Koltun , 2018 ) , Dynamic Weight Average ( DWA ) ( Liu et al. , 2019a ) , Projecting Conflicting Gradient ( PCGrad ) ( Yu et al. , 2020 ) , Gradient sign Dropout ( GradDrop ) ( Chen et al. , 2020 ) , Impartial Multi-Task Learning ( IMTL ) ( Liu et al. , 2021 ) , and Gradient Vaccine ( GradVac ) ( Wang et al. , 2021 ) . According to different ways of generating loss weights , we categorize those loss weighting strategies into three types : the learning approach , the solving approach , and the calculating approach . Both GradNorm and UW consider the loss weights { λt } Tt=1 in Eq . ( 1 ) as learnable parameters and explicitly optimize them by gradient descent . MGDA casts MTL as a multi-objective optimization problem and directly solves the loss weights { λt } Tt=1 in Eq . ( 1 ) by solving a quadratic programming problem . DWA , PCGrad , GradDrop and GradVac directly compute the weights { λt } Tt=1 by combining gradients and/or losses of all the tasks . IMTL is a hybrid strategy , which combines the learning and the calculating approaches . We summarize those strategies from the perspective of loss weighting in Table 5 in Appendix A . We now unify those eight SOTA methods as loss weighting strategies , i.e. , generating loss weights { λt } Tt=1 in Eq . ( 1 ) . Noticeably , almost all the existing strategies except EW need to incur intensive computation to generate loss weights in every iteration , such as solving a quadratic optimization problem in MGDA , and operating on high-dimensional gradients in PCGrad , GradDrop , IMTL , and GradVac . Different from those strategies , the proposed RLW strategy generates loss weights in a sampling way , thus it is as efficient as EW without bringing additionally computational costs . 3 THE RLW METHOD . In this section , we introduce the proposed RLW method . The RLW method is a simple loss weighting strategy and it considers the loss weights λ = ( λ1 , · · · , λT ) ∈ RT as random variables . Formally , the objective function of the RLW method is formulated as LRLW ( θ ) = Eλ [ λ > ` ( D ; θ ) ] = Eλ [ λ ] > ` ( D ; θ ) , ( 3 ) where E [ · ] denotes the expectation and ` ( D ; θ ) = ( ` 1 ( D1 ; θ ) , · · · , ` T ( DT ; θ ) ) where we omit the task-specific parameters { ψt } Tt=1 in Eq . ( 3 ) for brevity . To guarantee loss weights in λ to be nonnegative , we can first sample λ̃ = ( λ̃1 , · · · , λ̃T ) from any distribution p ( λ̃ ) and then normalize λ̃ into λ via a mapping f , where f : RT → ∆T is a normalization function for example softmax function and ∆T denotes a convex hull in RT , i.e . λ ∈ ∆T means ∑T t=1 λt = 1 and λt ≥ 0 for all t. Note that in most cases p ( λ ) is different from p ( λ̃ ) . In Eq . ( 3 ) , p ( λ ) is usually too complex to compute its expectation Eλ [ λ ] , thus a stochastic approximation scheme is adopted to minimize Eq . ( 3 ) . When the mini-batch SGD ( Bottou , 1991 ) or its variants is used to minimize Eq . ( 3 ) as most deep learning models did , Eq . ( 3 ) can be viewed as a doubly stochastic optimization problem , where the randomness is from both the mini-batch data sampling for each task and the randomly sampling of the loss weights . In the following , we show that the approximated gradient ∇θλ > ` ( D̃ ; θ ) is an unbiased estimation of the true gradient of LRLW ( θ ) , where D̃ denotes a mini-batch data sampled from all the tasks . Specifically , as D̃t is a mini-batch data sampled from Dt to calculate the stochastic gradient ∇θ ` t ( D̃t ; θ ) to approximate the full gradient ∇θ ` t ( Dt ; θ ) for task t , we have ED̃ [ ∇θ ` ( D̃ ; θ ) ] = ∇θ ` ( D ; θ ) . Therefore , when we further randomly sample a weight vector λ , we have Eλ [ ED̃ [ ∇θλ > ` ( D̃ ; θ ) ] ] = Eλ [ λ ] > ∇θ ` ( D ; θ ) = ∇θLRLW ( θ ) , which verifies that∇θλ > ` ( D̃ ; θ ) is an unbiased estimation . In practice , it is very easy to implement the RLW method without modifying network architecture or bringing additionally computational costs . Specifically , in each iteration , we first sample λ̃ from p ( λ̃ ) and normalize it to obtain λ via appropriate normalization function f , and then minimize the aggregated loss weighted by λ . The entire algorithm of RLW ( i.e. , minimizing Eq . ( 3 ) ) via SGD is shown in Algorithm 1 . Apparently , the only difference between the proposed RLW strategy and the widely used EW strategy is Line 7 in Algorithm 1 and it is very easy to implement with only one line of code . Algorithm 1 Optimization Algorithm for RLW by SGD Input : numbers of iterations K , numbers of tasks T , learning rate η , dataset { Dt } Tt=1 , weight distribution p ( λ̃ ) 1 : Randomly initialized θ0 ; 2 : for k = 1 to K do 3 : for t = 1 to T do 4 : Sample a mini-batch data D̃t from Dt ; 5 : Compute loss ` t ( D̃t ; θk ) ; 6 : end for 7 : Sample weights λ̃ from p ( λ̃ ) and normalize it into λ via f ; . RLW is only this 8 : θk+1 = θk − η∇θ ∑T t=1 λt ` t ( D̃t ; θk ) ; 9 : end for In this paper , we use six different distributions for p ( λ̃ ) in the proposed RLW method , including uniform distribution between 0 and 1 ( denoted by Uniform ) , standard normal distribution ( denoted by Normal ) , Dirichlet distribution with α = 1 ( denoted by Dirichlet ) , Bernoulli distribution with probability 1/2 ( denoted by Bernoulli ) , Bernoulli distribution with probability 1/2 and a constraint ∑T t=1 λ̃t = 1 ( denoted by constrained Bernoulli ) , and normal distribution with a random mean and a random variance both sampling from a uniform distribution U ( 0 , 1 ) for each task ( denoted by random Normal ) . We set f as a function of f ( λ̃ ) = λ̃/ ( ∑T t=1 λ̃t ) if p ( λ̃ ) is the Bernoulli distribution or the constrained Bernoulli distribution and a softmax function for the other types of distribution . When sampling from the first five types of distribution , E [ λ ] is simply proportional to ( 1T , · · · , 1 T ) , thus it is fair to compare with the EW strategy . When p ( λ̃ ) is a random Normal distribution , it means each λ̃t is sampled from a normal distribution with random mean and variance , thus it is intractable to compute the expectation for p ( λ ) and combining with such distribution can further show the effectiveness of RLW . When sampling from a Bernoulli distribution , the weights for all tasks are either 0 or 1 , i.e. , λt ∈ { 0 , 1 } for all t. In this way , just a subset of tasks contributes to updating the task-sharing parameters θ . This manner can be viewed as the mini-batch sampling on the task level . If considering an additional constraint that ∑T t=1 λ̃t = 1 , it implies only one task is involved in the update of the tasksharing parameters in each iteration . Although there are some works ( Dong et al. , 2015 ; Liu et al. , 2015a ; Søgaard & Goldberg , 2016 ; Subramanian et al. , 2018 ; Sanh et al. , 2019 ; Liu et al. , 2019b ) adopting this strategy to train a MTL model , it is a special case in the proposed RLW strategy and beyond existing works , we also provide theoretical analyses to show the effectiveness of the proposed RLW method .

Summary: This paper conducted an empirical study of existing loss weighting based multi-task learning techniques. By comparing them in both CV and NLP multi-task benchmark datasets, the authors also identified a simple random loss weighting method achieving similar performance. This paper provided both empirical and theoretical analysis of why random loss weighting can achieve similar performance efficiently.