Minimax Regret Optimization for Robust Machine Learning under Distribution Shift

Learning good models for scenarios where the evaluation happens under a test distribution that differs from the training distribution has become an increasingly important problem in machine learning and statistics. It is particularly motivated by the widespread practical deployment of machine learning models, and the observation that model performance deteriorates when the test distribution differs from the training distribution. Some motivating scenarios range from class imbalance between training and test data (Galar et al., 2011; Japkowicz, 2000) to algorithmic fairness (Dwork et al., 2012; Barocas and Selbst, 2016). Consequently, several recent formalisms have been proposed to develop algorithms for such settings, such as adversarial training (Goodfellow et al., 2014; Madry et al., 2017; Wang et al., 2019), Invariant Risk Minimization (Arjovsky et al., 2019; Ahuja et al., 2020a; Rosenfeld et al., 2020) and Distributionally Robust Optimization (Duchi and Namkoong, 2021; Namkoong and Duchi, 2016; Staib and Jegelka, 2019; Kuhn et al., 2019). While the specifics differ across these formulations, at a high-level they all minimize the worst-case risk of the learned model across some family of test distributions, and spend considerable effort on finding natural classes of test distributions and designing efficient algorithms to optimize the worst-case objectives.

In this paper, we consider this problem from a learning theory perspective, by asking if it is possible to obtain regret bounds that hold uniformly across all test distributions. Recall that in supervised learning, we are given a function class $\mathcal{F}$ and a loss function $\ell(z,f(z))$, where $z$ denotes a sample. For a test distribution $P$, we are interested in learning a prediction function $f\in\mathcal{F}$ using the training data, so that the risk of $f$ under $P$: $R_{P}(f)={\mathbb{E}}_{z\sim P}\ell(z,f(z))$, is small. In classical learning theory, an estimator $\hat{f}$ is learned from the training data, and we are interested in bounding the regret (or excess risk) of $f$ as follows

where $f_{P}$ minimizes $R_{P}(f)$ over $\mathcal{F}$. If the training data $(z_{1},\ldots,z_{n})$ are drawn from the same distribution as the test distribution $P$, then standard results imply that the empirical risk minimizaiton (ERM) method, which finds an $\hat{f}$ with a small training error, achieves small regret. A natural question we would like to ask in this paper is whether there exists an estimator that allows us to achieve small regret with distribution shift? More precisely, suppose we are given a family of distributions $\mathcal{P}$ with $P_{0}\in\mathcal{P}$. Given access to $n$ samples $(z_{1},\ldots,z_{n})$ from $P_{0}$, we would like to find an estimator $\hat{f}$ from the training data so that the uniform regret

is small. Here the model class $\mathcal{F}$ can be statistically misspecified in that it may not contain the optimal model $f^{\star}$ that attains the Bayes risk in a pointwise manner.

We note that this objective is different from the objective of Distributionally Robust Optimization (henceforth DRO), which directly minimizes the risk across all test distributions as follows:

We observe that the objective (1) is less sensitive than DRO to heterogeneity in the amount of noise in test distributions. Moreover, when the model is well-specified, our criterion directly measures the closeness of the estimator $\hat{f}$ and the optimal model $f^{\star}$, which is often desirable in applications. We call the criterion to minimize regret uniformly across test distributaions Minimax Regret Optimization (MRO), and its population formulation seeks to minimize the worst-case regret (1):

Compared to DRO, MRO evaluates the regret of a candidate model $f$ on each distribution $P\in\mathcal{P}$ as opposed to the raw risk. As we show in the following sections, regret is comparable across distribution in a more robust manner than the risk, since the subtraction of the minimum risk for each distribution takes away the variance due to noise. We note that a similar approach to remove the residual variance to deal with heterogenous noise arises in the offline Reinforcement Learning setting (Antos et al., 2008), which we comment more on in the discussion of related work. Similar to ERM, which replaces the regret minimization problem over a single test distribution by its empirical minimization counterpart, we consider an empirical minimization counterpart of MRO, and analyze its generalization performance.

We consider a class of learning problems where the learner is given a dataset sampled from a distribution $P_{0}$ of the form $z_{1},\ldots,z_{n}$ with $z_{i}\in\mathcal{Z}$. We also have a class $\mathcal{P}$ of distributions such that we want to do well relative to this entire class of distributions, even though we only have access to samples from $P_{0}$. In Minimax Regret Optimization, we formalize the property of doing uniformly well for all $P\in\mathcal{P}$ through the objective (3). As mentioned in the introduction, this is closely related to the DRO objective (2). The following result shows that small risk does not lead to small regret, which means DRO does not solve the MRO objective (3). The construction also shows that under heterogenous noise, DRO tends to focus on a distribution $P\in\mathcal{P}$ with large noise level, which can be undesirable for many applications.

