Balancing Average and Worst-case Accuracy in Multitask Learning
Authors: Paul Michel
Sebastian Ruder
Dani Yogatama
Venue: ICLR 2022 Submitted (2022)
Keywords: multitask learning
distributionally robust optimization

Abstract: 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 efficacy of L-DRO on a small controlled synthetic setting, we evaluate it on two realistic benchmarks: a multitask version of the CIFAR-100 image classification dataset and a large-scale multilingual language modeling experiment. Our empirical results show that L-DRO achieves a better trade-off between average and worst-case accuracy with little computational overhead compared to several strong baselines.

Under review as a conference paper at ICLR 2022
BALANCING AVERAGE AND WORST-CASE ACCURACY IN MULTITASK LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
When training and evaluating machine learning models 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 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 shellf 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 efficacy of L-DRO on a small controlled synthetic setting, evaluate it on two realistic benchmarks: multitask version of the CIFAR-100 100 image classification dataset and large-scale multilingual language modeling experiment. Our empirical results show that L- DRO achieves a better trade-off between average and worst-case accuracy with little computational overhead compared several strong baselines.
the
INTRODUCTION
Multitask learning the process by which a single model is trained to perform a variety of different tasks has become subject of increasing interest with many successful applications in variety of domains (Ruder, 2017; Yu et al., 2020; Wang et al., 2021). I). 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 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 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 ensure good worst-case performance. Rather, more natural choice is to use a type of min-max objective from the distribu- tionally robust optimization (DRO) literature (Rahimian ming max; (0). where f is the loss of the i-th task and 0 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).
Mehrotra, 2019; Sagawa et al., 2020):
Motivated by this observation, we propose novel optimization algorithm called Lookahead-DRe. (L-DRO). The intuition behind L-DRO is to anticipate the evolution 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 synthetic task, image classification task, and a language modeling task. In experiments on a multitask version of the CIFAR-101 -000 dataset, we show that L-DRO consistently allows to train models that achieve a good trade-off of worst and average task accuracy. On

---

