B ALANCING A VERAGE AND W ORST - CASE A CCURACY IN M ULTITASK L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
When training and evaluating machine learning models on a large number of tasks, it is important to not only look at average task accuracy—which may be biased by easy or redundant tasks—but also worst-case accuracy (i e. the performance on the task with the lowest accuracy).
In this work, we show how to use techniques from the distributionally robust optimization (DRO) literature to improve worst-case performance in multitask learning.
We highlight several failure cases of DRO when applied off-the-shelf and present an improved method, Lookahead-DRO (L-DRO), which mitigates these issues.
The core idea of L-DRO is to anticipate the interaction between tasks during training in order to choose a dynamic re-weighting of the various task losses, which will (i) lead to minimal worst-case loss and (ii) train on as many tasks as possible.
After demonstrating the efﬁcacy of L-DRO on a small controlled synthetic setting, we evaluate it on two realistic benchmarks: a multitask version of the CIFAR-100 image classiﬁcation dataset and a large-scale multilingual language modeling experiment.
Our empirical results show that LDRO achieves a better trade-off between average and worst-case accuracy with little computational overhead compared to several strong baselines
I NTRODUCTION
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 θ max i (cid:96) i ( θ ) , where (cid:96) 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 sacriﬁcing 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 classiﬁcation 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 ﬁnd 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 deﬁciencies; 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.
T RAINING M ULTITASK M ODELS
Consider a machine learning model parameterized by θ ∈ R d model , 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 deﬁne a task solely by its objective function (cid:96) 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. A VERAGE T ASK L OSS
The most common objective for multitask learning is to minimize the average of all the task losses:
[Equation]
A downside of this objective is that all tasks are treated equally.
In real-world scenarios, we often encounter a scenario where task difﬁculty is not uniform (i.e., the magnitude of (cid:96) 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 L avg 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, L avg will choose models which perform better on these languages at the expense of typologically distinct languages such as Mandarin or
Finnish.
2. S TATIC MIXING
Rather than simple averaging, prior work in multitask learning often uses arbitrary weightings w = { w 1 , . .
. , w N } in the probablility simplex ∆ N = { w | w i > 0 , (cid:80) i
w i = 1 } to modulate the losses of individual tasks:
[Equation]
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 w i ∝ | D i | .
This biases the model towards tasks that are over-represented in the data.
To counter-balance this, recent approaches use an inverse temperature α where w i ∝ | D i | α with α typically in the range of [0 . 2 , 0 . 3]
(Aharoni et al., 2019; Conneau et al., 2020; Xue et al., 2021).
2. W ORST T ASK L OSS
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):
[Equation]
[Figure]
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 L DRO 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 difﬁculty: if one task (cid:96) hard is much more difﬁcult than the others (i.e., min θ (cid:96) hard ( θ ) > min θ (cid:96) j
( θ ) for all other tasks j ), then L DRO is likely to get “stuck” trying to minimize (cid:96) 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 (cid:96) i may have a positive effect on another task (cid:96) 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
L DRO can lead to oscillatory behavior where the weights w alternate between two competing tasks.
2. L OOKAHEAD DRO
The core limitation of L DRO 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 , L DRO chooses weights maximizing the model’s loss w ∗ 0 := arg max w (cid:80) i w i (cid:96) i ( θ 0 ) .
After training with weights w ∗ 0 , we obtain parameters θ 1 with loss L DRO ( θ 1 ) = max w (cid:80) i w i (cid:96) i ( θ 1 ) .
However, note that for θ 1 , the worst-case weight allocation w ∗ 1 := arg max w (cid:80) i w i (cid:96) i ( θ 1 ) might be different than w ∗ 0 , so we have no guarantees that
L DRO ( θ 1 ) ≤ L DRO ( θ 0 ) .
In other words, training with the worst-case weights w ∗ 0 does not guarantee that L DRO will decrease.
This observation suggests an alternative strategy: suppose w , θ 
 → θ w , θ 0 is a ﬁxed—and, for simplicity, deterministic—gradient descent algorithm minimizing the w -weighted loss (cid:80)