It is easy to verify that $R_{1}(f_{1})=0.04$, $R_{1}(f_{2})=0.25$, $R_{2}(f_{1})=0.29$, and $R_{2}(f_{2})=0.26$. Then the DRO objective (2) evaluates to the optimal value of 0.26, with $f_{2}$ attaining this value. On the other hand, the MRO objective (3) is equal to $0.03$ and this value is attained for $f_{1}$, while $f_{2}$ incurs a much higher objective of $\sup_{P}R_{P}(f_{2})-R_{P}(f_{P})=0.21$. As a result, we have $f_{\text{DRO}}=f_{2}$ and $f_{\text{MRO}}=f_{1}$. Comparing regrets lets MRO remove the underlying noise level of $0.25$ for $P_{2}$.  

Note that while the worst-case regret of MRO is always better than that of DRO, this does not imply that the function $f_{\text{MRO}}$ is always preferable under each distribution $P\in\mathcal{P}$ to $f_{\text{DRO}}$. Next we give one such example where $f_{\text{DRO}}$ achieves an arguably better trade-off in balanced performance across distributions. This situation can happen in real applications when functions achieving minimum risks vary greatly across distributions, potentially due to overfitting. The example shows that both MRO and DRO may have advantages or disadvantages under different conditions.

More generally, one can always find a distribution $P\in\mathcal{P}$ under which the regret of $f_{\text{DRO}}$ is smaller than that of $f_{\text{MRO}}$ and vice versa, but the ordering in the worst-case is clear by definition. In the context of these observations, we next discuss how we might estimate $f_{\text{MRO}}$ from samples, for which it is useful to rewrite the objective (3) in an equivalent weight-based formulation we discuss next.

We begin with the most general case with fairly standard assumptions on the loss function, and the function and weight classes and show the following result on the regret of $\widehat{f}_{\text{MRO}}$ for any $w\in\mathcal{W}$.

The bound of Theorem 2 highlights the main difference of our objective compared with DRO style approaches. The bound states that we our solution $\widehat{f}_{\text{MRO}}$ has a small regret for each $w\in\mathcal{W}$, as long as at least one such function exists in the function class, up to the usual finite sample deviation terms. Thus, $\widehat{f}_{\text{MRO}}$ attains the uniform regret guarantee we asked for. To better interpret our result, we state an easy corollary, before making some additional remarks.

A limitation of Theorem 2 is that it does not leverage any structure of the loss function beyond Assumption 1, and as a result obtains a $1/\sqrt{n}$ dependence on the number of samples. For the case of a fixed distribution, it is well known that self-bounding properties of the loss variance can be used to obtain faster $1/n$ bounds on the regret of the ERM solution, and here we show that this improvement extends to our setting. For ease of exposition, we specialize to the squared loss setting in this section. That is, our samples $z$ take the form $(x,y)$ with $x\in\mathcal{X}\subseteq\mathbb{R}^{d}$, $y\in[-1,1]$ , our functions $f$ map from $\mathcal{X}$ to $[-1,1]$ and and $\ell(z,f(z))=(f(x)-y)^{2}$. We make an additional convexity assumption on the function class $\mathcal{F}$.

Note that convexity of $\mathcal{F}$ is quite different from convexity of $f$ in its parameters, and can always be satisfied by taking the convex hull of a base class. The assumption can be avoided by replacing our ERM based solution with aggregation approaches (see e.g. Tsybakov, 2003; Dalalyan and Salmon, 2012). However, we focus on the convex case here to illustrate the key ideas.

Theorem 4 crucially leverages the following property of convex classes $\mathcal{F}$ and squared loss.

Note that a similar result holds for the squared loss and any class $\mathcal{F}$, if we replace $f_{P}$ with the unconstrained minimizer of $R_{P}(f)$ over all measurable functions, but using this property in our analysis would result in an additional approximation error term in the bound of Theorem 4, which is undesirable when considering regret within the function class.

Our result so far guarantees uniform regret, with learning complexity that depends on the worst case complexity over all distributions $w$. For example, if the complexity of the distribution family is characterized by $B_{w}$ as in our analysis, then the bound depends on $\sup_{w}B_{w}$. In this section, we show that it is possible to improve such dependence using an adaptive scaling of the objective function. To better illustrate the consequences of this approach, let us assume that we can rewrite our guarantees in the form that for each $w\in\mathcal{W}$ and $f\in\mathcal{F}$, we have with probability at least $1-\delta$,

where the scaling function $c_{w}$ is a distribution $P$ dependent quantity that is known to the algorithm. For instance, under conditions of Theorem 2, we get this bound with $c=1$, $c_{w}=\sigma_{w}+B_{w}/\sqrt{n}$ and $\epsilon=d_{\mathcal{F},\mathcal{W}}(\delta)/\sqrt{n}$ (assuming $d_{\mathcal{F},\mathcal{W}}(\delta)\geq 1$). Under conditions of Theorem 4, we can use $c=3$, $c_{w}=B_{w}$ and $\epsilon=d_{\mathcal{F},\mathcal{W}}/n$, by taking the second $O(B/n)$ bound of Theorem 4. While this does worsen our dependence on $d_{\mathcal{F},\mathcal{W}}(\delta)$ in the first case, this is preferable to assuming that $d_{\mathcal{F},\mathcal{W}}(\delta)$ is known and setting of $c_{w}=\sigma_{w}+B_{w}\sqrt{d_{\mathcal{F},\mathcal{W}}(\delta)/n}$ and $\epsilon=\sqrt{d_{\mathcal{F},\mathcal{W}}(\delta)/n}$.