Under review as a conference paper at ICLR 2022
realistic large-scale multilingual language modeling dataset, find that L-DRO is able to improve performance on the worst- persfrrerfoming languages at very little cost to average performance.
Our main contributions 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) propose 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 synthetic and realistic experiments from two different modalities.
follows: (i) show how distributionally robust optimization (DRO)
TRAINING MULTITASK MODELS
Consider a machine learning model parameterized by 0 RRdme which is trained perform N tasks. In general, let a task refer to the combination of examples drawn from data distribution and an objective function. For simplicity, we define task solely by its objective function 1,=..... As running example, we consider a multilingual language model with dozens of tasks where each task corresponds to a language (Conneau et al., 2020; Xue al., 2021).
2.1 AVERAGE TASK LOsS
The most common objective for multitask learning is to minimize the average all the task losses:
min 1 N Σ (0)
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 & varies wildly; Hessel et al. 2019) or we want to pay more attention certain tasks. If the distribution of tasks is not homogeneous, training with Lavg results in a model with performance that is skewed towards certain clusters of tasks. our multilingual example, if a large number of languages come from the same language family such Indo-Europe 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 {wi.... wy) in the probablility simplex AN = {w >> 0, E 1} to modulate the losses of individual tasks:
min
w. (0).
The weights are typically decided based on some statisties the data. proportional sampling (Sanh et al., 2019), the weight of task i is proportional to the size of its data u; [D,]- This biases the model towards tasks that are over-represented use an inverse temperature where w; x |D with typically in the range of [0.2, 0.3] (Aharoni et al., 2019; Conneau et al., 2020; Xue et al., 2021).
the data. To counter-balance this, recent approaches
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 the distributional robustness literature (Sagawa et al., 2020):
min max (,()) min max
w
WEAN
Coro(
2

---

Under review as a conference paper at ICLR 2022
0
0
Best Solution
W2
Ow
min (0%)
W3
W3
Figure 1: Overview of the Lookahead-DRO procedure. Starting from parameters Oo, the optimal training weights w are chosen such that the resulting parameters O", minimize the worst case loss.
The above objective minimizes the worst-case loss of the model every step. Alternatively, 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 CDRO 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 Gard is much more difficult than the others (i.e., ming ( (0) ming { for all other tasks j), then LDRO is likely to get "stuck" trying to minimize Paard at the detriment of other tasks.
In addition, DRO is agmostic to positive (or negative) interaction between tasks. some cases, training on one task & may have a positive effect on another task E because they share some similarities-ffo 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 one has a higher loss than the other. On the other hand, when two tasks are odds with each other (i.e., training on one task decreases performance the other), the max objective in LRoO can lead to oscillatory behavior where the weights w alternate between two competing tasks.
2.4 LOOKAHEAD DRO
The core limitation of CDRO underpinning the aforrementioned 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 given value of the model parameters Og. LDRO chooses weights maximizing the model's s loss w == arg max E; w;l,( (0o). After training with weights w we obtain parameters 0, with loss [DR (01) = maxw E: (0)). However, note that for oJ the worst-case weight allocation w := arg max Ex (;(((0)) might be different than w so we have no guarantees that (O.). In other words, training with the worst-case weights w does not guarantee
< (0)) that LDRO will decrease.
This observation suggests alternative strategy: suppose w, Ou Ow.er d, is fixed-ane and, for simplicity, deterministic- tetemmnistic-gradienn descent algorithm minimizing the w-weighted loss Σe w,l; starting from parameters 0g. Rather than choosing to be the worst-case weights corresponding to Op we should pick weights that minimize the worst-case loss of l w. Lee after training. This can be formalized solving the min-max problem:
min max WEAN VEAN
Σ Eerl((Owee))
(1)
Intuitively, decoupling the weights w used for training from the weights v used for computing LDRO(O", W. e.,) removes the restriction of just training the worst-performing tasks and enables us to find a task weighting which able to continuously minimize the worst-case loss. Note that in this case there might be multiple solution 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 training tasks which performance
for instance, it might be that one more tasks have the same
3

---

Under review as a conference paper at ICLR 2022
still be improved (as long worst-case performance decreases). This can be formulated 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" the optimization trajectory from O, to identify the DRO loss maxy E, v,,;(Ow.0,)) call this weight selection principle Lookahead- DRO (L-DRO). We provide a graphical illustration Figure 1.
2.5 ONE-STEP LOOKAHEAD DRO
In practice we cannot explicitly solve the min-max max game in Eq. 1 because estimating the payoff max, E, (022(0w.,)) w,0,) would require unrolling the full optimization process to compute A 0., In- stead, we simplify and take only one gradient step with a small learning rate X, i.e. replacing o .o with 0, - E, 0e (0o). Such one-step approximations are often used in the meta-lcarning litera- ture (Finn al., 2017: Liu al., 2018). Using the first order Taylor extension of l enables us to rewrite the individual lookahead task losses as:
4
W (o)())))) ~ ((0)
Σ " VC((0))VE
==1
The resulting payoffs are now linear in w, which reduces the lookahead optimization problem in Eq. into bilinear matrix game:
min max v L w
(2)
w
with payoff matrix L, called the "lookahead matrix", which can be decomposed follows:
Lig
(,(O))
Loss of task i
"interaction" between tasks and
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 À treated hyperparameter. This results in training algorithm which we outline in Algorithm
Algorithm 1: Lookahead DRO
Input: Lookahead learning rate initial parameters Oo
+ 0.
while training not over do
L - (0) (V),
// Compute lookahead matrix
v t m-xmaatn_naaate solver(L)
Find Nash equilibrium of 0 + optimizer_step (4) // imizeeer stop on w'-weighted loss
end
To build more intuitions, we consider several edge cases and discuss how L.-RRO handles them:
Independent tasks. We have Ve,VL, 1izj. 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; E, with the highest loss, and the only possible best response for the column player (w) is to match this selection (since Li=: f À < Lij f for all / 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" DRO that training one task does not affect the others.
Redundant tasks. We have V = 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 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 to is maximal, but the model has already converged task (in other words Vlg 0). In this scenario acenario-assuming there
4

---

Under review as a conference paper at ICLR 2022
is strong negative interaction among the remaining tasks- the highest loss will always be to irrespective of the training weights w. Again, by the maximum entropy principle, L-DRO advocates for picking uniform weights. This makes it possible continue training the remaining tasks though fo has converged, bypassing the core limitations of distributionally robust optimization (DRO) identified in $2.3.
LOOKAHEAD DRO IN PRACTICE
3.1 GRADIENT AND LOss NOISES
In most realistic machine learning scenarios we do not have access to the full losses (0) and their gradients. Rather, we compute stochastic estimates on minibatches (0). In preliminary experiments, find 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 fluctuations 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. (2020). Instead of computing the optimal weights w' each step based on noisy losses and gradients, we keep a running set of weights w, which we update jointly with model parameters O. Specifically, we interleave gradient updates 0 with exponentiated gradient steps (Kivinen & Warmuth, 1997) on w to minimize max, v L w. We show the weight update rule in Algorithm 2.
Algorithm 2: Lookahead DRO online weight update
Function update weights (Lookahead matrix L, weights w, weight update rate 17): V - arg max, vT L w
17 Compute best response to w
Wh 4 w.e - L];
// Update weights
// Re-normalize weights
return
In practice, we find that this online weight update alleviates the instability issues mentioned above. Note that the choice of exponentiated gradient is not arbitrary. It corresponds to mirror descent step divergence (Kullback & Leibler, 1951) as 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).
(Nemirovskij & Yudin, 1983) using the Kullback-
3.2 LOOKAHEAD MATRIX COMPUTATION
The main computational bottleneck of L-DRO lies in the computation of the lookahead matrix L. Recall that Lij can be decomposed in two terms: the loss & and an "interaction" term TV), While the former can be estimated essentially "for free" during the forward pass of back propagationn computing the pairwise task gradients dot products is much more computationally demanding. Indeed, it necessitates computing (and holding in memory) model gradients, which be prohibitive for larger models. To reduce the computational overhead, we only update this second term (which we call the "interaction matrix" A (O) V (O)),,==....)) periodically, typically every 100 or 1000 training steps. In preliminary experiments. we also find that Lookahead- -RRO works more reliably when the lookahead gradients of each task are re-normalized to one. This means that the interaction matrix is actually computed be A, (O) these two design choices are investigated in ablation studies in Appendices A. and A.2. We show the final algorithm for online L-DRO in Algorithm
The effect of
4 SYNTHETIC EXPERIMENTE AS A PROOF OF CONCEPT
In this section, we demonstrate the efficacy of L-DRO synthetic multitask setting. We design toy experiment such that neither DRO nor static mixing is able to reach an optimal solution (Figure 2). We identify each task with point yı on the two dimensional plane R² and define their respective loss functions as the squared Euclidean distance of the model prediction f(0 to yis (0) = | f(0)
We distribute the tasks such that most points are clustered around (-1, 0), to simulate grouping of similar tasks (e.g. Euroouuno languages in our multilingual example). We then place a single point Youtier (1, 0) emulate the presence of outlier task.
5

---

Under review as a conference paper at ICLR 2022
Algorithm 3: Online lookahead DRO training loop
Input: Lookahead learning rate X, weight update rate y,, interaction matrix update interval
A 1-0
// Initialize interaction matrix at // Initialize w as uniform weights
w - ( (...))
while training is not over do
1% = then
A (O) [V Isijsn
// Update interaction matrix
L + (0) AA
// Compute lookahead matrix Update weights using Algorithm 1/ Opt zer zer step on w-weighted loss
< optimizer_step Σ,
w, w
We simulate the increased difficulty of the outlier task by adopting a parameterization s(0) that prevents the model from reaching any position closer than a fixed value I from V/outier. Specifically we adopt polar parameterization 0 := (r, Φ) and set
cos
((0) Youtlier (R +
sin σ
This essentially creates "forbidden region" of radius R around Youtlerr of the space that the model cannot penetrate, thus lower- twerrbbunding foutier to R. In experiments, we set = 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 the outlier. On the other hand, the DRO trajectory moves towards the latter but ends up stuck when it reaches a point where foutier cannot be improved despite being the highest loss among all tasks.
On the other hand, while the L.-RRO trajectory shown in Figure 2b follows a similar trajectory as DRO at first, it able to reach solution that still achieves the best possible foutier while achieving lower loss on the other tasks. Also notice that the online variant follows a similar trajectory, despite the fact that does not compute 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 both losses and gradients. Under these conditions, lookahead DRO becomes very unstable, an issue that is partially mitigated by the online update.
EXPERIMENTS
5.1
CIFAR100
As a first step towards a more realistic setting, we experiment with a multitask version of the CIFAR- 100 dataset (Krizhevsky al., 2009), in which the 20 coarse labels are treated separate 5-way classification tasks (Rosenbaum et al., 2018; Yu et al., 2020). In this setting, we train a small CNN model with four convolution layers with 32 x 3 filters each followed by ReLU activation and 2 x 2 max-pooling layer. The resulting representation is averaged and fed through 3-layer fully connected feed-forward network with hidden dimension 128 and ReLU activations. Predictions are made using 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 (2020). Finally we compare to PCGrad (Yu al., 2020), 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 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.001, 0.01, 0.1}. For L-DRO, sweep over 3 values of the weight update rate (0.01, 0.001, 0.0001} and values of the lookahead learning rate {0.001, 0.01, 0.13. For Static mixing and PCGrad, there are additional parameters sweep over.

---

Under review as a conference paper at ICLR 2022
-1
Average
LORO Orinn
LORD
-2
On L-RRD
DRO vs. Average loss
(b) L-DRO
(c) L-DRO (with noise)
Figure 2: Synthetic multitask experiment. The outlier task is shown in red while the others are in yellow. The grayed out area represents the region space that cannot be attained by the model with our chosen parameterization
To gain a better understanding of the various trade-offs of each approach terms of average and worst accuracy, we do not report a single number. Rather, we periodi- cally evaluate the trained models on the test set and report the Pareto frontier of all checkpoints (across epochs and hyper-parameter parameter configurations) of each method in the two dimensional plane of worst and average accuracy.
WorstrAverage accuracy parato frontiars MT-CIFARIOC
L-ORC
DMO
Average
#CGred
We show the results in Figure 2. (in blue) consistently achieves good trade-off between the two metrics, as its Pareto frontier is situated the upper right corner of the space (high average and worst- case accuracy). In particular it is able to reach similar worst-case performance as DRO, without incurring a hit in terms of average performance. We find that this trend is consistent across different random seeds (see Appendix B.
We can see that L-DRO
acculidity
Figure 3: Average/worst task accuracy Pareto frontiers on multitask CIFAR-IO0
An important difference in our experimental setting with Yu et al. (2020) is that the model we train comparatively much smaller. In fact, 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 finding suggests that L-DRO is primarily useful when the model's capacity a limiting factor compared to the number of tasks and the difficulty of individual tasks, which is not necessarily the case for the simple CIFAR-100 benchmark.
5.2 MULTILINGUAL LANGUAGE MODELING
To evaluate L-DRO in real-world scenario, we turn to large-scale multilingual language modeling. As NLP models are increasingly applied 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 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 large data imbalance across tasks: the highest resource language (English) contains 55x more data than the lowest resource language (Yorùbá) in our training corpus. Additionally, there is a difference in difficulty between the various languages. Formally, for a parametric language model pa. its loss the i-th language is the cross-entropy Hp,,, Pe) with the underlying distribution of the language, P:- This can be re- writteet H(p:) DKL (p: [pe). Since the KL divergence is non-negative, the loss by the language-specins entropy term H(p.). irrespective of the
of each language is capacity the fitness of the model. In other words, languages with higher entropy will be more difficult to learn (Mielke et al., 2019).
5.2. EXPERIMENTAL SETTING
We train a large autoregressive transformer model (Vaswani et al., 2017) on the multilingual C4 dataset (mC4; Xue et al., 2021). We limit ourselves to the 40 languages that are present in the
7

---

Under review as a conference paper at ICLR 2022
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 auro-eu and 42.5% are written in the Latin script. 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 validation and test set containing up 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 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:
as 0 (uniform sampling), 0.3 (a common value in the literature; Xue et al. 2021) and 1 (purely proportional sampling)
Static mixing ($2.2). We try different values for the inverse temperature
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 1j E {0.01, 0.001, 0.0001]. Baselined-DRO, variation on DRO proposed by Oren et al. (2019) where the losses of each language is adjusted with a "baselinc" loss to account for the entropy term which biases DRO towards selecting languages that inherently more difficult. In practice, we estimate (p:)) by training a separate monolingual model on each language, po; and using its cross-entropy Hp;: ) = ming Hp,,, (e) as a proxy. We then train to minimize max, (0) - H(pp.
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.2 RESULTS
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 monolingual model. Specifically, the relative perplexity of a model 0 on language is defined as:
Perplexity 0, 1) - Perplexittxx Perplexity(O:; ()),
Relative perplexity(0.
where 0 is monolingual model trained on language i. In other words, relative perplexity of given model on a language quantifics 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 (a 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 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 terms of worst-case relative perplexity, with a comparatively much smaller computational overhead (periodic computation of L), which makes it
In experiments, we omit the Tagalog language, which is not included mC4. We do not compare against PCGrad due to its underwhelming performance in the 100 experiments and the prohibitive cost of its gradient projection step in this setting.
8

---

Under review as a conference paper at ICLR 2022
absolute perplexity
relative perplexity
average
worst
average
worst
Static Mixing (a 1)
48.7
135.5
358.3
1170.7
Static Mixing (a 0.3)
24.0
41.2
130.6
213.7
Static Mixing (a 0)
23.9
48.7
128.2
198.1
DRO
24.5
35.1
142.7
286.7
Baselined-
23.9
50.7
127.3
168.5
L-DRO
23.6
55.5
126.0
191.7
Monolingual
10.6
25.3
0
0
Table Results on the multilingual language modeling experiment. We report results both in terms of absolute and relative perplexity.
much more appealing alternative. A more detailed discussion of the computational overhead of L-DRO can be found in appendix C.
6 RELATED WORK
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-ohjective optimization, with the goal of finding a Pareto-optimal solution. Other work proposes to directly modify the gradients. In order to minimize gradient interference, Yu et al. (2020) project conflicting gradients on the normal plane of the other. Wang et al. (2021) alter both the direction and magnitude of gradients to achieve certain gradient similarity.
Many other methods have been proposed to weight the task losses according to different criteria such as uncertainty (Kendall al., 2018), reward magnitude (Hessel et al., 2019), and learning speed (Chen et al., 2018; Liu et al., 2019; Zheng et al., 2019). Most 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 8; Jean et al., 2019) weight task losses based on performance, focusing on optimizing tasks with poor performance. Among these, weighting- methods. Such methods, however, do not consider the optimization procedure and 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 predefined sampling schedules while Sanh et al. (2019) proposed a proportional sampling strategy. multilingual language modeling and machine translation (Aharoni al., 2019; Conneau et al., 2020; Xue et al., 2021), proportional sampling with temperature is typically employed.
proposed method most similar in spirit to performance
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 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 mixture of sub-populations (Oren et al., 2019; Sagawa et al., 2020; Zhou et al., 2021).
7 CONCLUSION
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 difficult than others) or when there is an imbalance towards certain types of tasks. We proposed L-DRO, which enables efficient optimization of the worst-case loss without the pitfalls of existing DRO-based max-loss minimization approaches. We demonstrated the benefits of L-DRO in a synthetic scenario as well as in realistic image classification 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 multitask model.

---

Under review as a conference paper at ICLR 2022
ETHICS STATEMENT
Measuring models only based on their average performance obfuscates their limitations on under- represented 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.
REPRODUCIBILITY STATEMENT
We take several steps to ensure reproducibility of our work. Pseudo-code describing the final version of our algorithm given in Algorithm Moreover, we detail experimental settings and model hyper-paramettes in $5. Finally, all our experiments are performed on datasets that are openly available: CIFAR-I00' (Krizhevsky et al., 2009) and mC4 (Xue et al., 2021).
REFERENCES
Roce Aharoni, Melvin Johnson, and Orhan Firat. Massively multilingual neural machine translation. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Com- putational Linguisties: Human Language Technologies (NAACL-HL.T), pp. 3874-3884, Minneapo- lis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 11653//1//19- 1388. URL https:ttps........t.. acl www.annnollooooooooooooooooooooo ogy. org/Nl /N1 9- 1388.
Aharon Ben-Tal, Laurent El Ghaoui, and Arkadi Nemirovski. Robust optimization, volume 28. Princeton University Press, 2009.
George W Brown. Iterative solution of games by fictitious play. Activity analysis of production and allocation, 13(1):374-376, 1951.
Joy Buolamwini and Timnit Gebru. Gender shades: Intersectional accuracy disparities in commercial gender classification. In Conference on fairness, accountability and transparency, pp. 77-91. PMLR, 2018
Zhao Chen, Vijay Badrinarayanan, Chen- Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In International Conference on Machine Learning, pp. 794-803. PMLR, 2018.
Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Francisco Guzmán, Edouard Grave, Myle Ott, Luke Zettlemoyer, and Veselin Stoyanov. Un- supervised cross-lingual representation learning at scale. In Proceedings of the 8th Annual Meeting of the Association for Computational Linguistics (ACL), pp. 8440-845) 8451, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/v 11665//1/22200... 1/c12/..ccc-----i. aal-mainn.741 URL org/ 2020 acl mecc-mainn.. 747.
Michael Crawshaw. Multi-task learning with deep neural networks: A survey, arXiv preprint N0V::20009.0999, 2020.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Proceedings of the 34th International Conference on Machine Learning (ICML), pp. 111111111111 1135. PMLR, 2017.
Rui Gao and Anton J Kleywegt. Distributionally robust stochastic optimization with wasserstein distance, arXiv preprint 1604:0604 2016.
Michelle Guo, Albert Haque, De-An Huang, Serena Yeung, and Li Fei-Fei. Dynamic task prioriti- zation for multitask learning. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 270-287, 2018. 2018.
http://ww....oooooooeeetttt//t/e 'https://wwe httpaa///waa...eerrrroooooooaaaa
ingual
10

---

Under review as a conference paper at ICLR 2022
Matteo Hessel, Hubert Soyer. Lasse Espeholt, Wojeiech Czamecki, Simon Schmitt, and Hado van Hasselt. Multi-task deep reinforcement learning with popart. In Proceedings of the 33rd Meeting of the Association for Advancement of Artificial Intelligence (AAAI), pp. 5796-3803, 2019.
Junjie Hu, Sebastian Ruder, Aditya Siddhant, Graham Neubig, Orhan Firat, and Melvin Johnson. XTREME: A Massively Multilingual Multi-task Benchmark for Evaluating Cross-lingual Gen- cralization. In Proceedings of the 37th International Conference on Machine Learning (ICML), 2020.
Zhaolin Hu and L. Jeff Hong. Kullback-leibler divergence constrained distributionally robust opti- mization. Available at Optimization Online, 2013.
Edwin Jaynes. Information theory and statistical mechanics. Physical review, 106(4):620, 1957.
Sébastien Jean, Orhan Firat, and Melvin Johnson. Adaptive scheduling for multi-task learning. arXiv preprint urXiv:1909.064: 2019.
Fred Jelinek, Robert Mercer, Lalit R Bahl, and James K Baker. Perplexity-a measure of the difficulty of speech recognition tasks. The Journal of the Acoustical Society of America, 62(S1): S63-S63, 1977.
Guy Jumarie. Relative information what for? In Relative Information, pp. 1-11. Springer, 1990.
Jagat Narain Kapur and Hiremaglur' K Kesavan, Entropy optimization principles and their applications. In Entropy and energy dissipation in water resources, pp. 3-20. Springer, 1992.
Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. Proceedings of the 31st IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 7482-7491, 2018.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proceedings of the International Conference Learning Representations (ICLR), 2014.
Eliyahu Kiperwasser and Miguel Ballesteros. Scheduled multi-task learning: From syntax to translation. Transactions of the Association for Computational Linguistics, 6:225-240, 2018.
Jyrki Kivinen and Manfred K Warmuth. Exponentiated gradient versus gradient descent for linear predictors. information and computation, 132(1):1-63, 1997.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009.
Taku Kudo. Subword regularization: Improving neural network translation models with mul- tiple subword candidates. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 66-75, Melbourne, Australia, July 2018. Association for Computational Linguistics. doi: 10. 10.18653/11//18-1007. URL https: htt///aaaaaatltooaoeeeeeeea org/P18- 1007.
Solomon Kullback and Richard A Leibler. On information and sufficiency. The annals of mathemati- cal statistics, 22(1):799-5) -6, 1951.
Carlton Lemke. Bimatrix equilibrium points and mathematical programming. Management science, 11(7):681-689, 1965.
Daniel Levy, Yair Carmon, John C Duchi, and Aaron Sidford. Large-scale methods for distributionally robust optimization. Proceedings of the 34th Annual Conference on Neural Information Processing Systems (NeurIPS), 33, 2020.
Xi Lin, Hui-Ling Zhen, Zhenhua Li, Qing-Fu Zhang, and Sam Kwong. Pareto multi-task learning. Advances in neural information processing systems, 12:12000-12070, 2019.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. In
Proceedings of the International Conference Learning Representations (ICLR), 2018.
11

---

Under review as a conference paper at ICLR 2022
Shengchao Liu, Yingyu Liang, and Anthony Gitter. Loss-balanced task weighting to reduce neg- ative transfer in multi-task learning. In Proceedings of the 33rd Meeting of the Association for Advancement of Artificial Intelligence (AAAI), pp. 9977-9978, 2019.
Pingchuan Ma, Tao Du, and Wojciech Matusik. Effcient Continuous Pareto Exploration in Multi-Taan Task Learning. In Proceedings of ICML 2020, 2020.
Bryan McCann, Nitish Shirish Keskar, Caiming Xiong, and Richard Socher. The natural language decathlon: Multitask learning as question answering. arXiv preprint arXiv:1806 08730, 2018.
Sabrina J. Mielke, Ryan Cotterell, Kyle Gorman, Brian Roark, and Jason Eisner. What kind of language is hard to language-model? In Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, pp. 4975-4989, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 18653/ URL httpht:///// /aae:///clll aacchnnnhhoggyyy org/P19-22
Arkadij Semenovič Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. Wiley- Seley-Interscience Series in Discrete Mathematics, 1983.
Yonatan Oren, Shiori Sagawa, Tatsunori Hashimoto, and Percy Liang. Distributionally robust language modeling. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 4227-4237, 2019. URL https: /
org/
Nathologyy/D1--4433
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the Limits of Transfer Learning with a Unified Text-to-Tex: Transformer. Journal of Machine Learning Research, 21, 2020. URL http:
//arxiiv.rrg///.......a hbspg//bb/1999 10683.
Hamed Rahimian and Sanjay Mchrotra. Distributionally Robust Optimization: A Review. arXiv preprint a90iv: XXXx::908.00659, 2019.
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning. In Proceedings of the International Conference on Learning Representations (ICLR), 2018.
Sebastian Ruder. An Overview of Multi-Task Learning in Deep Neural Networks. In arXiv preprint 170500/706,00088, 2017.
Sebastian Ruder. Neural transfer learning for natural language processing. PhD thesis, NUI Galway, 2019.
Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In Proceedings of the International Conference on Learning Representations (ICIR), 2020. URL
https:/ aarxvv. org/ pdtp//d///// /9t.//91..... 08731. pdf.
Victor Sanh, Thomas Wolf, and Sebastian Ruder. hierarchical multi-task approach for learning embeddings from semantic tasks. In Proceedings of the 33rd Meeting of the Association for Advancement of Artificial Intelligence (AAAI), pp. 6949-6956, 2019.
Ozan Sener and Vladlen Koltun. Multi-Task Learning Multi-Objective Optimization. In Proceed- ings of NIPS 2018, 2018
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31st Annual Conference on Neural Information Processing Systems (NIPS), pp. 5998-6008, 2017.
Alex Wang, Amapreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue benchmark leaderboard. http://wwubenncmaaaa..oo///aaaaa //gl uebenchmart rk. wwwm/eaaeebbaa...ooo.o rd. Accessed: 2021-09-24.
2019. URL https:
12

---

Under review as a conference paper at ICLR 2022
Zirui Wang, Yulia Tsvetkov, Orhan Firat, and Yuan Cao. Gradient Vaccine: Investigating and Improving Multi-task Optimization in Massively Multilingual Models. In Proceedings of ICLR 2021, 2021.
Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami AI-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel. mTS: A massively multilingual pre-trained text-to-texi transformer. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT), pp. 483-498, Online, June 2021. Association for Computational Linguistics. doi: 10. 8653/v1/202/ .na.c-n maacl-main.41 41. URL https::tp/...m / acclcclcccccooco.ccc antho ogy org/ 2021 naac ma in ii..... 41.
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. In Proceedings of NeurIPS 2020, 2020.
Feng Zheng, Cheng Deng, Xing Sun, Xinyang Jiang, Xiaowei Guo, Zongqiao Yu, Feiyue Huang, and Rongrong Ji. Pyramidal person re-identification via multi-loss dynamic training. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8514-8522, 2019.
Chunting Zhou, Xuezhe Ma, Paul Michel, and Graham Neubig. Examining and combating spurious features under distribution shift, In Proceedings of the 38th International Conference on Machine Learning (ICML), 2021.
13

