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: Here, the authors present a novel method for multitasking loss weighting called Random Loss Weighting (RLW). The method is straightforward and clearly described; RLW considers the loss weights lambda as random variables. In practice, it works by sampling the loss weights from a distribution (authors evaluated six different distributions) with some normalization and then minimizing the aggregated loss weighted by the normalized random weights. The authors provided an extensive empirical evaluation including eight state-of-the-art loss weighting strategies on six computer vision datasets and the XTREME benchmark (four multilingual problems). RWL results are comparable to previous strategies, but with a lower computational burden. The authors also provided a theoretical analysis on how the extra randomness from the sampling affects the convergence and effectiveness of RLW compared with fixed loss weights optimized via SGD.