Since we assume $c_{w}$ is known, then we can define the estimator (Scaled MRO, or SMRO):

We now present our main results for the statistical properties of $\widehat{f}_{\text{SMRO}}$ under the assumptions of Theorems 2 and 4, before comparing the results and discussing some key consequences.

We can also state a version for squared loss and convex classes.

Comparing Theorems 7 and 8 with their counterparts of MRO, we notice a subtle but important difference. The finite sample deviation term $\epsilon^{\prime}_{w}$ in Theorem 7 depends on the weights $w$ for which the bound is stated, while the corresponding term in Theorem 7 is $\sup_{w^{\prime}\in\mathcal{W}}\epsilon_{w^{\prime}}$. In other words, the heterogeneous scaling in SMRO allows us to obtain a regret bound for each distribution depending on the deviation properties of that particular distribution, unlike in the basic MRO approach. To see the benefits of this approach, we state a corollary of Theorems 7 and 8 next.

Both results follow by choosing $f=f_{w}=f^{\star}$ in Theorems 7 and 8. This result shows that the rescaling allows our bounds to adapt to the closeness of a target distribution to the data collection distribution, simultaneously for all target distributions. The alignment condition of a shared $f^{\star}$ is always true in well-specified problems, and we illustrate a consequence of Corollary 9 for well-specified linear regression in Appendix D.

So far our development has focused on the statistical properties of MRO. In this section, we discuss how the MRO estimator can be computed from a finite dataset, given some reasonable computational assumptions on the function class $\mathcal{F}$.

While we assume access to an exact ERM oracle, we can weaken the notion to an approximate oracle with an optimization error comparable to the statistical error from finite samples. Given such an oracle, we can now approximately solve the MRO objective (6) (or (10)) by using a well-known strategy of solving minimax problems as two player zero-sum games. To do so, we recall our assumption of a finite class $\mathcal{W}$, and let $\Delta(\mathcal{W})$ denote the set of all distributions over the importance weight family. Similarly $\Delta(\mathcal{F})$ is the set of distributions over our function class. Then we can rewrite the objective in Equation (6) as:

The objective is bilinear in $P$ and $\rho$ so that the following result holds due to the exchangeability of minimum and maximum for convex-concave saddle point problems. It is also a direct consequence of Proposition 12.

For finite $\mathcal{W}$ which we consider in this paper, it is also possible to find the approximate $\hat{\rho}$ and $\widehat{f}_{\text{MRO}}$ efficiently, using the celebrated result of Freund and Schapire (1996) to solve this minimax problem through no-regret dynamics. We use the best response strategy for the $P$-player, as finding the best response distribution given a specific $\rho$ is equivalent to finding the function $f\in\mathcal{F}$ that minimizes ${\mathbb{E}}_{w\sim\rho}\widehat{R}_{w}(f)$, which can be accomplished through one call to the ERM oracle. For optimizing $\rho$, we use the exponentiated gradient algorithm of Kivinen and Warmuth (1997) (closely related to Hedge (Freund and Schapire, 1997)), which is a common no-regret strategy to learn a distribution over a finite collection of experts, with each $w\in\mathcal{W}$ being an ”expert” in our setting. More formally, we initialize $\rho_{1}$ to be the uniform distribution over $\mathcal{W}$ and repeatedly update:

Let us denote the distribution $P_{t}=(f_{1}+\ldots+f_{t})/t$ for the iterates generated by (12). Then the results of Freund and Schapire (1996) yield the following suboptimality bound on $P_{t}$.

At a high-level, Proposition 12 allows us to choose $T=n^{2}$ to ensure that the optimization error is no larger than the statistical error, allowing the same bounds to apply up to constants. The optimization strategy used here bears some resemblance to boosting techniques, which also seek to reweight the data in order to ensure a uniformly good performance on all samples, though the reweighting is not constrained to a particular class $\mathcal{W}$ unlike here. Note that while the optimization error scales logarithmically in $|\mathcal{W}|$, the computational complexity is linear in the size of this set, meaning that our strategy is computationally feasible for sets $\mathcal{W}$ of a modest size. On the other hand, our earlier discussion (8) suggests that alternative reformulations of the objective might be computationally preferred in the case of class $\mathcal{W}$ which are continuous. Handling the intermediate regime of a large discrete class $\mathcal{W}$ in a computationally efficient manner is an interesting question for future research.

In this paper, we introduce the MRO criterion for problems with distribution shift, and establish learning theoretic properties of optimizing the criterion from samples. We demonstrate the many benefits of reasoning with uniform regret as opposed to uniform risk guarantees, and we expect these observations to have implications beyond the setting of distribution shift.

On a technical side, it remains interesting to further develop scalable algorithms for large datasets and weight classes. Better understanding the statistical scaling with the size of the weight class and refining our techniques for important scenarios such as covariate shift are also important directions for future research.