i w i (cid:96) 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:
[Equation]
Intuitively, decoupling the weights w used for training from the weights v used for computing
L DRO ( θ ∗ w , θ 0 ) removes the restriction of just training on the worst-performing tasks and enables us to ﬁnd 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 ﬁnding the optimal w necessitates “looking ahead” in the optimization trajectory from θto identify the DRO loss max v (cid:80) i v i (cid:96) i ( θ w , θ 0 ) , we call this weight selection principle
LookaheadDRO (L-DRO).
We provide a graphical illustration in Figure 1.
2. O NE - STEP L OOKAHEAD DRO
In practice we cannot explicitly solve the min-max game in Eq. 1 because estimating the payoff max v (cid:80) i v i (cid:96) 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 , θwith θ 0 − λ (cid:80) i
w
i ∇ (cid:96) i ( θ 0 ) .
Such one-step approximations are often used in the meta-learning literature (Finn et al., 2017; Liu et al., 2018).
Using the ﬁrst order Taylor extension of (cid:96) i enables us to rewrite the individual lookahead task losses as:
[Equation]
The resulting payoffs are now linear in w , which reduces the lookahead optimization problem in
Eq.
1 into a bilinear matrix game:
[Equation]
with a payoff matrix L , called the “lookahead matrix”, which can be decomposed as follows:
[Equation]
We can ﬁnd a solution (Nash equilibrium) of this game using many available solvers such as ﬁctitious 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:
Input θ ← θwhile training is not over do
L ← [ (cid:96) i ( θ ) − λ ∇ (cid:96) i ( θ ) (cid:124) ∇ (cid:96) j ( θ )] ij // Compute lookahead matrix v ∗ , w ∗ ← maxent_nash_solver ( L ) //
Find Nash equilibrium of L θ ← optimizer_step( (cid:80) i w ∗ i (cid:96) 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 ∇ (cid:96) (cid:124) i ∇ (cid:96) j = 1 i = 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 max i (cid:96) i with the highest loss, and the only possible best response for the column player ( w ) is to match this selection (since L i ∗ i
= (cid:96) ∗ i −
λ <
L ij = (cid:96) i for all j (cid:54) = 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 ∇ (cid:96) (cid:124) i ∇ (cid:96) 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 (cid:96) 0 is maximal, but the model has already converged on task
(in other words ∇ (cid:96) 0 = 0 ).
In this scenario—assuming there is no strong negative interaction among the remaining tasks—the highest loss will always be (cid:96)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 (cid:96) 0 has converged, bypassing the core limitations of distributionally robust optimization (DRO) identiﬁed in §2.3.
L OOKAHEAD DRO IN P RACTICE 3. G RADIENT AND L OSS N OISES
In most realistic machine learning scenarios we do not have access to the full losses (cid:96) i ( θ ) and their gradients.
Rather, we compute stochastic estimates on minibatches ˆ (cid:96) i ( θ ) .
In preliminary experiments, we ﬁnd that this often leads to instability in the training algorithm, since even little noise in the lookahead matrix can lead to dramatically different Nash equilibria, which causes high ﬂuctuations of the training weights w ∗ .
To mitigate this issue, we adopt an online weight update procedure inspired by the online-DRO algorithm proposed by Sagawa et al.
Instead of computing the optimal weights w ∗ at each step based on noisy losses and gradients, we keep a running set of weights w , which we update jointly with model parameters θ .
Speciﬁcally, we interleave gradient updates on θ with exponentiated gradient steps (Kivinen & Warmuth, 1997) on w to minimize max v v (cid:124)
L w .
We show the weight update rule in Algorithm 2.
Algorithm 2:
DRO update Function update_weights( Lookahead matrix L , weights w , weight update rate η ) : v ∗ ← arg max v v (cid:124)
L w //
Compute best response to w ˜ w i ← w i e − η
[ v ∗ (cid:124) L ]
i // Update weights w i ← ˜ w i (cid:80) j ˜ w j // Re-normalize weights return w
In practice, we ﬁnd that this online weight update alleviates the instability issues mentioned above.
Note that the choice of exponentiated gradient is not arbitrary.
It corresponds to a mirror descent step (Nemirovskij & Yudin, 1983) using the Kullback-Leibler divergence (Kullback & Leibler, 1951) as a distance function.
Choosing weight updates that locally minimize the change in relative entropy plays into our maximum entropy guiding principle (Jumarie, 1990; Kapur & Kesavan, 1992).
3. L OOKAHEAD M ATRIX C OMPUTATION
The main computational bottleneck of L-DRO lies in the computation of the lookahead matrix L .
Recall that L ij can be decomposed in two terms: the loss (cid:96) i and an “interaction” term ∇ (cid:96) i ( θ ) (cid:124) ∇ (cid:96) j
( θ ) .
While the former can be estimated essentially “for free” during the forward pass of back-propagation, computing the pairwise task gradients dot products is much more computationally demanding.
Indeed, it necessitates computing (and holding in memory) n model gradients, which can be prohibitive for larger models.
To reduce the computational overhead, we only update this second term (which we call the “interaction matrix” A =
[ ∇ (cid:96) i
( θ ) (cid:124) ∇ (cid:96)
j ( θ )]
i,j =1 ...N ) periodically, typically everyor 1000 training steps.
In preliminary experiments, we also ﬁnd that Lookahead-DRO works more reliably when the lookahead gradients of each task are re-normalized to one.
This means that the interaction matrix is actually computed to be A ij = ∇ (cid:96) i ( θ ) (cid:124) ( ∇ (cid:96) j
( θ ) /
(cid:107)∇ (cid:96) j
( θ ) (cid:107) ) .
The effect of these two design choices are investigated in ablation studies in Appendices A.1 and A.2.
We show the ﬁnal algorithm for online L-DRO in Algorithm 3.
S YNTHETIC E XPERIMENTS AS A P ROOF OF C ONCEPT
In this section, we demonstrate the efﬁcacy of L-DRO on a synthetic multitask setting.
We design a toy experiment such that neither DRO nor static mixing is able to reach an optimal solution (Figure 2).
We identify each task i with a point y i on the two dimensional plane R 2 and deﬁne their respective loss functions as the squared Euclidean distance of the model prediction f ( θ ) to y i , (cid:96) i ( θ )
= (cid:107)
y
i − f
( θ ) (cid:107) 22 .
We distribute the tasks such that most points are clustered around ( − 1 , 0) , to simulate a grouping of similar tasks ( e.g. Indo-Europoean languages in our multilingual example).
We then place a single point y outlier at (1 , 0) to emulate the presence of an outlier task.
Input:
[Table]
rate
[Table] [Table]
[Table]
k A ←// Initialize interaction matrix atw ←
[Equation]
// Initialize w as uniform weights while training is not over do if i% k == 0 then
[Equation]
// Update interaction matrix
[Equation]
// Compute lookahead matrix
[Equation]
[Table]
weights (
[Equation]
[Table]
Update weights using Algorithmθ
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
// Optimizer step on w -weighted loss
We simulate the increased difﬁculty of the outlier task by adopting a parameterization f ( θ ) that prevents the model from reaching any position closer than a ﬁxed value R from y outlier .
Speciﬁcally we adopt a polar parameterization θ :=
( r, φ ) and set
[Equation]
This essentially creates a “forbidden region” of radius R around y outlier of the space that the model cannot penetrate, thus lower-bounding (cid:96) outlier to R .
In experiments, we set R = 1 . 25 .
Figure 2a depicts the optimization trajectories when the model is trained either with the average objective or with DRO.
As expected, minimizing the average loss leads to a solution that performs well on the majority of tasks but poorly on the outlier.
On the other hand, the DRO trajectory moves towards the latter but ends up stuck when it reaches a point where (cid:96) outlier cannot be improved despite being the highest loss among all tasks.
On the other hand, while the L-DRO trajectory shown in Figure 2b follows a similar trajectory as
DRO at ﬁrst, it is able to reach a solution that still achieves the best possible (cid:96) outlier while achieving lower loss on the other tasks.
Also notice that the online variant follows a similar trajectory, despite the fact that it does not compute a Nash equilibrium at every step.
Finally, Figure 2c highlights the importance of the online variant in the presence of noise.
We emulate the stochasticity of SGD by adding a small amount of Gaussian noise to both losses and gradients.
Under these conditions, lookahead DRO becomes very unstable, an issue that is partially mitigated by the online update.
E XPERIMENTS 5. CIFAR
As a ﬁrst step towards a more realistic setting, we experiment with a multitask version of the CIFAR100 dataset (Krizhevsky et al., 2009), in which the 20 coarse labels are treated as separate 5-way classiﬁcation tasks (Rosenbaum et al., 2018; Yu et al., 2020).
In this setting, we train a small CNN model with four convolution layers with 32 3 × 3 ﬁlters each followed by a ReLU activation and a 2 × 2 max-pooling layer.
The resulting representation is averaged and fed through a 3-layer fully connected feed-forward network with hidden dimension 128 and ReLU activations.
Predictions are made using a shared softmax layer across all tasks ( i.e. across all 100 classes).
In addition to L-DRO, we report results for three baselines.
First static-mixing , where all tasks are weighted equally (§2.2).
Another natural baseline is DRO (§2.3) using the online algorithm proposed by Sagawa et al.
Finally we compare to PCGrad (Yu et al., 2020), a multitask learning approach, which uses gradient level projections to mitigate interferences between tasks during learning.
We train using regular stochastic gradient descent with a batch size of 128, a learning rate of 0 .
025 and with Nesterov accelerated gradients with a momentum rate of 0 . 9 .
For
DRO, we sweep over 5 values of the hyper-parameter which controls the online weight update rate { 0 , 0 . 0001 , 0 .
For L-DRO, we sweep over 3 values of the weight update rate { 0 . 01 , 0 .
001 , 0 . 0001 } and 3 values of the lookahead learning rate { 0 .
For Static mixing and PCGrad, there are no additional parameters to sweep over.
AverageDRO (a) DRO vs. Average loss
L-DROOnline L-DRO (b) L-DRO
L-DROOnline L-DRO (c) L-DRO (with noise)
To gain a better understanding of the various trade-offs of each approach in terms of average and worst accuracy, we do not report a single number.
Rather, we periodically evaluate the trained models on the test set and report the Pareto frontier of all checkpoints (across epochs and hyper-parameter conﬁgurations) of each method in the two dimensional plane of worst and average accuracy.
We show the results in Figure 2.
We can see that L-DRO (in blue) consistently achieves a good trade-off between the two metrics, as its Pareto frontier is situated in the upper right corner of the space (high average and worstcase accuracy).
In particular it is able to reach similar worst-case performance as DRO, without incurring a hit in terms of average performance.
We ﬁnd that this trend is consistent across different random seeds (see Appendix B.1).
An important difference in our experimental setting with Yu et al.
(2020) is that the model we train is comparatively much smaller.
In fact, as we increase the size of the model, we observe that the difference between individual method grows thinner (detailed results are depicted in Appendix B.3).
This ﬁnding suggests that L-DRO is primarily useful when the model’s capacity is a limiting factor compared to the number of tasks and the difﬁculty of individual tasks, which is not necessarily the case for the simple CIFAR-100 benchmark.
5. M ULTILINGUAL L ANGUAGE M ODELING
To evaluate L-DRO in a real-world scenario, we turn to large-scale multilingual language modeling.
As NLP models are increasingly applied to new settings, pretraining on many languages is becoming a common scenario (Conneau et al., 2020; Xue et al., 2021).
Improved pretraining performance has generally been found to correlate with better downstream performance (Raffel et al., 2020).
Furthermore, multilingual language modeling is a particulary interesting testbed because it exhibits various properties of real world multitask learning scenarios.
For example, there is a large data imbalance across tasks: the highest resource language (English) contains 55 × more data than the lowest resource language (Yorùbá) in our training corpus.
Additionally, there is a difference in difﬁculty between the various languages.
Formally, for a parametric language model p θ , its loss on the i -th language is the cross-entropy H ( p i , p θ ) with the underlying distribution of the language, p i .
This can be re-written as H ( p i ) + D KL
( p i || p θ ) .
Since the KL divergence is non-negative, the loss of each language is lower-bounded by the language-speciﬁc entropy term H ( p i ) , irrespective of the capacity or the ﬁtness of the model.
In other words, languages with higher entropy will be more difﬁcult to learn (Mielke et al., 2019).
5.2. E XPERIMENTAL S ETTING
We train a large autoregressive transformer model (Vaswani et al., 2017) on the multilingual C dataset (mC4; Xue et al., 2021).
We limit ourselves to the 40 languages that are present in the
XTREME benchmark (Hu et al., 2020), which belong to 12 different language families and are written in 12 different scripts.
Out of all languages, 40% are Indo-European and 42.5% are written in the Latin script.
1 The mC4 dataset comes with a canonical “training” and “evaluation” split.
In order to distinguish between validation and test set, we modify the evaluation split as follows: for each language we split the evaluation set into a validation and test set containing up to 5M tokens each (using mT5 subword tokenization; Kudo 2018), or half the entire evaluation set, whichever is largest.
We limit the size of the validation and test sets to facilitate running validation of our models.
We train each model for 300 , 000 steps with a batch size of 256 elements of 512 tokens using the
Adam optimizer (Kingma & Ba, 2014) and a cosine annealed learning rate schedule (Vaswani et al., 2017).
For each run we select the best model in terms of average validation loss on all languages.
For evaluation, we set the sequence length to 1024 .
In this experiment we compare against three baselines:
• Static mixing (§2.2).
We try 3 different values for the inverse temperature hyper-parameter α : 0 (uniform sampling), 0 .
(a common value in the literature; Xue et al. 2021) and 1 (purely proportional sampling) • DRO (§2.3), which minimizes the worst-case loss among all languages using the online algorithm from Sagawa et al. (2020).
We sweep over the weight learning rate η ∈ { 0 . 01 , 0 . 001 , 0 . 0001 } .
• Baselined-DRO , a variation on DRO proposed by Oren et al.
(2019) where the losses of each language i is adjusted with a “baseline” loss to account for the entropy term which biases DRO towards selecting languages that are inherently more difﬁcult.
In practice, we estimate H ( p i ) by training a separate monolingual model on each language, p θ ∗ i and using its cross-entropy H
( p
[Equation]
p
[Equation]
=
[Equation]
θ ) as a proxy.
We then train to minimize max i (cid:96) i ( θ ) − H ( p i , p θ ∗ i ) using the same algorithm as DRO, sweeping over the same values of the weight learning rate.
Note that this particular variant necessitates training an additional model for each language.
5.2. R ESULTS
We report results in Table 1 in terms of perplexity (Jelinek et al., 1977), the standard metric for language modeling.
Perplexity is the exponential of the (token-wise) cross-entropy.
As a result, perplexities of different languages are not directly comparable.
Languages with inherently higher entropy will tend to contribute more to the average (or worst) perplexity.
Therefore, we also report average and worst-case perplexity relative to a monolingual model.
Speciﬁcally, the relative perplexity of a model θ on language i is deﬁned as:
[Equation]
where θ ∗ i is a monolingual model trained on language i .
In other words, relative perplexity of a given model on a language quantiﬁes how much worse the multilingual model is compared to a monolingual model trained on the same language (with the same data).
Our results show that L-DRO reaches the best average score both in terms of perplexity and relative perplexity.
While DRO is by far the best alternative in terms of worst perplexity, it ranks lowest in terms of relative perplexity.
This is because its focus on languages with high perplexity comes at the detriment of low-entropy languages.
Unsurprisingly, proportional mixing based methods ( α = 0 .
3 or 1 ) perform very well on high resource languages ( e.g. English or French) but poorly on almost all other languages, leading to low performance along almost all metrics.
On the other hand, L-DRO performs much better in terms of relative perplexity: it outperforms both static mixing and DRO in terms of average and worst-case performance.
Note that Baselined-DRO, which takes the entropy of individual languages into account when computing the maximum, reaches the lowest worst-case relative perplexity.
However, it is a much more computationally demanding approach: it requires training 39 additional monolingual models to estimate the language entropies.
In comparison, L-DRO comes in as a close second in terms of worst-case relative perplexity, with a comparatively much smaller computational overhead (periodic computation of L ), which makes it
[Table]
a much more appealing alternative.
A more detailed discussion of the computational overhead of
L-DRO can be found in appendix C.
R ELATED W ORK
Much work in multi-task learning has focused on improving the optimization process when learning many tasks.
Some prior work (Sener & Koltun, 2018; Lin et al., 2019; Ma et al., 2020) casts multi-task learning as multi-objective optimization, with the goal of ﬁnding a Pareto-optimal solution.
Other work proposes to directly modify the gradients.
In order to minimize gradient interference, Yu et al.
(2020) project conﬂicting gradients on the normal plane of the other.
Wang et al.
(2021) alter both the direction and magnitude of gradients to achieve a certain gradient similarity.
Many other methods have been proposed to weight the task losses according to different criteria such as uncertainty (Kendall et al., 2018), reward magnitude (Hessel et al., 2019), and learning speed (Chen et al., 2018; Liu et al., 2019; Zheng et al., 2019).
Most of the latter methods increase a task’s loss weight when the learning speed for that task is low.
Similar to learning speed, a few approaches (Guo et al., 2018; Jean et al., 2019) weight task losses based on performance, focusing on optimizing tasks with poor performance.
Among these, our proposed method is most similar in spirit to performance weighting-based methods.
Such methods, however, do not consider the optimization procedure and are not aware of task similarity.
Loss weighting can be seen as a continuous relaxation of the task scheduling and sampling so that most task scheduling methods can be easily adapted to loss weighting methods, and vice versa (Ruder, 2019; Crawshaw, 2020).
Kiperwasser & Ballesteros (2018) proposed different predeﬁned sampling schedules while Sanh et al.
(2019) proposed a proportional sampling strategy.
In multilingual language modeling and machine translation (Aharoni et al., 2019; Conneau et al., 2020; Xue et al., 2021), proportional sampling with a temperature is typically employed.
The distributional robustness literature focuses on training models that perform well on various domains (Ben-Tal et al., 2009; Rahimian & Mehrotra, 2019).
This is generally performed by minimizing the worst-case expected risk over a pre-determined family called the “uncertainty” set.
There has been substantial work on formulating appropriate uncertainty sets (Hu & Hong, 2013; Gao & Kleywegt, 2016; Levy et al., 2020).
Most relevant to our multitask setting are group-structured uncertainty sets, where the maximimum is taken over a mixture of sub-populations (Oren et al., 2019;
Sagawa et al., 2020; Zhou et al., 2021).
C ONCLUSION
In this paper, we argued for the importance of looking at worst-case performance when training and evaluating multi-task models, especially in cases where the distribution of tasks is not homogeneous (e.g., when some tasks are more difﬁcult than others) or when there is an imbalance towards certain types of tasks.
We proposed L-DRO, which enables efﬁcient optimization of the worst-case loss without the pitfalls of existing DRO-based max-loss minimization approaches.
We demonstrated the beneﬁts of L-DRO in a synthetic scenario as well as in realistic image classiﬁcation and language modeling experiments.
Overall, we encourage the community to evaluate not only on homogeneous task collections (e.g., CIFAR-100) but also “heterogeneous” multitask learning scenarios (e.g., multilingual language modeling) when assessing the performance of a multitask model.
E THICS S TATEMENT
Measuring models only based on their average performance obfuscates their limitations on underrepresented subpopulations and contributes to the marginalization of said groups.
Developing algorithms that work well not just for the average user but for every type of user is thus important from an ethical perspective.
R EPRODUCIBILITY S TATEMENT
We take several steps to ensure reproducibility of our work.
Pseudo-code describing the ﬁnal version of our algorithm is given in Algorithm 3.
Moreover, we detail experimental settings and model hyper-parameters in §5.
Finally, all our experiments are performed on datasets that are openly available: CIFAR-100 3 (Krizhevsky et al., 2009) and mC4 4 (Xue et al., 2021) A A BLATION STUDIES A. E FFECT OF P ERIODIC M ATRIX U PDATES
We investigate the effect of periodic matrix updates on the results of L-PDRO.
Figure 4 reports the Pareto frontier of L-DRO over a range of hyper-parameters η and λ on the CIFAR-100 experiment from Section 5.1. We compare frontiers for 3 different interaction matrix update intervals: 10 ,and 1000 , and report results for 3 random seeds. Although there is some variability across different runs, we do not observe that larger intervals are consistently worse than smaller intervals.
That being said, in experiments with larger models, we have observed some instability with the larger update intervals.
This is evidenced in Figure 5a for a
CIFAR-100 run with the larger CNN model used in Section B.3.
A. E FFECT OF R ENORMALIZATION
Note that there of the two components in each element in the interaction matrix A ij = ∇ (cid:96) (cid:124)
i ∇ (cid:96) j
and we only normalize one ( ∇ (cid:96) (cid:124) i ∇ (cid:96)
j → ∇ (cid:96) (cid:124)
i ∇ (cid:96) j |∇ (cid:96) j | ). This can be interpreted as the ﬁrst order approximation of taking a lookahead step along the *normalized* task gradients: [Equation] Our intuition with this re-normalization trick is to help mitigate the instability caused by the “staleness” of the interaction matrix with periodic updates.
Indeed, the approximation error of using a stale interaction matrix comes from two sources: outdated alignment between the gradient (cosine similarity) and outdated gradient norms.
With the re-normalization trick, we mitigate the latter by factoring out the norm of the lookahead step while keeping the original interpretation (as explained above).
In practice, we do observe that this helps improve stability for larger matrix update interval.
Indeed, in an ablation study in CIFAR-100 (see Figure 5b), we ﬁnd that the use of the re-normalization trick removes the instability observed with larger update intervals (1000).
B A DDITIONAL CIFAR-100 R ESULTS B. M ULTIPLE R ANDOM S EEDS
Figure 6 depicts results on the CIFAR-100 dataset for 3 different random seeds.
We ﬁnd that L-DRO consistently outperforms baselines. [Figure] [Figure]
Another natural baseline for trading off worst and average accuracy is to optimize a weighted average of average and worst-case task loss:
[Equation]
with some interpolation parameter α . A downside of this approach is that, for α > 0 , the model cannot, in general, reach the minimal worst-case loss because of the nonzero “average loss” component. Figure 7 shows results including this additional baseline with α = 0 . 5 (equal weights to both the average and worst-case loss).
We observe that while interpolating between DRO and Static mixing is occasionally better than either on their own, it doesn’t achieve the same average/worst-case performance trade-off as L-DRO.
In Figure 8, we report CIFAR-100 results across 3 random seeds using a larger CNN model: 3 layers of 160 3 × 3 ﬁlters each followed by a ReLU activation and a 2 × 2 max-pooling layer, followed by a 2 layer multilayer perceptron with hidden dimension 320 and ReLU activations leading into a ﬁnal fully connected layer mapping to the 100 classes.
We ﬁnd that the difference between the various method vanishes.
Although L-DRO tends to achieve similar robust accuracy as DRO at higher average accuracy, the difference is small but not consistent across the different runs. C T IME AND M EMORY O VERHEAD OF L-DRO
The main overhead of L-DRO is the computation of the interaction matrix A .
The most straightforward approach for doing so is to compute the N task gradients and take their pairwise dot-product.
The overhead incurred is equivalent to taking N forward/backward passes (one to obtain each of the task gradients), as the cost of the pairwise dot-product is marginal in comparison.
In theory, this means that training with L-DRO becomes approximately N/T times slower, where T is the interaction matrix update interval.
We employed this strategy for the CIFAR-100 experiments in Section 5.1.
With N = 20 and T =in that setting, the theoretical slowdown is approximately × 1 . 2 .
In practice, for our biggest CNN model on CIFAR-100, training speed was around 27 steps/s for L-DRO vs 28.5 steps/s for ERM (for reference, PCGrad ran at 21 steps/s).
This means that the effective slowdown is approximately 1 . 06 .
We attribute this improvement over the theoretical limit to the optimized Jacobian computation in our
JAX implementation.
The downside of this ﬁrst approach is that it necessitates holding N gradients in device memory.
This can quickly become cumbersome when the number of tasks and the size of the model are large.
For the language modeling experiment in Section 5.2, we adopt the alternate —more memory efﬁcient— strategy of re-computing each dot-product separately.
This means the memory overhead drops to only holding 2 gradients in memory at each time, but we need to compute O ( n 2 ) backward passes.
In practice for the language modeling experiment, there N = 49 tasks and, and the interaction matrix was updated every T = 1000 steps, leading to a theoretical slowdown of approximately × 3 . 4 .
By optimizing the interaction matrix computation ( e.g. computing only the upper triangle, using a ( 2 × ) smaller batch size to estimate the gradients), the ﬁnal training speed was 2.2 steps/s for L-DRO, as opposed to approximately 3.4 steps/s for ERM, i.e. an effective slowdown of × 1 . 55 .
Note that another way to reduce device memory usage would be to use the ﬁrst approach and ofﬂoad the computation of the gradient dot-products to the CPU (where memory is much less of an issue), in which case the additional overhead comes from the cost of moving gradients from device to CPU as well as computing large dot-products on CPU.
However, we did not experiment with this last approach.