Towards Impartial Multi-task Learning

1 INTRODUCTION . Recent deep networks in computer vision can match or even surpass human beings on some specific tasks separately . However , in reality multiple tasks ( e.g. , semantic segmentation and depth estimation ) must be solved simultaneously . Multi-task learning ( MTL ) ( Caruana , 1997 ; Evgeniou & Pontil , 2004 ; Ruder , 2017 ; Zhang & Yang , 2017 ) aims at sharing the learned representation among tasks ( Zamir et al. , 2018 ) to make them benefit from each other and achieve better results and stronger robustness ( Zamir et al. , 2020 ) . However , sharing the representation can lead to a partial learning issue : some specific tasks are learned well while others are overlooked , due to the different loss scales or gradient magnitudes of various tasks and the mutual competition among them . Several methods have been proposed to mitigate this issue either via gradient balance such as gradient magnitude normalization ( Chen et al. , 2018 ) and Pareto optimality ( Sener & Koltun , 2018 ) , or loss balance like homoscedastic uncertainty ( Kendall et al. , 2018 ) . Gradient balance can evenly learn task-shared parameters while ignoring task-specific ones . Loss balance can prevent MTL from being biased in favor of tasks with large loss scales but can not ensure the impartial learning of the shared parameters . In this work , we find that gradient balance and loss balance are complementary , and combining the two balances can further improve the results . To this end , we propose impartial MTL ( IMTL ) via simultaneously balancing gradients and losses across tasks . For gradient balance , we propose IMTL-G ( rad ) to learn the scaling factors such that the aggregated gradient of task-shared parameters has equal projections onto the raw gradients of individual tasks ∗Corresponding author ( see Fig . 1 ( d ) ) . We show that the scaling factor optimization problem is equivalent to finding the angle bisector of gradients from all tasks in geometry , and derive a closed-form solution to it . In contrast with previous gradient balance methods such as GradNorm ( Chen et al. , 2018 ) , MGDA ( Sener & Koltun , 2018 ) and PCGrad ( Yu et al. , 2020 ) , which have learning biases in favor of tasks with gradients close to the average gradient direction , those with small gradient magnitudes , and those with large gradient magnitudes , respectively ( see Fig . 1 ( a ) , ( b ) and ( c ) ) , in our IMTL-G task-shared parameters can be updated without bias to any task . For loss balance , we propose IMTL-L ( oss ) to automatically learn a loss weighting parameter for each task so that the weighted losses have comparable scales and the effect of different loss scales from various tasks can be canceled-out . Compared with uncertainty weighting ( Kendall et al. , 2018 ) , which has biases towards regression tasks rather than classification tasks , our IMTL-L treats all tasks equivalently without any bias . Besides , we model the loss balance problem from the optimization perspective without any distribution assumption that is required by ( Kendall et al. , 2018 ) . Therefore , ours is more general and can be used in any kinds of losses . Moreover , the loss weighting parameters and the network parameters can be jointly learned in an end-to-end fashion in IMTL-L. Further , we find the above two balances are complementary and can be combined to improve the performance . Specifically , we apply IMTL-G on the task-shared parameters and IMTL-L on the task-specific parameters , leading to the hybrid balance method IMTL . Our IMTL is scale-invariant : the model can converge to similar results even when the same task is designed to have different loss scales , which is common in practice . For example , the scale of the cross-entropy loss in semantic segmentation may have different scales when using “ average ” or “ sum ” reduction over locations in the loss computation . We empirically validate that our IMTL is more robust against heavy loss scale changes than its competitors . Meanwhile , our IMTL only adds negligible computational overheads . We extensively evaluate our proposed IMTL on standard benchmarks : Cityscapes , NYUv2 and CelebA , where the experimental results show that IMTL achieves superior performances under all settings . Besides , considering there lacks a fair and practical benchmark for comparing MTL methods , we unify the experimental settings such as image resolution , data augmentation , network structure , learning rate and optimizer option . We re-implement and compare with the representative MTL methods in a unified framework , which will be publicly available . Our contributions are : • We propose a novel closed-form gradient balance method , which learns task-shared parameters without any task bias ; and we develop a general learnable loss balance method , where no distribution assumption is required and the scale parameters can be jointly trained with the network parameters . • We unveil that gradient balance and loss balance are complementary and accordingly propose a hybrid balance method to simultaneously balance gradients and losses . • We validate that our proposed IMTL is loss scale-invariant and is more robust against loss scale changes compared with its competitors , and we give in-depth theoretical and experimental analyses on its connections and differences with previous methods . • We extensively verify the effectiveness of our IMTL . For fair comparisons , a unified codebase will also be publicly available , where more practical settings are adopted and stronger performances are achieved compared with existing code-bases . 2 RELATED WORK . Recent advances in MTL mainly come from two aspects : network structure improvements and loss weighting developments . Network-structure methods based on soft parameter-sharing usually lead to high inference cost ( review in Appendix A ) . Loss weighting methods find loss weights to be multiplied on the raw losses for model optimization . They employ a hard parameter-sharing paradigm ( Ruder , 2017 ) , where several light-weight task-specific heads are attached upon the heavy-weight task-agnostic backbone . There are also efforts that learn to group tasks and branch the network in the middle layers ( Guo et al. , 2020 ; Standley et al. , 2020 ) , which try to achieve better accuracyefficiency trade-off and can be seen as semi-hard parameter-sharing . We believe task grouping and loss weighting are orthogonal and complementary directions to facilitate multi-task learning and can benefit from each other . In this work we focus on loss weighting methods which are the most economic as almost all of the computations are shared across tasks , leading to high inference speed . Task Prioritization ( Guo et al. , 2018 ) weights task losses by their difficulties to focus on the harder tasks during training . Uncertainty weighting ( Kendall et al. , 2018 ) models the loss weights as dataagnostic task-dependent homoscedastic uncertainty . Then loss weighting is derived from maximum likelihood estimation . GradNorm ( Chen et al. , 2018 ) learns the loss weights to enforce the norm of the scaled gradient for each task to be close . MGDA ( Sener & Koltun , 2018 ) casts multi-task learning as multi-object optimization and finds the minimum-norm point in the convex hull composed by the gradients of multiple tasks . Pareto optimality is supposed to be achieved under mild conditions . GLS ( Chennupati et al. , 2019 ) instead uses the geometric mean of task-specific losses as the target loss , we will show it actually weights the loss by its reciprocal value . PCGrad ( Yu et al. , 2020 ) avoids interferences between tasks by projecting the gradient of one task onto the normal plane of the other . DSG ( Lu et al. , 2020 ) dynamically makes a task “ stop or go ” by its converging state , where a task is updated only once for a while if it is stopped . Although many loss weighting methods have been proposed , they are seldom open-sourced and rarely compared thoroughly under practical settings where strong performances are achieved , which motivates us to give an in-depth analysis and a fair comparison about them . 3 IMPARTIAL MULTI-TASK LEARNING . In MTL , we map a sample x ∈ X to its labels { yt ∈ Yt } t∈ [ 1 , T ] of all T tasks through multiple taskspecific mappings { ft : X→ Yt } . In most loss weighting methods , the hard parameter-sharing paradigm is employed , such that ft is parameterized by heavy-weight task-shared parameters θ and light-weight task-specific parameters θt . All tasks take the same shared intermediate feature z = f ( x ; θ ) as input , and the t-th task head outputs the prediction as ft ( x ) = ft ( z ; θt ) . We aim to find the scaling factors { αt } for all T task losses { Lt ( ft ( x ) , yt ) } , so that the weighted sum loss L = ∑ t αtLt can be optimized to make all tasks perform well . This poses great challenges because : 1 ) losses may have distinguished forms such as cross-entropy loss and cosine similarity ; 2 ) the dynamic ranges of losses may differ by orders of magnitude . In this work , we propose a hybrid solution for both the task-shared parameters θ and the task-specific parameters { θt } , as Fig . 2 . 3.1 GRADIENT BALANCE : IMTL-G For task-shared parameters θ , we can receive T gradients { gt = ∇θLt } via back-propagation from all of the T raw losses { Lt } , and these gradients represent optimal update directions for individual tasks . As the parameters θ can only be updated with a single gradient , we should compute an aggregated gradient g by the linear combination of { gt } . It also implies to find the scaling factors { αt } of raw losses { Lt } , since g = ∑ t αtgt = ∇θL = ∇θ ( ∑ t αtLt ) . Motivated by the principle of balance among tasks , we propose to make the projections of g onto { gt } to be equal , as Fig . 1 ( d ) . In this way , Algorithm 1 Training by Impartial Multi-task Learning Input : input sample x , task-specific labels { yt } and learning rate η Output : task-shared/-specific parameters θ/ { θt } , scale parameters { st } 1 : compute task-shared feature z = f ( x ; θ ) 2 : for t = 1 to T do 3 : compute task prediction by head network ft ( x ) = fnett ( z ; θt ) 4 : compute raw loss by loss function Lrawt = Lfunct ( ft ( x ) , yt ) 5 : compute scaled loss Lt = bastLrawt − st ( default a = e , b = 1 ) . loss balance 6 : compute gradient of shared feature z : gt = ∇zLt 7 : compute unit-norm gradient ut = gt‖gt‖ 8 : end for 9 : compute gradient differencesD > = [ g > 1 − g > 2 , · · · , g > 1 − g > T ] 10 : compute unit-norm gradient differences U > = [ u > 1 − u > 2 , · · · , u > 1 − u > T ] 11 : compute scaling factors for tasks 2 to T : α2 : T = g1U > ( DU > ) −1 . gradient balance 12 : compute scaling factors for all tasks : α = [ 1− 1α > 2 : T , α2 : T ] 13 : update task-shared parameters θ = θ − η∇θ ( ∑ t αtLt ) 14 : for t = 1 to T do 15 : update task-specific parameters θt = θt − η∇θtLt 16 : update loss scale parameter st = st − η ∂Lt∂st 17 : end for we treat all tasks equally so that they progress in the same speed and none is left behind . Formally , let { ut = gt/ ‖gt‖ } denote the unit-norm vector of { gt } which are row vectors , then we have : gu > 1 = gu > t ⇔ g ( u1 − ut ) > = 0 , ∀ 2 6 t 6 T. ( 1 ) The above problem is under-determined , but we can obtain the closed-form results of { αt } by constraining ∑ t αt = 1 . Assume α = [ α2 , · · · , αT ] , U > = [ u > 1 − u > 2 , · · · , u > 1 − u > T ] , D > = [ g > 1 − g > 2 , · · · , g > 1 − g > T ] and 1 = [ 1 , · · · , 1 ] , from Eq . ( 1 ) we can obtain : α = g1U > ( DU > ) −1 . ( IMTL-G ) ( 2 ) The detailed derivation is in Appendix B.1 . After obtaining α , the scaling factor of the first task can be computed by α1 = 1 − 1α > since ∑ t αt = 1 . The optimized { αt } are used to compute L =∑ t αtLt , which is ultimately minimized by SGD to update the model . By now , back-propagation needs to be executed T times to obtain the gradient of each task loss with respect to the heavy-weight task-shared parameters θ , which is time-consuming and non-scalable . We replace the parameterlevel gradients { gt = ∇θLt } with feature-level gradients { ∇zLt } to compute { αt } . This implies to achieve gradient balance with respect to the last shared feature z as a surrogate of task-shared parameters θ , since it is possible for the network to back-propagate this balance all the way through the task-shared backbone starting from z . This relaxation allows us to do back propagation through the backbone only once after obtaining { αt } , and thus the training time can be dramatically reduced .

Summary: The authors propose to balance multi-task training using IMTL-G on the shared backbone and IMTL-L on the task-specific branches. IMTL-G enforces equal gradient projections between tasks with a close-form formulation to calculate the desired gradient weightings $\alpha$. IMTL-L learns the loss weightings $e^s$ with a regularization term $-s$. Additional constraint by making all loss weightings sum to one is used. The paper compares the effectiveness of the proposed IMTLs with their counterparts on Cityscapes, NYUv2, and CelebA and claims state-of-the-art performance.