Balancing Average and Worst-case Accuracy in Multitask Learning

1 INTRODUCTION . Multitask learning—the process by which a single model is trained to perform a variety of different tasks—has become a subject of increasing interest with many successful applications in a variety of domains ( Ruder , 2017 ; Yu et al. , 2020 ; Wang et al. , 2021 ) . By and large , multitask learning models are evaluated by reporting their average performance on all tasks ( McCann et al. , 2018 ; Wang et al. , 2019 ) . However , average accuracy does not paint the full picture of a multitask model ’ s performance . Ensuring comparable accuracy across tasks , including on tasks from under-represented training data distributions , is an important issue from both practical and fairness perspectives . For example , in natural language processing , multilingual models often performs worse on languages with smaller amounts of resources on the web ( e.g. , Wikipedia ) due to the resulting scarcity of relevant training data ( Hu et al. , 2020 ) . In computer vision , facial recognition models have been shown to perform worse on racial groups that are underrepresented in the training data ( Buolamwini & Gebru , 2018 ) . In this paper , we argue that multitask models should also be evaluated by their worst-case accuracy . Yet commonly used task re-weighting based objectives are inadequate to ensure good worst-case performance . Rather , a more natural choice is to use a type of min-max objective from the distributionally robust optimization ( DRO ) literature ( Rahimian & Mehrotra , 2019 ; Sagawa et al. , 2020 ) : minθ maxi ` i ( θ ) , where ` i is the loss of the i-th task and θ are the model parameters . Training with this objective guarantees at least a certain level of performance on all tasks . However , we observe that DRO does not always work well in practice . Since the min-max objective optimizes the worst loss , it tends to do so by sacrificing performance on other tasks , which hurts average accuracy ( §2.3 ) . Motivated by this observation , we propose a novel optimization algorithm called Lookahead-DRO ( L-DRO ) . The intuition behind L-DRO is to anticipate the evolution of the loss of each task depending on which task we train on . Using this information , L-DRO chooses a weight allocation that ( i ) leads to minimal worst-case loss and ( ii ) trains on the maximal number of tasks possible ( §2.4 ) . We evaluate L-DRO on a synthetic task , an image classification task , and a language modeling task . In experiments on a multitask version of the CIFAR-100 dataset , we show that L-DRO consistently allows us to train models that achieve a good trade-off of worst and average task accuracy . On a realistic large-scale multilingual language modeling dataset , we find that L-DRO is able to improve performance on the worst-performing languages at very little cost to average performance . Our main contributions are as follows : ( i ) we show how distributionally robust optimization ( DRO ) can be applied to multitask learning to maximize worst-case accuracy ; ( ii ) we highlight failure cases of the standard DRO formulation in the multitask setting ; ( iii ) we propose a new algorithm , L-DRO that addresses DRO ’ s deficiencies ; and ( iv ) we demonstrate that L-DRO achieves a better trade-off between average and worst-case accuracy than strong baselines in a synthetic and two realistic experiments from two different modalities . 2 TRAINING MULTITASK MODELS . Consider a machine learning model parameterized by θ ∈ Rdmodel , which is trained to perform N tasks . In general , let a task refer to the combination of examples drawn from a data distribution and an objective function . For simplicity , we define a task solely by its objective function ` i=1 ... N . As a running example , we consider a multilingual language model with dozens of tasks where each task corresponds to a language ( Conneau et al. , 2020 ; Xue et al. , 2021 ) . 2.1 AVERAGE TASK LOSS . The most common objective for multitask learning is to minimize the average of all the task losses : min θ 1 N N∑ i=1 ` i ( θ ) A downside of this objective is that all tasks are treated equally . In real-world scenarios , we often encounter a scenario where task difficulty is not uniform ( i.e. , the magnitude of ` i varies wildly ; Hessel et al . 2019 ) or we want to pay more attention to certain tasks . If the distribution of all tasks is not homogeneous , training with Lavg results in a model with performance that is skewed towards certain clusters of tasks . In our multilingual example , if a large number of languages come from the same language family such as Indo-European languages , Lavg will choose models which perform better on these languages at the expense of typologically distinct languages such as Mandarin or Finnish . 2.2 STATIC MIXING . Rather than simple averaging , prior work in multitask learning often uses arbitrary weightings w = { w1 , . . . , wN } in the probablility simplex ∆N = { w | wi > 0 , ∑ i wi = 1 } to modulate the losses of individual tasks : min θ N∑ i=1 wi ` i ( θ ) . The weights are typically decided based on some statistics of the data . In proportional sampling ( Sanh et al. , 2019 ) , the weight of a task i is proportional to the size of its data wi ∝ |Di| . This biases the model towards tasks that are over-represented in the data . To counter-balance this , recent approaches use an inverse temperature α where wi ∝ |Di|α with α typically in the range of [ 0.2 , 0.3 ] ( Aharoni et al. , 2019 ; Conneau et al. , 2020 ; Xue et al. , 2021 ) . 2.3 WORST TASK LOSS . A third option that considers the task on which the model performs worst is the min-max objective . We refer to this objective as DRO because of its prominence in the distributional robustness literature ( Sagawa et al. , 2020 ) : min θ max i ` i ( θ ) = min θ max w∈∆N N∑ i=1 wi ` i ( θ ) ︸ ︷︷ ︸ LDRO ( θ ) . The above objective minimizes the worst-case loss of the model at every step . Alternatively , we can view the objective similar to the previous ones as minimizing the loss under a certain task weighting w. In this case , the weights w are not static , but rather chosen dynamically to maximize the weighted loss by allocating the maximum weight ( s ) to the task ( s ) on which the model performs worst . The main advantage of minimizing LDRO is that it ensures a certain level of performance on all tasks . However , it suffers from a number of drawbacks . First , it is sensitive to task difficulty : if one task ` hard is much more difficult than the others ( i.e. , minθ ` hard ( θ ) > minθ ` j ( θ ) for all other tasks j ) , then LDRO is likely to get “ stuck ” trying to minimize ` hard at the detriment of other tasks . In addition , DRO is agnostic to positive ( or negative ) interaction between tasks . In some cases , training on one task ` i may have a positive effect on another task ` j because they share some similarities—for example if they represent related languages . In such cases it may be desirable to train on both tasks rather than on one or the other ( as it may reduce the variance of the gradient estimator ) , even if one has a higher loss than the other . On the other hand , when two tasks are at odds with each other ( i.e. , training on one task decreases performance on the other ) , the max objective in LDRO can lead to oscillatory behavior where the weights w alternate between two competing tasks . 2.4 LOOKAHEAD DRO . The core limitation of LDRO underpinning the aforementioned issues is the fact that it does not take the optimization procedure into account ( i.e. , it will bluntly try to minimize the worst loss ) . This is suboptimal because in certain cases , taking gradient steps on the worst task might hurt other tasks disproportionately . If the worst-performing task has converged already , this weight allocation will also miss out on optimizing other tasks , which can attain lower losses and are still not converged . Formally , for a given value of the model parameters θ0 , LDRO chooses weights maximizing the model ’ s loss w∗0 : = arg maxw ∑ i wi ` i ( θ0 ) . After training with weights w ∗ 0 , we obtain parameters θ1 with loss LDRO ( θ1 ) = maxw ∑ i wi ` i ( θ1 ) . However , note that for θ1 , the worst-case weight allocation w∗1 : = arg maxw ∑ i wi ` i ( θ1 ) might be different than w ∗ 0 , so we have no guarantees that LDRO ( θ1 ) ≤ LDRO ( θ0 ) . In other words , training with the worst-case weights w∗0 does not guarantee that LDRO will decrease . This observation suggests an alternative strategy : suppose w , θ0 → θw , θ0 is a fixed—and , for simplicity , deterministic—gradient descent algorithm minimizing the w-weighted loss ∑ i wi ` i starting from parameters θ0 . Rather than choosing w to be the worst-case weights corresponding to θ0 , we should pick weights that minimize the worst-case loss of θ∗w , θ0 , i.e . after training . This can be formalized as solving the min-max problem : min w∈∆N max v∈∆N ∑ i vi ` i ( θw , θ0 ) ( 1 ) Intuitively , decoupling the weights w used for training from the weights v used for computing LDRO ( θ∗w , θ0 ) removes the restriction of just training on the worst-performing tasks and enables us to find a task weighting which is able to continuously minimize the worst-case loss . Note that in this case there might be multiple solution w : for instance , it might be that one or more tasks have the same effect on the resulting worst-case task , and as such any permutation of their weights is an optimal solution . In such scenarios , we do not want to miss out on training on tasks on which performance can still be improved ( as long as worst-case performance decreases ) . This can be formulated as a maximum entropy principle ( Jaynes , 1957 ) : if multiple optimal training weights exist , we should pick the one with the highest entropy . Because finding the optimal w necessitates “ looking ahead ” in the optimization trajectory from θ0 to identify the DRO loss maxv ∑ i vi ` i ( θw , θ0 ) , we call this weight selection principle LookaheadDRO ( L-DRO ) . We provide a graphical illustration in Figure 1 . 2.5 ONE-STEP LOOKAHEAD DRO . In practice we can not explicitly solve the min-max game in Eq . 1 because estimating the payoff maxv ∑ i vi ` i ( θw , θ0 ) would require unrolling the full optimization process to compute θ ∗ w , θ0 . Instead , we simplify and take only one gradient step with a small learning rate λ , i.e . replacing θ∗w , θ0 with θ0−λ ∑ i wi∇ ` i ( θ0 ) . Such one-step approximations are often used in the meta-learning literature ( Finn et al. , 2017 ; Liu et al. , 2018 ) . Using the first order Taylor extension of ` i enables us to rewrite the individual lookahead task losses as : ` i ( θ0−λ [ N∑ j=1 wj∇ ` j ( θ0 ) ] ) ≈ ` i ( θ0 ) − λ N∑ j=1 wj∇ ` i ( θ0 ) ᵀ∇ ` j ( θ0 ) The resulting payoffs are now linear in w , which reduces the lookahead optimization problem in Eq . 1 into a bilinear matrix game : min w max v vᵀ Lw ( 2 ) with a payoff matrix L , called the “ lookahead matrix ” , which can be decomposed as follows : Lij = ` i ( θ0 ) ︸ ︷︷ ︸ Loss of task i −λ∇ ` i ( θ0 ) ᵀ∇ ` j ( θ0 ) ︸ ︷︷ ︸ “ interaction ” between tasks i and j We can find a solution ( Nash equilibrium ) of this game using many available solvers such as fictitious play ( Brown , 1951 ) or linear programming ( Lemke , 1965 ) . In such cases where multiple equilibria exist , we choose the one with the highest entropy as dictated by our maximum entropy principle . The lookahead learning rate λ is treated as a hyperparameter . This results in a training algorithm which we outline in Algorithm 1 . Algorithm 1 : Lookahead DRO Input : Lookahead learning rate λ , initial parameters θ0 θ ← θ0 while training is not over do L← [ ` i ( θ ) − λ∇ ` i ( θ ) ᵀ∇ ` j ( θ ) ] ij // Compute lookahead matrix v∗ , w∗ ← maxent_nash_solver ( L ) // Find Nash equilibrium of L θ ← optimizer_step ( ∑ i w ∗ i ` i ) // Optimizer step on w ∗-weighted loss end To build more intuitions , we consider several edge cases and discuss how L-DRO handles them : Independent tasks . We have∇ ` ᵀi∇ ` j = 1i=j . In this case , for small enough values of the lookahead learning rate , the optimal strategy for the row player ( v ) is to pick the task i∗ = arg maxi ` i with the highest loss , and the only possible best response for the column player ( w ) is to match this selection ( since Li∗i = ` ∗i −λ < Lij = ` i for all j 6= i ) . In other words , we recover the original DRO objective of training against the worst-case loss . This can be interpreted as an “ implicit assumption ” of DRO that training on one task does not affect the others . Redundant tasks . We have ∇ ` ᵀi∇ ` j = 1 . In this case , the gradients of the tasks are perfectly aligned . In practice , this means that the payoff is independent of the value of w , so any choice of the training weights is valid . By the maximum entropy principle , we should pick uniform weights , thus recovering the average task loss baseline . The highest loss task has already converged . Consider the case where ` 0 is maximal , but the model has already converged on task 0 ( in other words ∇ ` 0 = 0 ) . In this scenario—assuming there is no strong negative interaction among the remaining tasks—the highest loss will always be ` 0 irrespective of the training weights w. Again , by the maximum entropy principle , L-DRO advocates for picking uniform weights . This makes it possible to continue training on the remaining tasks even though ` 0 has converged , bypassing the core limitations of distributionally robust optimization ( DRO ) identified in §2.3 .

Summary: This paper propose one-step look ahead DRO for multitask learning considering the worst task’s performance as the objective. Due to the task’s interaction with each other, directly using gradient-based approach is suboptimal. To overcome, look-ahead DRO estimates the task interaction using Taylor expansion and returns the locally optimal weighting by solving a linear min-max subproblem. Several tricks for mini-batch and online updating version of DRO are proposed.