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: In this work, the authors study the problem of multi-task learning in the presence of heterogenous tasks when there is an imbalance in the distribution of tasks. In such settings, we can learn models for multi-task learning using two approaches -- a) learn to minimize average loss across tasks, b) learn to minimize the worst case loss. Turns out both the approaches are not ideal. Minimizing average loss can compromise one task a lot and minimizing worst case loss can focus only on one very difficult task at the expense of improvements possible in other tasks. The authors propose a look ahead distributionally optimization based approach to balance the two -- worst case performance and the average performance. In the proposed approach, instead of solving a standard min-max optimization, the authors propose to constrain the min learner to find the optimal weight combination for weighted risk minimization, which leads to the minimum worst case risk. The authors carry out synthetic and real datasets to demonstrate the efficacy of the approach.