Balancing Average and Worst-case Accuracy in Multitask Learning

Extracted source text · arXiv 2110.05838

\documentclass{article} % For LaTeX2e
\usepackage{iclr2022_conference,times}

% Optional math commands from https://github.com/goodfeli/dlbook_notation.
\input{math_commands.tex}


\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc}    % use 8-bit T1 fonts
\usepackage[colorlinks=true,linkcolor=red,citecolor=blue]{hyperref}%
\usepackage{url}            % simple URL typesetting
\usepackage{booktabs}       % professional-quality tables
\usepackage{amsfonts}       % blackboard math symbols
\usepackage{nicefrac}       % compact symbols for 1/2, etc.
\usepackage{microtype}      % microtypography
\usepackage{xcolor}         % colors
\usepackage{amsmath}         % colors
\usepackage{amssymb}        % colors
\usepackage{graphicx}
\usepackage{subcaption}
\usepackage{bbm}
\usepackage{wrapfig}
\usepackage{setspace}
\usepackage{enumitem}
\usepackage[acronym,shortcuts,acronymlists={hidden}]{glossaries}

\renewcommand*{\glstextformat}[1]{\textcolor{black}{#1}}

\usepackage[ruled,vlined]{algorithm2e}
\newcommand\mycommfont[1]{\footnotesize\ttfamily\textcolor{blue}{#1}}
\SetCommentSty{mycommfont}


\newcommand{\eg}{\textit{e.g.}}
\newcommand{\ie}{\textit{i.e.}}
\newcommand{\etc}{\textit{etc}}



\newcommand{\dycomment}[1]{\textcolor{brown}{\bf \small [#1 --DY]}}
\newcommand{\paul}[1]{\textcolor{magenta}{\bf \small [#1 --PM]}}
\newcommand{\seb}[1]{\textcolor{olive}{\bf \small [#1 --SR]}}
\newcommand{\vect}[1]{\mathbf{#1}}
\newcommand{\set}[1]{\mathcal{#1}}
\newcommand{\struct}[1]{\boldsymbol{#1}}
% \DeclareMathOperator{\v}{\mathbf{q}}
% \DeclareMathOperator{\w}{\vect{p}}
\DeclareMathOperator{\q}{\vect{v}}
\DeclareMathOperator{\vecv}{\vect{v}}
\DeclareMathOperator{\w}{\vect{w}}
\DeclareMathOperator{\btheta}{\bm{\theta}}
\DeclareMathOperator{\Ell}{\mathcal{L}}

\let\textv\v
\renewcommand\v{\TextOrMath\textv\vecv}

\newacronym{dro}{DRO}{distributionally robust optimization}
\newacronym{erm}{ERM}{empirical risk minimization}
\newacronym{p-dro}{P-DRO}{Parametric DRO}
\newcommand\dro{\text{\acrshort{dro}}}


\bibliographystyle{iclr2022_conference}

\title{Balancing Average and Worst-case Accuracy in Multitask Learning}

% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to LaTeX to determine where to break the
% lines. Using \AND forces a line break at that point. So, if LaTeX puts 3 of 4
% authors names on the first line, and the last on the second line, try using
% \AND instead of \And before the third author name.

\author{%
  Paul Michel\thanks{Work done during an internship at DeepMind} \\
  \'{E}cole normale superieure PSL\\
  Paris, France \\
  \texttt{pmichel31415@gmail.com} \\
  % examples of more authors
   \And
   Sebastian Ruder \& Dani Yogatama \\
   DeepMind \\
   London, UK \\
   \texttt{\{ruder,dyogatama\}@deepmind.com} \\
}

\iclrpreprintcopy % Uncomment for camera-ready version, but NOT for submission.


\begin{document}

\maketitle

\begin{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.
\end{abstract}

\section{Introduction}

Multitask learning---the process by which a single model is trained to perform a variety of different tasks---has become a subject of increasing interest with many successful applications in a variety of domains \citep{Ruder2017overview,Yu2020gradient_surgery,Wang2021gradient_vaccine}. By and large, multitask learning models are evaluated by reporting their average performance on all tasks \citep{mccann2018natural,glueleaderboard2020}. 
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 \citep{Hu2020xtreme}.
In computer vision, facial recognition models 
have been shown to perform worse on racial groups that 
are underrepresented in the training data \citep{buolamwini2018gender}.


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 \citep{rahimian2019distributionally,sagawa2019distributionally}: $\min_{\btheta}\max_i \ell_i(\btheta)$, where $\ell_i$ is the loss of the $i$-th task and $\btheta$ 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 (\S{\ref{sec:worst_case}}). 

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 (\S{\ref{sec:lookahead}}).

We evaluate L-DRO on a synthetic task, an image classification task, and a language modeling task. 
In experiments on a multitask version of the CIFAR-100 dataset, we show that L-DRO consistently allows us to train models that achieve a good trade-off of worst and average task accuracy. 
On a realistic large-scale multilingual language modeling dataset, we find that L-DRO is able to improve performance on the worst-performing languages at very little cost to average performance.

Our main contributions are as follows: (i) we show how distributionally robust optimization (DRO) can be applied to multitask learning to maximize worst-case accuracy; (ii) we highlight failure cases of the standard DRO formulation in the multitask setting; (iii) we propose a new algorithm, L-DRO that addresses DRO's deficiencies; and (iv) we demonstrate that L-DRO achieves a better trade-off between average and worst-case accuracy than strong baselines in a synthetic and two realistic experiments from two different modalities.

\section{Training Multitask Models}

Consider a machine learning model parameterized by $\btheta\in\mathbb R^{d_{\text{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 define a task solely by its objective function $\ell_{i=1\ldots N}$. As a running example, we consider a multilingual language model with dozens of tasks where each task corresponds to a language \citep{conneau-etal-2020-unsupervised,xue-etal-2021-mt5}.

\subsection{Average Task Loss}

The most common objective for multitask learning is to minimize the average of all the task losses:
\begin{equation*}
    \min_{\btheta}\frac 1 N \sum_{i=1}^N \ell_i(\btheta)
\end{equation*}
A downside of this objective is that all tasks are treated equally. In real-world scenarios,
we often encounter a scenario where task difficulty is not uniform (i.e., 
the magnitude of $\ell_i$ varies wildly; \citealt{hessel2019multi}) or we want to pay more attention to
certain tasks. If the distribution of all tasks is not homogeneous, 
training with $\Ell_{\text{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, $\Ell_{\text{avg}}$ will choose 
models which perform better on these languages at the expense of typologically distinct languages such as Mandarin or Finnish.

\subsection{Static mixing} \label{sec:static_mixing}

Rather than simple averaging, prior work in multitask learning often uses arbitrary weightings $\w=\{w_1,\ldots,w_N\}$ in the probablility simplex $\Delta_N=\{\w\mid w_i>0, \sum_iw_i=1\}$ to modulate the losses of individual tasks:
\begin{equation*}
    \min_{\btheta}\sum_{i=1}^N w_i \: \ell_i(\btheta).
\end{equation*}
The weights are typically decided based on some statistics of the data.
In proportional sampling \citep{sanh2019hierarchical}, the weight of a task $i$ is proportional to the size of its data $w_i \propto |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 $\alpha$ where $w_i  \propto|D_i|^\alpha$ with $\alpha$ typically in the range of $[0.2, 0.3]$ \citep{aharoni-etal-2019-massively,conneau-etal-2020-unsupervised,xue-etal-2021-mt5}.



\subsection{Worst Task Loss} \label{sec:worst_case}


A third option that considers the task on which the model performs worst is the min-max objective. We refer to this objective as \acrshort{dro} because of its prominence in the distributional robustness literature \citep{sagawa2019distributionally}:
\begin{equation*}
    \min_{\btheta}\max_i\ell_i(\theta)=\min_{\btheta}\underbrace{\max_{\w\in\Delta_N}\sum_{i=1}^Nw_i\ell_i(\theta)}_{\Ell_{\dro{}}(\btheta)}.
\end{equation*}
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 $\Ell_{\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 difficulty: if one task $\ell_{\text{hard}}$ is much more difficult than the others (i.e., $\min_{\btheta}\ell_{\text{hard}}(\btheta) > \min_{\btheta}\ell_{j}(\btheta)$ for all other tasks $j$), then $\Ell_{\dro{}}$ is likely to get ``stuck'' trying to minimize $\ell_{\text{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 $\ell_i$ may have a positive effect on another task $\ell_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 $\Ell_{\dro{}}$ can lead to oscillatory behavior where the weights $\w$ alternate between two competing tasks.


\subsection{Lookahead DRO} \label{sec:lookahead}

The core limitation of $\Ell_{\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 $\btheta_0$, $\Ell_{\dro{}}$ chooses weights maximizing the model's loss $\w^*_0:=\argmax_{\w}\sum_iw_i\ell_i(\btheta_0)$. After training with weights $\w^*_0$, we obtain parameters $\btheta_1$ with loss  $\Ell_{\dro{}}(\btheta_1)=\max_{\w}\sum_iw_i\ell_i(\btheta_1)$. However, note that for $\btheta_1$, the worst-case weight allocation  $\w^*_1:=\argmax_{\w}\sum_iw_i\ell_i(\btheta_1)$ might be different than $\w^*_0$, so we have no guarantees that $\Ell_{\dro{}}(\btheta_1)\leq\Ell_{\dro{}}(\btheta_0)$. In other words, training with the worst-case weights $\w^*_0$ does not guarantee that $\Ell_{\dro{}}$ will decrease.

\input{figures/ldro_diagram_fig}

This observation suggests an alternative strategy: suppose we have a fixed---and, for simplicity, deterministic---learning algorithm $\w,\btheta_0\rightarrow \btheta_{\w,\btheta_0}$ for minimizing the $\w$-weighted loss $\sum_i w_i\ell_{i}$ starting from parameters $\btheta_0$. Rather than choosing $\w$ to be the worst-case weights corresponding to $\btheta_0$, we should pick weights that minimize the worst-case loss of $\btheta^*_{\w,\btheta_0}$, \ie{} \emph{after training}. This can be formalized as solving the min-max problem:
\begin{equation}\label{eq:lookahead_full_minmax}
    \min_{\w\in\Delta_N}\max_{\v\in\Delta_N}\sum_iv_i\ell_i(\btheta_{\w, \btheta_0})
\end{equation}
Intuitively, decoupling the weights $\w$ used for training from the weights $\v$ used for computing $\Ell_{\dro{}}(\btheta^*_{\w,\btheta_0})$ removes the restriction of just training on the worst-performing tasks and enables us to find a task weighting which is able to continuously minimize the worst-case loss. Note that in this case there might be multiple solution $\w$: for instance, it might be that one or more tasks have the same effect on the resulting worst-case task, and as such any permutation of their weights is an optimal solution. In such scenarios, we do not want to miss out on training on tasks on which performance can still be improved (as long as worst-case performance decreases). This can be formulated as a maximum entropy principle \citep{jaynes1957information}: \emph{if multiple optimal training weights exist, we should pick the one with the highest entropy}.

Because finding the optimal $\w$ necessitates ``looking ahead'' in the optimization trajectory from $\btheta_0$ to identify the DRO loss $\max_{\q}\sum_iv_i\ell_i(\btheta_{\w, \btheta_0})$, we call this weight selection principle \textbf{Lookahead-DRO} (L-DRO). We provide a graphical illustration in Figure \ref{fig:ldro_diagram}.



\subsection{One-step Lookahead DRO}

In practice we cannot explicitly solve the min-max game in Eq.~\ref{eq:lookahead_full_minmax} because estimating the payoff $\max_{\q}\sum_iv_i\ell_i(\btheta_{\w,\btheta_0})$ would require unrolling the full optimization process to compute $\btheta^*_{\w,\btheta_0}$. Instead, we simplify and take only one gradient step with a small learning rate $\lambda$, \ie{} replacing $\btheta^*_{\w,\btheta_0}$ with $\btheta_0 - \lambda\sum_iw_i\nabla\ell_i(\btheta_0)$. Such one-step approximations are often used in the meta-learning literature \citep{finn2017model,liu2018darts}. Using the first order Taylor extension of $\ell_i$ enables us to rewrite the individual lookahead task losses as:
\begin{equation*}
\ell_i(\btheta_0-\lambda[\sum_{j=1}^N w_j\nabla\ell_{j}(\btheta_0)])
\approx\ell_i(\btheta_0) - \lambda \sum_{j=1}^Nw_j\nabla\ell_{i}(\btheta_0)^\intercal\nabla\ell_{j}(\btheta_0)
\end{equation*}
The resulting payoffs are now linear in $\w$, which reduces the lookahead optimization problem in Eq.~\ref{eq:lookahead_full_minmax} into a bilinear matrix game:
\begin{equation}\label{eq:one_step_lookahead}
\begin{split}
    \min_{\w}\max_{\q}&\q^\intercal \mathbf{L} \w
\end{split}
\end{equation}
with a payoff matrix $\mathbf{L}$, called the ``lookahead matrix'', which can be decomposed as follows:
\begin{equation*}
    L_{ij}=\underbrace{\ell_i(\btheta_0)}_{\text{Loss of task $i$}} - \lambda\underbrace{{\nabla\ell_{i}(\btheta_0)}^\intercal\nabla\ell_j(\btheta_0)}_{\substack{\text{``interaction''} \\\text{between tasks $i$ and $j$}}}
\end{equation*}
We can find a solution (Nash equilibrium) of this game using many available solvers such as fictitious play \citep{brown1951iterative} or linear programming \citep{lemke1965bimatrix}. 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 $\lambda$ is treated as a hyperparameter. This results in a training algorithm which we outline in Algorithm~\ref{alg:ldro}.

\begin{algorithm}[H]
\SetAlgoLined
\DontPrintSemicolon
\KwIn{Lookahead learning rate $\lambda$, initial parameters $\btheta_0$}
 $\btheta\leftarrow \btheta_0$\;
 \While{training is not over}{
   $L\leftarrow \left[\ell_i(\btheta) - \lambda{\nabla\ell_{i}(\btheta)}^\intercal\nabla\ell_j(\btheta)\right]_{ij}$\tcp*{Compute lookahead matrix}
   $\v^*,\w^*\leftarrow \texttt{maxent\_nash\_solver}(L)$\tcp*{Find Nash equilibrium of $L$}
   $\btheta\leftarrow \texttt{optimizer\_step(}\sum_iw^*_i\ell_i\texttt{)}$\tcp*{Optimizer step on $\w^*$-weighted loss}
 }
 \caption{Lookahead DRO}
 \label{alg:ldro}
\end{algorithm}

To build more intuitions, we consider several edge cases and discuss how L-DRO handles them:

\textbf{Independent tasks.} $\:$ We have $\nabla \ell_i^\intercal\nabla\ell_j = \mathbbm{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^*=\argmax_{i}\ell_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}=\ell_i^*-\lambda<L_{ij}=\ell_i$ for all $j\neq 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.

\textbf{Redundant tasks.} $\:$ We have $\nabla \ell_i^\intercal\nabla\ell_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.

\textbf{The highest loss task has already converged.} $\:$ Consider the case where $\ell_0$ is maximal, but the model has already converged on task $0$ (in other words $\nabla\ell_0=0$). In this scenario---assuming there is no strong negative interaction among the remaining tasks---the highest loss will always be $\ell_0$ irrespective of the training weights $\w$. Again, by the maximum entropy principle, L-DRO advocates for picking uniform weights. This makes it possible to continue training on the remaining tasks even though $\ell_0$ has converged, bypassing the core limitations of \ac{dro} identified in \textsection \ref{sec:worst_case}.


\section{Lookahead DRO in Practice}

\subsection{Gradient and Loss Noises}

In most realistic machine learning scenarios we do not have access to the full losses $\ell_i(\btheta)$ and their gradients. Rather, we compute stochastic estimates on minibatches $\hat{\ell}_i(\btheta)$. In preliminary experiments, we 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-\ac{dro} algorithm proposed by \citet{sagawa2019distributionally}. 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 $\btheta$. Specifically, we interleave gradient updates on $\btheta$ with exponentiated gradient steps \citep{kivinen1997exponentiated} on $\w$ to minimize $\max_{\v}\v^\intercal L \w$. We show the weight update rule in Algorithm \ref{alg:weight_update}.

\begin{algorithm}[H]
\setstretch{0}
  \DontPrintSemicolon
  \SetKwFunction{Fupdate}{update\_weights}
  \SetKwProg{Fn}{Function}{:}{}
  \Fn{\Fupdate{Lookahead matrix $\vect{L}$, weights $\w$, weight update rate $\eta$}}{
 $\v^* \leftarrow \argmax_{\v} \v^\intercal \textbf{L}\w$ \tcp*{Compute best response to $\w$}
 $\tilde{w}_i\leftarrow w_ie^{-\eta \left[\v^{*\intercal}\textbf{L}\right]_i}$ \tcp*{Update weights}
 $w_i\leftarrow \frac{\tilde{w}_i}{\sum_j\tilde{w}_j}$ \tcp*{Re-normalize weights}
\KwRet $\w$
  }
 \caption{\label{alg:weight_update}Lookahead DRO online weight update}
\end{algorithm}

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 a mirror descent step \citep{nemirovskij1983problem} using the Kullback-Leibler divergence \citep{kullback1951information} as a distance function. Choosing weight updates that locally minimize the change in relative entropy plays into our maximum entropy guiding principle \citep{jumarie1990relative,kapur1992entropy}.

\subsection{Lookahead Matrix Computation}


The main computational bottleneck of L-DRO lies in the computation of the lookahead matrix $\mathbf{L}$. Recall that $L_{ij}$ can be decomposed in two terms: the loss $\ell_i$ and an ``interaction'' term ${\nabla\ell_{i}(\btheta)}^\intercal\nabla\ell_j(\btheta)$. 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'' $\mathbf{A}=\left[{\nabla\ell_{i}(\btheta)}^\intercal\nabla\ell_j(\btheta)\right]_{i,j=1\ldots N}$) periodically, typically every 100 or 1000 training steps.
In preliminary experiments, we also find 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}={\nabla\ell_{i}(\btheta)}^\intercal\left(\nabla\ell_j(\btheta)/\Vert\nabla\ell_j(\btheta)\Vert\right)$. We show the final algorithm for online L-DRO in Algorithm~\ref{alg:odro}.

\begin{algorithm}[tb]
\setstretch{0}

\DontPrintSemicolon
\KwIn{Lookahead learning rate $\lambda$, weight update rate $\eta$, interaction matrix update interval $k$}
 $A\leftarrow 0$\tcp*{Initialize interaction matrix at 0}
 $\w\leftarrow \left(\frac 1 N\cdots \frac 1 N\right)$\tcp*{Initialize $\w$ as uniform weights}
 \While{training is not over}{
  \If{i\% k == 0}{
   $A\leftarrow \left[\nabla\ell_i(\btheta)^\intercal\frac{\nabla\ell_j(\btheta)}{\Vert\nabla\ell_j(\btheta)\Vert}\right]_{1\leq i,j\leq n}$\tcp*{Update interaction matrix}
   }
%   /Compute task losses\;
   $L\leftarrow \ell(\btheta) - \lambda A$\tcp*{Compute lookahead matrix}
   $\w\leftarrow \texttt{update\_weights}(\w, L)$ \tcp*{Update weights using Algorithm \ref{alg:weight_update}}
   $\btheta\leftarrow \texttt{optimizer\_step(}\sum_iw_i\ell_i\texttt{)}$\tcp*{Optimizer step on $\w$-weighted loss}
 }
 \caption{Online lookahead DRO training loop}
 \label{alg:odro}
% \vspace{-1em}
\end{algorithm}

\section{Synthetic Experiments as a Proof of Concept}
\vspace{-0.5em}

In this section, we demonstrate the efficacy 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~\ref{fig:synthetic_experiments}). We identify each task $i$ with a point $y_i$ on the two dimensional plane $\mathbb R^2$ and define their respective loss functions as the squared Euclidean distance of the model prediction $f(\btheta)$ to $y_i$,  $\ell_i(\btheta)=\Vert y_i-f(\btheta)\Vert_2^2$.
We distribute the tasks such that most points are clustered around $(-1, 0)$, to simulate a grouping of similar tasks (\eg{} Indo-Europoean languages in our multilingual example). We then place a single point $y_{\text{outlier}}$ at $(1, 0)$ to emulate the presence of an outlier task.

We simulate the increased difficulty of the outlier task by adopting a parameterization $f(\btheta)$ that prevents the model from reaching any position closer than a fixed value $R$ from $y_{\text{outlier}}$. Specifically we adopt a polar parameterization $\btheta:=(r, \phi)$ and set
\begin{equation*}
    f(\btheta)=y_{\text{outlier}}+(R+r^2)\begin{bmatrix}\cos\phi\\\sin\phi\end{bmatrix}.
\end{equation*}
This essentially creates a ``forbidden region'' of radius $R$ around $y_{\text{outlier}}$ of the space that the model cannot penetrate, thus lower-bounding $\ell_{\text{outlier}}$ to $R$. In experiments, we set $R=1.25$.

Figure \ref{fig:dro_vs_avg} depicts the optimization trajectories when the model is trained either with the average objective or with \ac{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 \ac{dro} trajectory moves towards the latter but ends up stuck when it reaches a point where $\ell_{\text{outlier}}$ cannot be improved despite being the highest loss among all tasks.

On the other hand, while the L-DRO trajectory shown in Figure \ref{fig:lookahead} follows a similar trajectory as \ac{dro} at first, it is able to reach a solution that still achieves the best possible $\ell_{\text{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 \ref{fig:lookahead_noise} 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.


\section{Experiments}
\label{sec:experiments}

\subsection{CIFAR100}

As a first step towards a more realistic setting, we experiment with a multitask version of the CIFAR-100 dataset \citep{krizhevsky2009learning}, in which the 20 coarse labels are treated as separate 5-way classification tasks \citep{rosenbaum2017routing,Yu2020gradient_surgery}. In this setting, we train a small CNN model with four convolution layers with 32 $3\times 3$ filters each followed by a ReLU activation and a $2\times 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 (\ie{} across all 100 classes).


\input{figures/synthetic_figures}

In addition to L-DRO, we report results for three baselines. First \textbf{static-mixing}, where all tasks are weighted equally (\S{\ref{sec:static_mixing}}). Another natural baseline is \textbf{DRO} (\S{\ref{sec:worst_case}}) using the online algorithm proposed by \citet{sagawa2019distributionally}. Finally we compare to \textbf{PCGrad} \citep{Yu2020gradient_surgery}, 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.001, 0.01, 0.1\}$. 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.001, 0.01, 0.1\}$. For Static mixing and PCGrad, there are no additional parameters to sweep over.

\input{figures/cifar_100}
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 configurations) of each method in the two dimensional plane of worst and average accuracy.

We show the results in Figure \ref{fig:synthetic_experiments}. 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 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 \ref{sec:appendix_cifar_seeds}).

An important difference in our experimental setting with \citet{Yu2020gradient_surgery} 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 \ref{sec:appendix_cifar_bigcnn}). This finding suggests that L-DRO is primarily useful when the model's capacity is 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.



\vspace{-0.5em}
\subsection{Multilingual Language Modeling}

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 \citep{conneau-etal-2020-unsupervised,xue-etal-2021-mt5}. Improved pretraining performance has generally been found to correlate with better downstream performance \citep{Raffel2020t5}.

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\times$ 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 $p_{\btheta}$, its loss on the $i$-th language is the cross-entropy $\mathbb{H}(p_i,p_{\btheta})$ with the underlying distribution of the language, $p_i$. This can be re-written as $\mathbb{H}(p_i) + \KL(p_i||p_{\btheta})$. Since the KL divergence is non-negative, the loss of each language is lower-bounded by the language-specific entropy term $\mathbb{H}(p_i)$, irrespective of the capacity or the fitness of the model. In other words, languages with higher entropy will be more difficult to learn \citep{mielke-etal-2019-kind}.

\subsubsection{Experimental Setting}

We train a large autoregressive transformer model \citep{vaswani2017attention} on the multilingual C4 dataset \citep[mC4;][]{xue-etal-2021-mt5}. We limit ourselves to the 40 languages that are present in the XTREME benchmark \citep{Hu2020xtreme}, 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.\footnote{In experiments, we omit the Tagalog language, which is not included in mC4.} 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; \citealt{kudo-2018-subword}), 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 \citep{Kingma2014Adam} and a cosine annealed learning rate schedule \citep{vaswani2017attention}. 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:\footnote{We do not compare against PCGrad due to its underwhelming performance in the CIFAR-100 experiments and the prohibitive cost of its gradient projection step in this setting.}
\begin{itemize}[nolistsep,leftmargin=*]
\item \textbf{Static mixing} (\S\ref{sec:static_mixing}). We try 3 different values for the inverse temperature hyper-parameter $\alpha$: $0$ (uniform sampling), $0.3$ (a common value in the literature; \citealt{xue-etal-2021-mt5}) and $1$ (purely proportional sampling) 
\item \textbf{DRO} (\S\ref{sec:worst_case}), which minimizes the worst-case loss among all languages using the online algorithm from \citet{sagawa2019distributionally}. We sweep over the weight learning rate $\eta\in\{0.01, 0.001, 0.0001\}$.
\item \textbf{Baselined-DRO}, a variation on DRO proposed by \citet{oren2019distributionally} 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 difficult. In practice, we estimate $\mathbb{H}(p_i)$ by training a separate \emph{monolingual} model on each language, $p_{\btheta^*_i}$ and using its cross-entropy $\mathbb{H}(p_i,p_{\btheta^*_i})=\min_{\btheta}\mathbb{H}(p_i,p_{\btheta})$ as a proxy. We then train to minimize $\max_{i}\ell_i(\btheta)-\mathbb{H}(p_i,p_{\btheta^*_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.
\end{itemize}

\subsubsection{Results}

We report results in Table~\ref{tab:lm_results} in terms of perplexity  \citep{jelinek1977perplexity}, 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 \emph{relative} to a monolingual model. Specifically, the relative perplexity of a model $\btheta$ on language $i$ is defined as:
\begin{equation}
    \text{Relative perplexity}(\btheta, i)=\frac{\text{Perplexity}(\btheta, i) - \text{Perplexity}(\btheta^*_i,i)}{\text{Perplexity}(\btheta^*_i, i)}
\end{equation}
where $\btheta^*_i$ is a monolingual model trained on language $i$. In other words, relative perplexity of a given model on a language quantifies 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 ($\alpha=0.3$ or $1$) perform very well on high resource languages (\eg{} 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 $\mathbf{L}$), which makes it a much more appealing alternative.

\input{tables/lm_results}


\section{Related Work}

Much work in multi-task learning has focused on improving the optimization process when learning many tasks. Some prior work \citep{Sener2018,lin2019pareto,Ma2020} casts multi-task learning as multi-objective optimization, with the goal of finding a Pareto-optimal solution. Other work proposes to directly modify the gradients. In order to minimize gradient interference, \citet{Yu2020gradient_surgery} project conflicting gradients on the normal plane of the other. \citet{Wang2021gradient_vaccine} 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 \citep{kendall2018multi}, reward magnitude \citep{hessel2019multi}, and learning speed \citep{chen2018gradnorm,liu2019loss,zheng2019pyramidal}. 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 \citep{guo2018dynamic,jean2019adaptive} 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 \citep{ruder2019neural,crawshaw2020multi}. \citet{kiperwasser2018scheduled} proposed different predefined sampling schedules while \citet{sanh2019hierarchical} proposed a proportional sampling strategy. In multilingual language modeling and machine translation \citep{aharoni-etal-2019-massively,conneau-etal-2020-unsupervised,xue-etal-2021-mt5}, proportional sampling with a temperature is typically employed.

The distributional robustness literature focuses on training models that perform well on various domains \citep{bental2009robust,rahimian2019distributionally}. 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 \citep{hu2013kullback,gao2016distributionally,levy2020large}. Most relevant to our multitask setting are group-structured uncertainty sets, where the maximimum is taken over a mixture of sub-populations \citep{oren2019distributionally,sagawa2019distributionally,zhou2021examining}.



\section{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 a multitask model.

\section*{Acknowledgements}

The authors would like to thank the Language team at DeepMind for helpful discussion throughout the project, in particular Angeliki Lazaridou, Cyprien de Masson d'Autume, Adam Liska, Tayfun Terzi, Adhiguna Kuncoro, Kris Cao, and Phil Blunsom.

\section*{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. 

\section*{Reproducibility Statement}

We take several steps to ensure reproducibility of our work. Pseudo-code describing the final version of our algorithm is given in Algorithm \ref{alg:odro}. Moreover, we detail experimental settings and model hyper-parameters in \S\ref{sec:experiments}. Finally, all our experiments are performed on datasets that are openly available: CIFAR-100\footnote{\url{https://www.cs.toronto.edu/~kriz/cifar.html}} \citep{krizhevsky2009learning} and mC4\footnote{\url{https://www.tensorflow.org/datasets/catalog/c4\#c4multilingual}} \citep{xue-etal-2021-mt5}.


\bibliography{bibliography/myplain.bib,bibliography/references.bib}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\clearpage

\appendix

\section{Additional CIFAR-100 Results}

\subsection{Multiple Random Seeds} \label{sec:appendix_cifar_seeds}

Figure \ref{fig:cifar_100_all} depicts results on the CIFAR-100 dataset for 3 different random seeds. We find that L-DRO consistently outperforms baselines.

\input{figures/cifar_100_all_seeds}

\subsection{Experiments with larger CNN} \label{sec:appendix_cifar_bigcnn}

In Figure \ref{fig:cifar_100_big_all}, we report CIFAR-100 results across 3 random seeds using a larger CNN model: 3 layers of 160 $3\times 3$ filters each followed by a ReLU activation and a $2\times 2$ max-pooling layer, followed by a 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 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.


\input{figures/cifar_100_big_all_seeds}

\end{document}

%%%%% NEW MATH DEFINITIONS %%%%%

\usepackage{amsmath,amsfonts,bm}

% Mark sections of captions for referring to divisions of figures
\newcommand{\figleft}{{\em (Left)}}
\newcommand{\figcenter}{{\em (Center)}}
\newcommand{\figright}{{\em (Right)}}
\newcommand{\figtop}{{\em (Top)}}
\newcommand{\figbottom}{{\em (Bottom)}}
\newcommand{\captiona}{{\em (a)}}
\newcommand{\captionb}{{\em (b)}}
\newcommand{\captionc}{{\em (c)}}
\newcommand{\captiond}{{\em (d)}}

% Highlight a newly defined term
\newcommand{\newterm}[1]{{\bf #1}}


% Figure reference, lower-case.
\def\figref#1{figure~\ref{#1}}
% Figure reference, capital. For start of sentence
\def\Figref#1{Figure~\ref{#1}}
\def\twofigref#1#2{figures \ref{#1} and \ref{#2}}
\def\quadfigref#1#2#3#4{figures \ref{#1}, \ref{#2}, \ref{#3} and \ref{#4}}
% Section reference, lower-case.
\def\secref#1{section~\ref{#1}}
% Section reference, capital.
\def\Secref#1{Section~\ref{#1}}
% Reference to two sections.
\def\twosecrefs#1#2{sections \ref{#1} and \ref{#2}}
% Reference to three sections.
\def\secrefs#1#2#3{sections \ref{#1}, \ref{#2} and \ref{#3}}
% Reference to an equation, lower-case.
\def\eqref#1{equation~\ref{#1}}
% Reference to an equation, upper case
\def\Eqref#1{Equation~\ref{#1}}
% A raw reference to an equation---avoid using if possible
\def\plaineqref#1{\ref{#1}}
% Reference to a chapter, lower-case.
\def\chapref#1{chapter~\ref{#1}}
% Reference to an equation, upper case.
\def\Chapref#1{Chapter~\ref{#1}}
% Reference to a range of chapters
\def\rangechapref#1#2{chapters\ref{#1}--\ref{#2}}
% Reference to an algorithm, lower-case.
\def\algref#1{algorithm~\ref{#1}}
% Reference to an algorithm, upper case.
\def\Algref#1{Algorithm~\ref{#1}}
\def\twoalgref#1#2{algorithms \ref{#1} and \ref{#2}}
\def\Twoalgref#1#2{Algorithms \ref{#1} and \ref{#2}}
% Reference to a part, lower case
\def\partref#1{part~\ref{#1}}
% Reference to a part, upper case
\def\Partref#1{Part~\ref{#1}}
\def\twopartref#1#2{parts \ref{#1} and \ref{#2}}

\def\ceil#1{\lceil #1 \rceil}
\def\floor#1{\lfloor #1 \rfloor}
\def\1{\bm{1}}
\newcommand{\train}{\mathcal{D}}
\newcommand{\valid}{\mathcal{D_{\mathrm{valid}}}}
\newcommand{\test}{\mathcal{D_{\mathrm{test}}}}

\def\eps{{\epsilon}}


% Random variables
\def\reta{{\textnormal{$\eta$}}}
\def\ra{{\textnormal{a}}}
\def\rb{{\textnormal{b}}}
\def\rc{{\textnormal{c}}}
\def\rd{{\textnormal{d}}}
\def\re{{\textnormal{e}}}
\def\rf{{\textnormal{f}}}
\def\rg{{\textnormal{g}}}
\def\rh{{\textnormal{h}}}
\def\ri{{\textnormal{i}}}
\def\rj{{\textnormal{j}}}
\def\rk{{\textnormal{k}}}
\def\rl{{\textnormal{l}}}
% rm is already a command, just don't name any random variables m
\def\rn{{\textnormal{n}}}
\def\ro{{\textnormal{o}}}
\def\rp{{\textnormal{p}}}
\def\rq{{\textnormal{q}}}
\def\rr{{\textnormal{r}}}
\def\rs{{\textnormal{s}}}
\def\rt{{\textnormal{t}}}
\def\ru{{\textnormal{u}}}
\def\rv{{\textnormal{v}}}
\def\rw{{\textnormal{w}}}
\def\rx{{\textnormal{x}}}
\def\ry{{\textnormal{y}}}
\def\rz{{\textnormal{z}}}

% Random vectors
\def\rvepsilon{{\mathbf{\epsilon}}}
\def\rvtheta{{\mathbf{\theta}}}
\def\rva{{\mathbf{a}}}
\def\rvb{{\mathbf{b}}}
\def\rvc{{\mathbf{c}}}
\def\rvd{{\mathbf{d}}}
\def\rve{{\mathbf{e}}}
\def\rvf{{\mathbf{f}}}
\def\rvg{{\mathbf{g}}}
\def\rvh{{\mathbf{h}}}
\def\rvu{{\mathbf{i}}}
\def\rvj{{\mathbf{j}}}
\def\rvk{{\mathbf{k}}}
\def\rvl{{\mathbf{l}}}
\def\rvm{{\mathbf{m}}}
\def\rvn{{\mathbf{n}}}
\def\rvo{{\mathbf{o}}}
\def\rvp{{\mathbf{p}}}
\def\rvq{{\mathbf{q}}}
\def\rvr{{\mathbf{r}}}
\def\rvs{{\mathbf{s}}}
\def\rvt{{\mathbf{t}}}
\def\rvu{{\mathbf{u}}}
\def\rvv{{\mathbf{v}}}
\def\rvw{{\mathbf{w}}}
\def\rvx{{\mathbf{x}}}
\def\rvy{{\mathbf{y}}}
\def\rvz{{\mathbf{z}}}

% Elements of random vectors
\def\erva{{\textnormal{a}}}
\def\ervb{{\textnormal{b}}}
\def\ervc{{\textnormal{c}}}
\def\ervd{{\textnormal{d}}}
\def\erve{{\textnormal{e}}}
\def\ervf{{\textnormal{f}}}
\def\ervg{{\textnormal{g}}}
\def\ervh{{\textnormal{h}}}
\def\ervi{{\textnormal{i}}}
\def\ervj{{\textnormal{j}}}
\def\ervk{{\textnormal{k}}}
\def\ervl{{\textnormal{l}}}
\def\ervm{{\textnormal{m}}}
\def\ervn{{\textnormal{n}}}
\def\ervo{{\textnormal{o}}}
\def\ervp{{\textnormal{p}}}
\def\ervq{{\textnormal{q}}}
\def\ervr{{\textnormal{r}}}
\def\ervs{{\textnormal{s}}}
\def\ervt{{\textnormal{t}}}
\def\ervu{{\textnormal{u}}}
\def\ervv{{\textnormal{v}}}
\def\ervw{{\textnormal{w}}}
\def\ervx{{\textnormal{x}}}
\def\ervy{{\textnormal{y}}}
\def\ervz{{\textnormal{z}}}

% Random matrices
\def\rmA{{\mathbf{A}}}
\def\rmB{{\mathbf{B}}}
\def\rmC{{\mathbf{C}}}
\def\rmD{{\mathbf{D}}}
\def\rmE{{\mathbf{E}}}
\def\rmF{{\mathbf{F}}}
\def\rmG{{\mathbf{G}}}
\def\rmH{{\mathbf{H}}}
\def\rmI{{\mathbf{I}}}
\def\rmJ{{\mathbf{J}}}
\def\rmK{{\mathbf{K}}}
\def\rmL{{\mathbf{L}}}
\def\rmM{{\mathbf{M}}}
\def\rmN{{\mathbf{N}}}
\def\rmO{{\mathbf{O}}}
\def\rmP{{\mathbf{P}}}
\def\rmQ{{\mathbf{Q}}}
\def\rmR{{\mathbf{R}}}
\def\rmS{{\mathbf{S}}}
\def\rmT{{\mathbf{T}}}
\def\rmU{{\mathbf{U}}}
\def\rmV{{\mathbf{V}}}
\def\rmW{{\mathbf{W}}}
\def\rmX{{\mathbf{X}}}
\def\rmY{{\mathbf{Y}}}
\def\rmZ{{\mathbf{Z}}}

% Elements of random matrices
\def\ermA{{\textnormal{A}}}
\def\ermB{{\textnormal{B}}}
\def\ermC{{\textnormal{C}}}
\def\ermD{{\textnormal{D}}}
\def\ermE{{\textnormal{E}}}
\def\ermF{{\textnormal{F}}}
\def\ermG{{\textnormal{G}}}
\def\ermH{{\textnormal{H}}}
\def\ermI{{\textnormal{I}}}
\def\ermJ{{\textnormal{J}}}
\def\ermK{{\textnormal{K}}}
\def\ermL{{\textnormal{L}}}
\def\ermM{{\textnormal{M}}}
\def\ermN{{\textnormal{N}}}
\def\ermO{{\textnormal{O}}}
\def\ermP{{\textnormal{P}}}
\def\ermQ{{\textnormal{Q}}}
\def\ermR{{\textnormal{R}}}
\def\ermS{{\textnormal{S}}}
\def\ermT{{\textnormal{T}}}
\def\ermU{{\textnormal{U}}}
\def\ermV{{\textnormal{V}}}
\def\ermW{{\textnormal{W}}}
\def\ermX{{\textnormal{X}}}
\def\ermY{{\textnormal{Y}}}
\def\ermZ{{\textnormal{Z}}}

% Vectors
\def\vzero{{\bm{0}}}
\def\vone{{\bm{1}}}
\def\vmu{{\bm{\mu}}}
\def\vtheta{{\bm{\theta}}}
\def\va{{\bm{a}}}
\def\vb{{\bm{b}}}
\def\vc{{\bm{c}}}
\def\vd{{\bm{d}}}
\def\ve{{\bm{e}}}
\def\vf{{\bm{f}}}
\def\vg{{\bm{g}}}
\def\vh{{\bm{h}}}
\def\vi{{\bm{i}}}
\def\vj{{\bm{j}}}
\def\vk{{\bm{k}}}
\def\vl{{\bm{l}}}
\def\vm{{\bm{m}}}
\def\vn{{\bm{n}}}
\def\vo{{\bm{o}}}
\def\vp{{\bm{p}}}
\def\vq{{\bm{q}}}
\def\vr{{\bm{r}}}
\def\vs{{\bm{s}}}
\def\vt{{\bm{t}}}
\def\vu{{\bm{u}}}
\def\vv{{\bm{v}}}
\def\vw{{\bm{w}}}
\def\vx{{\bm{x}}}
\def\vy{{\bm{y}}}
\def\vz{{\bm{z}}}

% Elements of vectors
\def\evalpha{{\alpha}}
\def\evbeta{{\beta}}
\def\evepsilon{{\epsilon}}
\def\evlambda{{\lambda}}
\def\evomega{{\omega}}
\def\evmu{{\mu}}
\def\evpsi{{\psi}}
\def\evsigma{{\sigma}}
\def\evtheta{{\theta}}
\def\eva{{a}}
\def\evb{{b}}
\def\evc{{c}}
\def\evd{{d}}
\def\eve{{e}}
\def\evf{{f}}
\def\evg{{g}}
\def\evh{{h}}
\def\evi{{i}}
\def\evj{{j}}
\def\evk{{k}}
\def\evl{{l}}
\def\evm{{m}}
\def\evn{{n}}
\def\evo{{o}}
\def\evp{{p}}
\def\evq{{q}}
\def\evr{{r}}
\def\evs{{s}}
\def\evt{{t}}
\def\evu{{u}}
\def\evv{{v}}
\def\evw{{w}}
\def\evx{{x}}
\def\evy{{y}}
\def\evz{{z}}

% Matrix
\def\mA{{\bm{A}}}
\def\mB{{\bm{B}}}
\def\mC{{\bm{C}}}
\def\mD{{\bm{D}}}
\def\mE{{\bm{E}}}
\def\mF{{\bm{F}}}
\def\mG{{\bm{G}}}
\def\mH{{\bm{H}}}
\def\mI{{\bm{I}}}
\def\mJ{{\bm{J}}}
\def\mK{{\bm{K}}}
\def\mL{{\bm{L}}}
\def\mM{{\bm{M}}}
\def\mN{{\bm{N}}}
\def\mO{{\bm{O}}}
\def\mP{{\bm{P}}}
\def\mQ{{\bm{Q}}}
\def\mR{{\bm{R}}}
\def\mS{{\bm{S}}}
\def\mT{{\bm{T}}}
\def\mU{{\bm{U}}}
\def\mV{{\bm{V}}}
\def\mW{{\bm{W}}}
\def\mX{{\bm{X}}}
\def\mY{{\bm{Y}}}
\def\mZ{{\bm{Z}}}
\def\mBeta{{\bm{\beta}}}
\def\mPhi{{\bm{\Phi}}}
\def\mLambda{{\bm{\Lambda}}}
\def\mSigma{{\bm{\Sigma}}}

% Tensor
\DeclareMathAlphabet{\mathsfit}{\encodingdefault}{\sfdefault}{m}{sl}
\SetMathAlphabet{\mathsfit}{bold}{\encodingdefault}{\sfdefault}{bx}{n}
\newcommand{\tens}[1]{\bm{\mathsfit{#1}}}
\def\tA{{\tens{A}}}
\def\tB{{\tens{B}}}
\def\tC{{\tens{C}}}
\def\tD{{\tens{D}}}
\def\tE{{\tens{E}}}
\def\tF{{\tens{F}}}
\def\tG{{\tens{G}}}
\def\tH{{\tens{H}}}
\def\tI{{\tens{I}}}
\def\tJ{{\tens{J}}}
\def\tK{{\tens{K}}}
\def\tL{{\tens{L}}}
\def\tM{{\tens{M}}}
\def\tN{{\tens{N}}}
\def\tO{{\tens{O}}}
\def\tP{{\tens{P}}}
\def\tQ{{\tens{Q}}}
\def\tR{{\tens{R}}}
\def\tS{{\tens{S}}}
\def\tT{{\tens{T}}}
\def\tU{{\tens{U}}}
\def\tV{{\tens{V}}}
\def\tW{{\tens{W}}}
\def\tX{{\tens{X}}}
\def\tY{{\tens{Y}}}
\def\tZ{{\tens{Z}}}


% Graph
\def\gA{{\mathcal{A}}}
\def\gB{{\mathcal{B}}}
\def\gC{{\mathcal{C}}}
\def\gD{{\mathcal{D}}}
\def\gE{{\mathcal{E}}}
\def\gF{{\mathcal{F}}}
\def\gG{{\mathcal{G}}}
\def\gH{{\mathcal{H}}}
\def\gI{{\mathcal{I}}}
\def\gJ{{\mathcal{J}}}
\def\gK{{\mathcal{K}}}
\def\gL{{\mathcal{L}}}
\def\gM{{\mathcal{M}}}
\def\gN{{\mathcal{N}}}
\def\gO{{\mathcal{O}}}
\def\gP{{\mathcal{P}}}
\def\gQ{{\mathcal{Q}}}
\def\gR{{\mathcal{R}}}
\def\gS{{\mathcal{S}}}
\def\gT{{\mathcal{T}}}
\def\gU{{\mathcal{U}}}
\def\gV{{\mathcal{V}}}
\def\gW{{\mathcal{W}}}
\def\gX{{\mathcal{X}}}
\def\gY{{\mathcal{Y}}}
\def\gZ{{\mathcal{Z}}}

% Sets
\def\sA{{\mathbb{A}}}
\def\sB{{\mathbb{B}}}
\def\sC{{\mathbb{C}}}
\def\sD{{\mathbb{D}}}
% Don't use a set called E, because this would be the same as our symbol
% for expectation.
\def\sF{{\mathbb{F}}}
\def\sG{{\mathbb{G}}}
\def\sH{{\mathbb{H}}}
\def\sI{{\mathbb{I}}}
\def\sJ{{\mathbb{J}}}
\def\sK{{\mathbb{K}}}
\def\sL{{\mathbb{L}}}
\def\sM{{\mathbb{M}}}
\def\sN{{\mathbb{N}}}
\def\sO{{\mathbb{O}}}
\def\sP{{\mathbb{P}}}
\def\sQ{{\mathbb{Q}}}
\def\sR{{\mathbb{R}}}
\def\sS{{\mathbb{S}}}
\def\sT{{\mathbb{T}}}
\def\sU{{\mathbb{U}}}
\def\sV{{\mathbb{V}}}
\def\sW{{\mathbb{W}}}
\def\sX{{\mathbb{X}}}
\def\sY{{\mathbb{Y}}}
\def\sZ{{\mathbb{Z}}}

% Entries of a matrix
\def\emLambda{{\Lambda}}
\def\emA{{A}}
\def\emB{{B}}
\def\emC{{C}}
\def\emD{{D}}
\def\emE{{E}}
\def\emF{{F}}
\def\emG{{G}}
\def\emH{{H}}
\def\emI{{I}}
\def\emJ{{J}}
\def\emK{{K}}
\def\emL{{L}}
\def\emM{{M}}
\def\emN{{N}}
\def\emO{{O}}
\def\emP{{P}}
\def\emQ{{Q}}
\def\emR{{R}}
\def\emS{{S}}
\def\emT{{T}}
\def\emU{{U}}
\def\emV{{V}}
\def\emW{{W}}
\def\emX{{X}}
\def\emY{{Y}}
\def\emZ{{Z}}
\def\emSigma{{\Sigma}}

% entries of a tensor
% Same font as tensor, without \bm wrapper
\newcommand{\etens}[1]{\mathsfit{#1}}
\def\etLambda{{\etens{\Lambda}}}
\def\etA{{\etens{A}}}
\def\etB{{\etens{B}}}
\def\etC{{\etens{C}}}
\def\etD{{\etens{D}}}
\def\etE{{\etens{E}}}
\def\etF{{\etens{F}}}
\def\etG{{\etens{G}}}
\def\etH{{\etens{H}}}
\def\etI{{\etens{I}}}
\def\etJ{{\etens{J}}}
\def\etK{{\etens{K}}}
\def\etL{{\etens{L}}}
\def\etM{{\etens{M}}}
\def\etN{{\etens{N}}}
\def\etO{{\etens{O}}}
\def\etP{{\etens{P}}}
\def\etQ{{\etens{Q}}}
\def\etR{{\etens{R}}}
\def\etS{{\etens{S}}}
\def\etT{{\etens{T}}}
\def\etU{{\etens{U}}}
\def\etV{{\etens{V}}}
\def\etW{{\etens{W}}}
\def\etX{{\etens{X}}}
\def\etY{{\etens{Y}}}
\def\etZ{{\etens{Z}}}

% The true underlying data generating distribution
\newcommand{\pdata}{p_{\rm{data}}}
% The empirical distribution defined by the training set
\newcommand{\ptrain}{\hat{p}_{\rm{data}}}
\newcommand{\Ptrain}{\hat{P}_{\rm{data}}}
% The model distribution
\newcommand{\pmodel}{p_{\rm{model}}}
\newcommand{\Pmodel}{P_{\rm{model}}}
\newcommand{\ptildemodel}{\tilde{p}_{\rm{model}}}
% Stochastic autoencoder distributions
\newcommand{\pencode}{p_{\rm{encoder}}}
\newcommand{\pdecode}{p_{\rm{decoder}}}
\newcommand{\precons}{p_{\rm{reconstruct}}}

\newcommand{\laplace}{\mathrm{Laplace}} % Laplace distribution

\newcommand{\E}{\mathbb{E}}
\newcommand{\Ls}{\mathcal{L}}
\newcommand{\R}{\mathbb{R}}
\newcommand{\emp}{\tilde{p}}
\newcommand{\lr}{\alpha}
\newcommand{\reg}{\lambda}
\newcommand{\rect}{\mathrm{rectifier}}
\newcommand{\softmax}{\mathrm{softmax}}
\newcommand{\sigmoid}{\sigma}
\newcommand{\softplus}{\zeta}
\newcommand{\KL}{D_{\mathrm{KL}}}
\newcommand{\Var}{\mathrm{Var}}
\newcommand{\standarderror}{\mathrm{SE}}
\newcommand{\Cov}{\mathrm{Cov}}
% Wolfram Mathworld says $L^2$ is for function spaces and $\ell^2$ is for vectors
% But then they seem to use $L^2$ for vectors throughout the site, and so does
% wikipedia.
\newcommand{\normlzero}{L^0}
\newcommand{\normlone}{L^1}
\newcommand{\normltwo}{L^2}
\newcommand{\normlp}{L^p}
\newcommand{\normmax}{L^\infty}

\newcommand{\parents}{Pa} % See usage in notation.tex. Chosen to match Daphne's book.

\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}

\DeclareMathOperator{\sign}{sign}
\DeclareMathOperator{\Tr}{Tr}
\let\ab\allowbreak

\begin{figure}
    \centering
    \includegraphics[width=0.85\columnwidth]{figures/ldro_diagram.pdf}
    \caption{Overview of the Lookahead-DRO procedure. Starting from parameters $\btheta_0$, the optimal training weights $\w$ are chosen such that the resulting parameters $\btheta^*_{\w}$ minimize the worst case loss.}
    \label{fig:ldro_diagram}
    \vspace{-1em}
\end{figure}

\begin{figure*}[t!]
\centering
\begin{subfigure}[t]{0.3\textwidth}
\centering
\includegraphics[width=\columnwidth]{figures/dro_vs_avg.pdf}
\caption{\label{fig:dro_vs_avg} \acrshort{dro} vs. Average loss
}
\end{subfigure}
~
\begin{subfigure}[t]{0.3\textwidth}
\centering
\includegraphics[width=\columnwidth]{figures/lookahead.pdf}
\caption{\label{fig:lookahead} L-DRO
}
\end{subfigure}
~
\begin{subfigure}[t]{0.3\textwidth}
\centering
\includegraphics[width=\columnwidth]{figures/lookahead_noise.pdf}
\caption{\label{fig:lookahead_noise} L-DRO (with noise)
}
\end{subfigure}
\caption{\label{fig:synthetic_experiments} Synthetic multitask experiment. The outlier task is shown in \textcolor{red}{red} while the others are in \textcolor{yellow}{yellow}. The \textcolor{gray}{grayed} out area represents the region in space that cannot be attained by the model with our chosen parameterization}
\vspace{-2em}
\end{figure*}

\begin{wrapfigure}{R}{0.4\textwidth}
    \centering
    \includegraphics[width=0.4\columnwidth]{figures/cifar100small_main.png}
    \caption{Average/worst task accuracy Pareto frontiers on multitask CIFAR-100}
    \label{fig:cifar_100}
    \vspace{-1em}
\end{wrapfigure}
\begin{table}[]
    \centering
    \begin{tabular}{l|cc|cc}
    \toprule
         & \multicolumn{2}{c}{absolute perplexity} & \multicolumn{2}{c}{relative perplexity}  \\
         \cmidrule{2-5}
         & average & worst & average & worst \\
         \midrule
        Static Mixing ($\alpha=1$)& 48.7 & 135.5 & 358.3 & 1170.7 \\
        Static Mixing ($\alpha=0.3$)& 24.0 & 41.2 & 130.6 & 213.7 \\
        Static Mixing ($\alpha=0$)& 23.9 & 48.7 & 128.2 & 198.1 \\
        DRO & 24.5 & {\bf 35.1} & 142.7 & 286.7 \\
        Baselined-DRO & 23.9 & 50.7 & 127.3 &  {\bf 168.5} \\
        L-DRO & {\bf 23.6} & 55.5 & {\bf 126.0} & 191.7 \\
        \midrule
        Monolingual & 10.6 & 25.3 & 0 & 0 \\
        \bottomrule
    \end{tabular}
    \caption{Results on the multilingual language modeling experiment. We report results both in terms of absolute and relative perplexity.}
    \label{tab:lm_results}
    \vspace{-2em}
\end{table}
\begin{figure}
    \centering
    \includegraphics[width=\columnwidth]{figures/cifar100small_all.png}
    \caption{Average/worst task accuracy Pareto frontiers across 3 random seeds on multitask CIFAR-100}
    \label{fig:cifar_100_all}
\end{figure}
\begin{figure}
    \centering
    \includegraphics[width=\columnwidth]{figures/cifar100big.png}
    \caption{Average/worst task accuracy Pareto frontiers across 3 random seeds on multitask CIFAR-100}
    \label{fig:cifar_100_big_all}
\end{figure}
\begin{thebibliography}{50}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
  \providecommand{\doi}[1]{doi: #1}\else
  \providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi

\bibitem[Aharoni et~al.(2019)Aharoni, Johnson, and
  Firat]{aharoni-etal-2019-massively}
Roee Aharoni, Melvin Johnson, and Orhan Firat.
\newblock Massively multilingual neural machine translation.
\newblock In \emph{Proceedings of the 2019 Conference of the North American
  Chapter of the Association for Computational Linguistics: Human Language
  Technologies (NAACL-HLT)}, pp.\  3874--3884, Minneapolis, Minnesota, June
  2019. Association for Computational Linguistics.
\newblock \doi{10.18653/v1/N19-1388}.
\newblock URL \url{https://aclanthology.org/N19-1388}.

\bibitem[Ben-Tal et~al.(2009)Ben-Tal, El~Ghaoui, and
  Nemirovski]{bental2009robust}
Aharon Ben-Tal, Laurent El~Ghaoui, and Arkadi Nemirovski.
\newblock \emph{Robust optimization}, volume~28.
\newblock Princeton University Press, 2009.

\bibitem[Brown(1951)]{brown1951iterative}
George~W Brown.
\newblock Iterative solution of games by fictitious play.
\newblock \emph{Activity analysis of production and allocation}, 13\penalty0
  (1):\penalty0 374--376, 1951.

\bibitem[Buolamwini \& Gebru(2018)Buolamwini and Gebru]{buolamwini2018gender}
Joy Buolamwini and Timnit Gebru.
\newblock Gender shades: Intersectional accuracy disparities in commercial
  gender classification.
\newblock In \emph{Conference on fairness, accountability and transparency},
  pp.\  77--91. PMLR, 2018.

\bibitem[Chen et~al.(2018)Chen, Badrinarayanan, Lee, and
  Rabinovich]{chen2018gradnorm}
Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, and Andrew Rabinovich.
\newblock Gradnorm: Gradient normalization for adaptive loss balancing in deep
  multitask networks.
\newblock In \emph{International Conference on Machine Learning}, pp.\
  794--803. PMLR, 2018.

\bibitem[Conneau et~al.(2020)Conneau, Khandelwal, Goyal, Chaudhary, Wenzek,
  Guzm{\'a}n, Grave, Ott, Zettlemoyer, and
  Stoyanov]{conneau-etal-2020-unsupervised}
Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume
  Wenzek, Francisco Guzm{\'a}n, Edouard Grave, Myle Ott, Luke Zettlemoyer, and
  Veselin Stoyanov.
\newblock Unsupervised cross-lingual representation learning at scale.
\newblock In \emph{Proceedings of the 8th Annual Meeting of the Association for
  Computational Linguistics (ACL)}, pp.\  8440--8451, Online, July 2020.
  Association for Computational Linguistics.
\newblock \doi{10.18653/v1/2020.acl-main.747}.
\newblock URL \url{https://aclanthology.org/2020.acl-main.747}.

\bibitem[Crawshaw(2020)]{crawshaw2020multi}
Michael Crawshaw.
\newblock Multi-task learning with deep neural networks: A survey.
\newblock \emph{arXiv preprint arXiv:2009.09796}, 2020.

\bibitem[Finn et~al.(2017)Finn, Abbeel, and Levine]{finn2017model}
Chelsea Finn, Pieter Abbeel, and Sergey Levine.
\newblock Model-agnostic meta-learning for fast adaptation of deep networks.
\newblock In \emph{Proceedings of the 34th International Conference on Machine
  Learning (ICML)}, pp.\  1126--1135. PMLR, 2017.

\bibitem[Gao \& Kleywegt(2016)Gao and Kleywegt]{gao2016distributionally}
Rui Gao and Anton~J Kleywegt.
\newblock Distributionally robust stochastic optimization with wasserstein
  distance.
\newblock \emph{arXiv preprint arXiv:1604.02199}, 2016.

\bibitem[Guo et~al.(2018)Guo, Haque, Huang, Yeung, and Fei-Fei]{guo2018dynamic}
Michelle Guo, Albert Haque, De-An Huang, Serena Yeung, and Li~Fei-Fei.
\newblock Dynamic task prioritization for multitask learning.
\newblock In \emph{Proceedings of the European Conference on Computer Vision
  (ECCV)}, pp.\  270--287, 2018.

\bibitem[Hessel et~al.(2019)Hessel, Soyer, Espeholt, Czarnecki, Schmitt, and
  van Hasselt]{hessel2019multi}
Matteo Hessel, Hubert Soyer, Lasse Espeholt, Wojciech Czarnecki, Simon Schmitt,
  and Hado van Hasselt.
\newblock Multi-task deep reinforcement learning with popart.
\newblock In \emph{Proceedings of the 33rd Meeting of the Association for
  Advancement of Artificial Intelligence (AAAI)}, pp.\  3796--3803, 2019.

\bibitem[Hu et~al.(2020)Hu, Ruder, Siddhant, Neubig, Firat, and
  Johnson]{Hu2020xtreme}
Junjie Hu, Sebastian Ruder, Aditya Siddhant, Graham Neubig, Orhan Firat, and
  Melvin Johnson.
\newblock {XTREME: A Massively Multilingual Multi-task Benchmark for Evaluating
  Cross-lingual Generalization}.
\newblock In \emph{Proceedings of the 37th International Conference on Machine
  Learning (ICML)}, 2020.

\bibitem[Hu \& Hong(2013)Hu and Hong]{hu2013kullback}
Zhaolin Hu and L~Jeff Hong.
\newblock Kullback-leibler divergence constrained distributionally robust
  optimization.
\newblock \emph{Available at Optimization Online}, 2013.

\bibitem[Jaynes(1957)]{jaynes1957information}
Edwin~T Jaynes.
\newblock Information theory and statistical mechanics.
\newblock \emph{Physical review}, 106\penalty0 (4):\penalty0 620, 1957.

\bibitem[Jean et~al.(2019)Jean, Firat, and Johnson]{jean2019adaptive}
S{\'e}bastien Jean, Orhan Firat, and Melvin Johnson.
\newblock Adaptive scheduling for multi-task learning.
\newblock \emph{arXiv preprint arXiv:1909.06434}, 2019.

\bibitem[Jelinek et~al.(1977)Jelinek, Mercer, Bahl, and
  Baker]{jelinek1977perplexity}
Fred Jelinek, Robert~L Mercer, Lalit~R Bahl, and James~K Baker.
\newblock Perplexity—a measure of the difficulty of speech recognition tasks.
\newblock \emph{The Journal of the Acoustical Society of America}, 62\penalty0
  (S1):\penalty0 S63--S63, 1977.

\bibitem[Jumarie(1990)]{jumarie1990relative}
Guy Jumarie.
\newblock Relative information—what for?
\newblock In \emph{Relative Information}, pp.\  1--11. Springer, 1990.

\bibitem[Kapur \& Kesavan(1992)Kapur and Kesavan]{kapur1992entropy}
Jagat~Narain Kapur and Hiremaglur~K Kesavan.
\newblock Entropy optimization principles and their applications.
\newblock In \emph{Entropy and energy dissipation in water resources}, pp.\
  3--20. Springer, 1992.

\bibitem[Kendall et~al.(2018)Kendall, Gal, and Cipolla]{kendall2018multi}
Alex Kendall, Yarin Gal, and Roberto Cipolla.
\newblock Multi-task learning using uncertainty to weigh losses for scene
  geometry and semantics.
\newblock In \emph{Proceedings of the 31st IEEE Conference on Computer Vision
  and Pattern Recognition (CVPR)}, pp.\  7482--7491, 2018.

\bibitem[Kingma \& Ba(2014)Kingma and Ba]{Kingma2014Adam}
Diederik~P. Kingma and Jimmy Ba.
\newblock Adam: A method for stochastic optimization.
\newblock In \emph{Proceedings of the International Conference on Learning
  Representations (ICLR)}, 2014.

\bibitem[Kiperwasser \& Ballesteros(2018)Kiperwasser and
  Ballesteros]{kiperwasser2018scheduled}
Eliyahu Kiperwasser and Miguel Ballesteros.
\newblock Scheduled multi-task learning: From syntax to translation.
\newblock \emph{Transactions of the Association for Computational Linguistics},
  6:\penalty0 225--240, 2018.

\bibitem[Kivinen \& Warmuth(1997)Kivinen and Warmuth]{kivinen1997exponentiated}
Jyrki Kivinen and Manfred~K Warmuth.
\newblock Exponentiated gradient versus gradient descent for linear predictors.
\newblock \emph{information and computation}, 132\penalty0 (1):\penalty0 1--63,
  1997.

\bibitem[Krizhevsky et~al.(2009)]{krizhevsky2009learning}
Alex Krizhevsky et~al.
\newblock Learning multiple layers of features from tiny images.
\newblock Technical report, Citeseer, 2009.

\bibitem[Kudo(2018)]{kudo-2018-subword}
Taku Kudo.
\newblock Subword regularization: Improving neural network translation models
  with multiple subword candidates.
\newblock In \emph{Proceedings of the 56th Annual Meeting of the Association
  for Computational Linguistics (Volume 1: Long Papers)}, pp.\  66--75,
  Melbourne, Australia, July 2018. Association for Computational Linguistics.
\newblock \doi{10.18653/v1/P18-1007}.
\newblock URL \url{https://aclanthology.org/P18-1007}.

\bibitem[Kullback \& Leibler(1951)Kullback and
  Leibler]{kullback1951information}
Solomon Kullback and Richard~A Leibler.
\newblock On information and sufficiency.
\newblock \emph{The annals of mathematical statistics}, 22\penalty0
  (1):\penalty0 79--86, 1951.

\bibitem[Lemke(1965)]{lemke1965bimatrix}
Carlton~E Lemke.
\newblock Bimatrix equilibrium points and mathematical programming.
\newblock \emph{Management science}, 11\penalty0 (7):\penalty0 681--689, 1965.

\bibitem[Levy et~al.(2020)Levy, Carmon, Duchi, and Sidford]{levy2020large}
Daniel Levy, Yair Carmon, John~C Duchi, and Aaron Sidford.
\newblock Large-scale methods for distributionally robust optimization.
\newblock \emph{Proceedings of the 34th Annual Conference on Neural Information
  Processing Systems (NeurIPS)}, 33, 2020.

\bibitem[Lin et~al.(2019)Lin, Zhen, Li, Zhang, and Kwong]{lin2019pareto}
Xi~Lin, Hui-Ling Zhen, Zhenhua Li, Qing-Fu Zhang, and Sam Kwong.
\newblock Pareto multi-task learning.
\newblock \emph{Advances in neural information processing systems},
  32:\penalty0 12060--12070, 2019.

\bibitem[Liu et~al.(2018)Liu, Simonyan, and Yang]{liu2018darts}
Hanxiao Liu, Karen Simonyan, and Yiming Yang.
\newblock Darts: Differentiable architecture search.
\newblock In \emph{Proceedings of the International Conference on Learning
  Representations (ICLR)}, 2018.

\bibitem[Liu et~al.(2019)Liu, Liang, and Gitter]{liu2019loss}
Shengchao Liu, Yingyu Liang, and Anthony Gitter.
\newblock Loss-balanced task weighting to reduce negative transfer in
  multi-task learning.
\newblock In \emph{Proceedings of the 33rd Meeting of the Association for
  Advancement of Artificial Intelligence (AAAI)}, pp.\  9977--9978, 2019.

\bibitem[Ma et~al.(2020)Ma, Du, and Matusik]{Ma2020}
Pingchuan Ma, Tao Du, and Wojciech Matusik.
\newblock {Effcient Continuous Pareto Exploration in Multi-Task Learning}.
\newblock In \emph{Proceedings of ICML 2020}, 2020.

\bibitem[McCann et~al.(2018)McCann, Keskar, Xiong, and
  Socher]{mccann2018natural}
Bryan McCann, Nitish~Shirish Keskar, Caiming Xiong, and Richard Socher.
\newblock The natural language decathlon: Multitask learning as question
  answering.
\newblock \emph{arXiv preprint arXiv:1806.08730}, 2018.

\bibitem[Mielke et~al.(2019)Mielke, Cotterell, Gorman, Roark, and
  Eisner]{mielke-etal-2019-kind}
Sabrina~J. Mielke, Ryan Cotterell, Kyle Gorman, Brian Roark, and Jason Eisner.
\newblock What kind of language is hard to language-model?
\newblock In \emph{Proceedings of the 57th Annual Meeting of the Association
  for Computational Linguistics}, pp.\  4975--4989, Florence, Italy, July 2019.
  Association for Computational Linguistics.
\newblock \doi{10.18653/v1/P19-1491}.
\newblock URL \url{https://aclanthology.org/P19-1491}.

\bibitem[Nemirovskij \& Yudin(1983)Nemirovskij and
  Yudin]{nemirovskij1983problem}
Arkadij~Semenovi{\v{c}} Nemirovskij and David~Borisovich Yudin.
\newblock Problem complexity and method efficiency in optimization.
\newblock \emph{Wiley-Interscience Series in Discrete Mathematics}, 1983.

\bibitem[Oren et~al.(2019)Oren, Sagawa, Hashimoto, and
  Liang]{oren2019distributionally}
Yonatan Oren, Shiori Sagawa, Tatsunori Hashimoto, and Percy Liang.
\newblock Distributionally robust language modeling.
\newblock In \emph{Proceedings of the 2019 Conference on Empirical Methods in
  Natural Language Processing (EMNLP)}, pp.\  4227--4237, 2019.
\newblock URL \url{https://www.aclweb.org/anthology/D19-1432}.

\bibitem[Raffel et~al.(2020)Raffel, Shazeer, Roberts, Lee, Narang, Matena,
  Zhou, Li, and Liu]{Raffel2020t5}
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael
  Matena, Yanqi Zhou, Wei Li, and Peter~J. Liu.
\newblock {Exploring the Limits of Transfer Learning with a Unified
  Text-to-Text Transformer}.
\newblock \emph{Journal of Machine Learning Research}, 21, 2020.
\newblock URL \url{http://arxiv.org/abs/1910.10683}.

\bibitem[Rahimian \& Mehrotra(2019)Rahimian and
  Mehrotra]{rahimian2019distributionally}
Hamed Rahimian and Sanjay Mehrotra.
\newblock Distributionally {R}obust {O}ptimization: A {R}eview.
\newblock \emph{arXiv preprint arXiv:1908.05659}, 2019.

\bibitem[Rosenbaum et~al.(2018)Rosenbaum, Klinger, and
  Riemer]{rosenbaum2017routing}
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer.
\newblock Routing networks: Adaptive selection of non-linear functions for
  multi-task learning.
\newblock In \emph{Proceedings of the International Conference on Learning
  Representations (ICLR)}, 2018.

\bibitem[Ruder(2017)]{Ruder2017overview}
Sebastian Ruder.
\newblock {An Overview of Multi-Task Learning in Deep Neural Networks}.
\newblock In \emph{arXiv preprint arXiv:1706.05098}, 2017.

\bibitem[Ruder(2019)]{ruder2019neural}
Sebastian Ruder.
\newblock \emph{Neural transfer learning for natural language processing}.
\newblock PhD thesis, NUI Galway, 2019.

\bibitem[Sagawa et~al.(2020)Sagawa, Koh, Hashimoto, and
  Liang]{sagawa2019distributionally}
Shiori Sagawa, Pang~Wei Koh, Tatsunori~B Hashimoto, and Percy Liang.
\newblock Distributionally robust neural networks for group shifts: On the
  importance of regularization for worst-case generalization.
\newblock In \emph{Proceedings of the International Conference on Learning
  Representations (ICLR)}, 2020.
\newblock URL \url{https://arxiv.org/pdf/1911.08731.pdf}.

\bibitem[Sanh et~al.(2019)Sanh, Wolf, and Ruder]{sanh2019hierarchical}
Victor Sanh, Thomas Wolf, and Sebastian Ruder.
\newblock A hierarchical multi-task approach for learning embeddings from
  semantic tasks.
\newblock In \emph{Proceedings of the 33rd Meeting of the Association for
  Advancement of Artificial Intelligence (AAAI)}, pp.\  6949--6956, 2019.

\bibitem[Sener \& Koltun(2018)Sener and Koltun]{Sener2018}
Ozan Sener and Vladlen Koltun.
\newblock {Multi-Task Learning as Multi-Objective Optimization}.
\newblock In \emph{Proceedings of NIPS 2018}, 2018.

\bibitem[Vaswani et~al.(2017)Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez,
  Kaiser, and Polosukhin]{vaswani2017attention}
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones,
  Aidan~N Gomez, {\L}ukasz Kaiser, and Illia Polosukhin.
\newblock Attention is all you need.
\newblock In \emph{Proceedings of the 31st Annual Conference on Neural
  Information Processing Systems (NIPS)}, pp.\  5998--6008, 2017.

\bibitem[Wang et~al.(2019)Wang, Singh, Michael, Hill, Levy, and
  Bowman]{glueleaderboard2020}
Alex Wang, Amapreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel~R
  Bowman.
\newblock Glue benchmark leaderboard.
\newblock https://gluebenchmark.com/leaderboard, 2019.
\newblock URL \url{https://gluebenchmark.com/leaderboard}.
\newblock Accessed: 2021-09-24.

\bibitem[Wang et~al.(2021)Wang, Tsvetkov, Firat, and
  Cao]{Wang2021gradient_vaccine}
Zirui Wang, Yulia Tsvetkov, Orhan Firat, and Yuan Cao.
\newblock {Gradient Vaccine: Investigating and Improving Multi-task
  Optimization in Massively Multilingual Models}.
\newblock In \emph{Proceedings of ICLR 2021}, 2021.

\bibitem[Xue et~al.(2021)Xue, Constant, Roberts, Kale, Al-Rfou, Siddhant,
  Barua, and Raffel]{xue-etal-2021-mt5}
Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami Al-Rfou, Aditya
  Siddhant, Aditya Barua, and Colin Raffel.
\newblock m{T}5: A massively multilingual pre-trained text-to-text transformer.
\newblock In \emph{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.
\newblock \doi{10.18653/v1/2021.naacl-main.41}.
\newblock URL \url{https://aclanthology.org/2021.naacl-main.41}.

\bibitem[Yu et~al.(2020)Yu, Kumar, Gupta, Levine, Hausman, and
  Finn]{Yu2020gradient_surgery}
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and
  Chelsea Finn.
\newblock {Gradient surgery for multi-task learning}.
\newblock In \emph{Proceedings of NeurIPS 2020}, 2020.

\bibitem[Zheng et~al.(2019)Zheng, Deng, Sun, Jiang, Guo, Yu, Huang, and
  Ji]{zheng2019pyramidal}
Feng Zheng, Cheng Deng, Xing Sun, Xinyang Jiang, Xiaowei Guo, Zongqiao Yu,
  Feiyue Huang, and Rongrong Ji.
\newblock Pyramidal person re-identification via multi-loss dynamic training.
\newblock In \emph{Proceedings of the IEEE/CVF Conference on Computer Vision
  and Pattern Recognition}, pp.\  8514--8522, 2019.

\bibitem[Zhou et~al.(2021)Zhou, Ma, Michel, and Neubig]{zhou2021examining}
Chunting Zhou, Xuezhe Ma, Paul Michel, and Graham Neubig.
\newblock Examining and combating spurious features under distribution shift.
\newblock In \emph{Proceedings of the 38th International Conference on Machine
  Learning (ICML)}, 2021.

\end{thebibliography}