---

Under review as a conference paper at ICLR 2022
Wortshverage FAL100
NT-CIRR100
MI-CIFAALCO
23
15
1.5
35
14
Morit tast ECCuTRCy
(4or2 bast CCCuTRCy
Wondl wask accurscy
Figure 4: Average/worst task accuracy Pareto frontiers on multitask CIFAR- 100 for L-DRO with different interaction matrix update intervals. Each plot corresponds to different random seed.
A ABLATION STUDIES
A.1 EFFECT OF PERTODIC MATRIX UPDATES
We investigate the effect of periodic matrix updates on the results L-PDRO.. Figure 4 reports the Pareto frontier of L-DRO over range of dlyeer-parameee 1 and À the CIFAR-I00 experiment from Section 5.1. We compare frontiers for different interaction matrix update intervals: 10, 100 and 1000, and report results for 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, have observed some instability with the larger update intervals. This is evidenced in Figure 5a for CIFAR-100 run with the larger CNN model used in Section B.3.
A.2 EFFECT OF RENORMALIZATION
Note that there of the two components each element the interaction matrix A, = Ve and we only normalize one VV + V V4
This can be interpreted the first order approximation of taking lookahead step along the *normal- ized* task gradients:
V (0,) [ 1) ~ (((0)) - x Σ wy VL,( Ve,(0) [2,(0)0]
W
j=1
Our intuition with this re-normalization trick is to help mitigate the instability caused by the "stal- eness" of the interaction matrix with periodic updates. Indeed, the approximation error of using stale interaction matrix comes from two sources: outdated alignment between the gradient (cosine similarity) and outdated gradient norms. With the re-normalization trick, 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, an ablation study CIFAR-IOS (see Figure 5b), find that the use of the re-normalization trick removes the instability observed with larger update intervals (1000).
B ADDITIONAL CIFAR-100 RESULTS
B. MULTIPLE RANDOM SEEDS
Figure depicts results on the CIFAR-100 dataset for 3 different random seeds. We find that L-DRO consistently outperforms baselines.
14

---

Under review as a conference paper at ICLR 2022
(a) Without renormalization.
(b) With renormalization.
Figure 5: Worst-case accuracy training curves for CIFAR-100 with f) 1, À 0.001 and interaction matrix update intervals 10 (blue), 100 (orange) and 1000 (green). The curves are smoothed with running exponential average for readability.
Wensokverage accuracy Dere3o franbirrs CIAANOO0
WosstAnsrag accuracy pareto frontiors
Versstsvveaage accuracy pereto frorbers MIFARAA100
32
#cGend
26-
Figure 6: Average/worst task accuracy Pareto frontiers across random seeds on multitask CIFAR-
100
B.2 ADDITIONAL BASELINE
Another natural baseline for trading off worst and average accuracy is to optimize a weighted average of average and worst-case task loss:
N 4(0) +1 a1 x max (0)
min Σ 111
with some interpolation parameter A downside of this approach is that, for a O, the model cannot, in general, reach the minimal worst-case loss because 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 their own, it doesn't achieve the same average/worst-case performance off as L-DRO.
B.3 EXPERIMENTS WITH LARGER CNN
In Figure 8, we report CIFAR-100 results across 3 random seeds using larger CNN model: 3 layers of 160 by 2 layer multilayer perceptron with hidden dimension 320 and ReLU activations leading into a final fully connected layer mapping to the 100 classes. We find that the difference between the
filters each followed by ReLU activation and a max-pooling layer, followed
Worsthkuerage
Wonstttvergge on MT-CIRAFLO0
CITA5100
CIFAIIRR100
Lazouhoodd ERO
Lockatsed DAD
Lostuhccd Eno
20
13
Figure 7: Average/worst task accuracy Pareto frontiers across random seeds on multitask CIFAR- 100, including additional baseline.
15

---

Under review as a conference paper at ICLR 2022
on MII-CIPALLOO
on NI-CIAAM
on RRR-GIAAAO
Instchheed tak
DeG Stale roong iXbrad
accurecy
Morti sccumayy
Skans acouaayy
Figure 8: Average/worst task accuracy Pareto frontiers across random seeds on multitask CIFAR- 100
various method vanishes. Although L-DRO tends to achieve similar robust accuracy as DRO higher average accuracy, the difference small but not consistent across the different runs.
C TIME AND MEMORY OVERHEAD OF L-DRO
The main overhead of L-DRO is the computation of the interaction matrix A. The most straightforward approach for doing 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 the task gradients), 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- 1 00 experiments in Section 5.1. With N = 20 and 7 100 in that setting, the theoretical slowdown is approximately x1.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 first approach is that it necessitates holding N gradients in device memory. This can quickly become cumbersome when the number of tasks and the size the model are large. For the language modeling experiment in Section 5.2, we adopt the alternate mmore memory efficient- 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²) 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 theoretical slowdown of approximately x3.4. By optimizing the interaction matrix computation (e.g. computing only the upper triangle, using a (2x) smaller batch size to estimate the gradients), the final 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 x x1.55.
Note that another way reduce device memory usage would be to use the first approach and offload 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.
16


Reviews:
[{"review_id": "cMFb0EIVK6q", "reviewer": "Reviewer_j3ej", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n1. The idea sounds legit, interesting and novel.\n\n2. The writing is clear and easy to follow.\n\nWeaknesses:\n\n1. The empirical improvement does not seem to be very significant, especially on large models.\n\n2. Another concern is with all the approximations, it is no longer clear whether the merit of the proposed method is still held, i.e., can the proposed L-DRO guarantee that $L_{DRO}(\\theta_1) < L_{DRO}(\\theta_0)$, empirically?\n\nQuestions:\n\n1. Does Figure 5 imply the advantage of the proposed L-DRO is mitigated by model with larger capacity?\n\n2. How large is the variance of the experiments? Is 3 runs enough?\n\n3. Can the authors visualize the weights in algorithm 2? Will the trick make the weight $w$ almost consistent across the training?\n\n---- Post Rebuttal\n\nI thank the authors for the response. I would like to keep my original rating as weakly accept, with the reason that 1) I think the novel viewpoint could potentially benefit the community and 2) the diminishing performance of the proposed method when the model gets larger makes the approach less interesting. ", "summary_of_the_paper": "This submission proposes Lookahead-DRO, of which the core idea is to do DRO on the updated weights, i.e., to choose a weight that minimizes the DRO loss in the next step. The computation is not tractable so the authors propose to use first order Taylor expansion approximation. Experiments on synthetic datasets and two real-world datasets demonstrate its effectiveness.", "summary_of_the_review": "Interesting idea, good writing, moderate empirical results", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "HR7r7OGqm36", "reviewer": "Reviewer_TtaP", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths:** The problem is very important. The authors bring the attention of the community of an important shortcoming of both average loss and worst case loss in the settings when there is imbalance in task distribution. The paper for most part is well written. \n\n**Weaknesses:** I have several concerns that I highlight below: \n1. **Look ahead DRO and role of deterministic algorithm** The exact approach of look ahead DRO is presented in Section 2.4. The approach has several problems that are not clear. If we assume that the deterministic algorithm finds $\\theta_{w,\\theta_0}$ achieves the global minimum for the weight $w$, then I cannot see how this approach is any different from doing regular min-max optimization. The reason I say is that if you do weighted sum risk minimization, then the loss values for different tasks achieved at the optima for different values of weights cover the entire Pareto boundary of the graph of loss of tasks. The problem amounts to selecting the point on the Pareto boundary that achieves the best worst case risk. \nThe authors should show using at least simple 2D illustrations, what point on the Pareto plane DRO achieves, and what point exact look ahead DRO intends to achieve, and contrast with standard sum risk minimization. \nI suspect the exact lookahead DRO is implicitly assumed that it will not find the exact minimum. If that is the case, then there should be a very detailed discussion on it. \n\n2. **Two baseline approaches** The title of the paper is about balancing worst case accuracy and average accuracy. The title itself suggests that the authors should compare with a simple combination of minimization of worst case together with average risk. This can be stated as \n$$min_{\\theta}\\Big( max_{k}l^{k}(\\theta) + \\lambda \\sum_{k} l^{k}(\\theta)\\Big)$$\nThis gradients of the above can done with the approach of either Oren et al. and also perhaps Sagawa et al.\nThe above objective exactly tries to optimize what the authors require. It is not clear why there are no comparisons with the above approach since it is the most natural baseline to study. \n\nThe objective in equation (8) in http://proceedings.mlr.press/v139/krueger21a/krueger21a.pdf can also solve the problem that the authors propose. The objective minimizes the average loss across tasks and regularizes it using the variance of the loss across tasks. Hence, if for each task if you compute the loss as the average loss over the data by dividing it by the size of the data for that task, this would be a useful approach to compare against. \n\n3. **Regarding experiments** In the main real experiment, the baseline-DRO performs the best. The authors claim that its much more computationally expensive. I have two main points to make here\na) Baseline DRO is still a DRO based approach. Why is it able to achieve a much better tradeoff? The entire paper up until this section says that you will achieve a better tradeoff irrespective of the compute budget. This seems to be going against the main idea presented. Even your synthetic experiment at least agrees with your main hypothesis. \n\nb) You claim that the computational expense is much higher. If that is the case and that is the main source of gain, why not report the wall clock computation times to make the point clear? There is one more reason I ask this is because in Oren et al. (Section 4.2) the authors state \" In practice, we use a bigram model, which was fast enough to scale and worked sufficiently well in experiment\". What you say seems to contradict their claim. \n\n\n\n\n\n\n ", "summary_of_the_paper": "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.\nThe authors carry out synthetic and real datasets to demonstrate the efficacy of the approach.", "summary_of_the_review": "Overall, the paper addresses an important problem. It proposes an interesting approach to solve it as well. The synthetic experiment is promising. However, there are quite a few weaknesses -- the target that exact lookahead DRO approach seems to be no different than regular DRO (as I explained above), missing natural baselines (without such baselines its hard to even understand if such a lookahead is needed or not),  problems with real data experiments. If the authors can effectively address these concerns, I would be willing to change my score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6F_ij5Wldmd", "reviewer": "Reviewer_Bw6i", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, the idea of this paper is well presented and intuition is well explained. Below please find my detailed comments.\n\nMajor:\n1. Seems for most cases when the worst task’s loss is larger than other losses with a certain margin. This algorithm only performance gradient descent on this worst task. \n\nThe advantage of this algorithm only appears when the top-n worst tasks are very close to each other. In this case, considering task interactions might help. But I was wondering whether such a scenario happens often in practice? \n \n2. Seems there is a significantly simpler algorithm functioning similar to the proposed one:\nAt each iteration, we find the worst task, performance gradient on this worst task. If the gradient is not zero, we apply the gradient. Else, we apply the uniformly weighted gradient of all the other tasks.\n\nThere might be a slight difference between this very simple algorithm and the proposed one. I was wondering how do they compare.\n\n3. No convergence guarantee is shown here. I.e., does this algorithm converges to the Pareto stationary set? What’s the convergence rate interns of worst-case loss?\n\nMinor:\n1. There are various criteria to decide the final objective for multi-task learning. There is not much explanation why minimizing the worst loss is a good choice.\n", "summary_of_the_paper": "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.", "summary_of_the_review": "I think the general idea of this paper is interesting. However, there are still several important points that need to be clarified. I am happy to have more discussion and adjust my score accordingly if things are properly addrssed.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "5n9LP5tMtA", "reviewer": "Reviewer_NzDd", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I think the submission is novel and addresses an important question in multi-task learning. I believe the algorithm has a clear theoretical motivation but also want to note that the author makes a key simplification (i.e. the first order approximation) in estimating the task interaction, which most likely does not hold for non-convex neural network optimization. My main concerns are with the empirical evaluation of the proposed method, which are detailed below. Currently, my stand for the submission is \"marginally below the acceptance threshold\" but I can see the potential of the proposed method and could be convinced otherwise if the authors can help answer my questions/requests.\n\nI have the following list of questions:\n\n1. In practice, the loss values from different tasks can be very different. For example, if we learn a neural network to perform image recognition and segmentation simultaneously, the loss values will be of different scales from these two tasks. What would the authors recommend for such use case?\n2. The authors propose to renormalize the gradients into one when estimating the \"interaction matrix\" but provides little intuition and empirical validation so I am not convinced that this is a principled technique. Normalizing the gradients breaks the analysis section 2.5 which creates a gap between the conceptual motivation and the actual algorithm, which is quite disappointing. Can the authors justify this decision? Also, I would like to see ablation study on this design decision. \n3. The authors propose to reduce the computational cost by estimating the interaction matrix periodically. Can I see an ablation study on this hyper-parameter to understand the trade-off here? How much performance do we lose when we set the estimation period to be 1000 steps?\n4. The CIFAR100 experiment in the appendix (Figure 5) is quite concerning. I appreciate the honest documentation of this experiment. I wonder if this is because weighting approach in general does not work with overparametrized networks [1]? Can the authors provide more insights on this experiment?\n5. Can the authors document the actual wall clock time and memory overhead of the proposed method on their CIFAR and Multilingual LM experiments? \n6. For Table 1, the difference between absolute perplexity and relative perplexity is quite confusing. We observe that DRO achieves the best worst case absolute perplexity but not so much for relative perplexity. I am not entirely sure the \"relative perplexity\" is the correct metric: perplexity is defined as the inverse of the average probability of predicting the right word so computing the difference and the division does not have a clear definition. I understand the difficult in comparing perplexities across languages. Can the authors provide a table of the full list of perplexity values on different languages in the Appendix so I can better judge the results?\nAside from my current confusion mentioned above, the empirical gain of L-DRO over uniform sampling (static mixing with alpha=0) is quite marginal, which is concerning.", "summary_of_the_paper": "This paper studies balancing average case and worst case performance in a multi-task setting. The authors propose a new algorithm that extends the group DRO algorithm by 1. ) taking task interactions into account 2.) enforcing the maximum entropy principle when breaking ties. The interactions between tasks are modeled by a first order approximation to change in training loss after a single gradient update. Empirically, the authors propose methods to stabilize their estimation of the \"task interaction\" estimation by a online weight update procedure and also introduce an additional technique that re-normalize gradients to unit norm before calculating the \"task interaction\" estimation. In addition, the authors observed that the proposed algorithm only needs to update the task interaction estimation periodically during training, which reduces the computational overhead. \n\nThe submission empirically test the proposed algorithm on CIFAR100 and Multilingual Language Modeling. On CIFAR100, the proposed method achieves improved performance for smaller-scale neural networks but does not achieve significant improvement on larger neural networks. On Multilingual Language Modeling, the proposed method achieves best average performance and second best worst case performance (trailing a computationally expensive baseline).", "summary_of_the_review": "Novel and interesting algorithm that studies a important problem in multi-task learning. Current empirical evaluation of the paper can be furthered improved. Several design decisions in the final version of the algorithm should be better justified. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: In general, the reviewers recognized the importance of the question and the innovation in the proposed algorithm, but do not seem to be super excited about the overall contribution of the paper. (One or two reviewers did not seem to respond authors' response after the AC's reminder.) The AC read the reviews and responses and observed that the main concern appears to be the empirical performance --- the improvements are not as strong for larger models or if more computational time is allowed. Modern models are indeed typically large, and it would be good to discuss this point more thoroughly. If the work's focus is limited resource setting, the paper might want to state that upfront. Indeed, one reviewer is still concerned post-rebuttal about a clock-time comparison. Given these considerations, the AC will recommend reject for the paper but encourage the authors to resubmit to a top venue conference after revising the paper.
Decision: Reject