.tocmtchapter mtchaptersubsection mtappendixnone

#  UMix: Improving Importance Weighting for Subpopulation Shift via
Uncertainty-Aware Mixup

Zongbo Han1 333 , Zhipeng Liang2111 444 , Fan Yang3111, Liu Liu3, Lanqing Li3,
Yatao Bian3,
Peilin Zhao3, Bingzhe Wu3222, Changqing Zhang1222, Jianhua Yao3222
1College of Intelligence and Computing, Tianjin University,
2 Hong Kong University of Science and Technology, 3Tencent AI Lab Equal
contribution. ‡ Supported by 2021 Tencent Rhino-Bird Research Elite Training
Program. § Work done during an internship at Tencent AI Lab. $\dagger$
Corresponding authors<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Subpopulation shift widely exists in many real-world machine learning
applications, referring to the training and test distributions containing the
same subpopulation groups but varying in subpopulation frequencies. Importance
reweighting is a normal way to handle the subpopulation shift issue by
imposing constant or adaptive sampling weights on each sample in the training
dataset. However, some recent studies have recognized that most of these
approaches fail to improve the performance over empirical risk minimization
especially when applied to over-parameterized neural networks. In this work,
we propose a simple yet practical framework, called uncertainty-aware mixup
(UMix), to mitigate the overfitting issue in over-parameterized models by
reweighting the “mixed” samples according to the sample uncertainty. The
training-trajectories-based uncertainty estimation is equipped in the proposed
UMix for each sample to flexibly characterize the subpopulation distribution.
We also provide insightful theoretical analysis to verify that UMix achieves
better generalization bounds over prior works. Further, we conduct extensive
empirical studies across a wide range of tasks to validate the effectiveness
of our method both qualitatively and quantitatively. Code is available at this
URL.

##  1 Introduction

Empirical risk minimization (ERM) typically faces challenges from distribution
shift, which refers to the difference between training and test distributions
[61, 27, 3]. One common type of distribution shift is subpopulation shift
wherein the training and test distributions consist of the same subpopulation
groups but differ in subpopulation frequencies [6, 8]. Many practical research
problems (e.g., fairness of machine learning and class imbalance) can all be
considered as a special case of subpopulation shift [32, 21, 28]. For example,
in the setting of fair machine learning, we train the model on a training
dataset with biased demographic subpopulations and test it on an unbiased test
dataset [32, 21]. Therefore the essential goal of fair machine learning is to
mitigate the subpopulation shift between training and test datasets.

Many approaches have been proposed for solving this problem. Among these
approaches, importance weighting (IW) is a classical yet effective technique
by imposing static or adaptive weights on each sample when building weighted
empirical loss. Therefore each subpopulation group contributes comparably to
the final training objective. Specifically, there are normally two ways to
achieve importance reweighting. Early works propose to reweight the sample
inverse proportionally to the subpopulation frequencies (i.e., static weights)
[61, 59, 13, 58, 11, 42], such as class-imbalanced learning approaches [13,
11, 42]. Alternatively, a more flexible way is to reweight individual samples
adaptively according to training dynamics [66, 74, 47, 72, 35, 48, 40, 62].
Distributional robust optimization (DRO) is one of the most representative
methods in this line, which minimizes the loss over the worst-case
distribution in a neighborhood of the empirical training distribution. A
commonly used dual form of DRO can be seen as a special case of importance
reweighting wherein the sampling weights are updated based on the current loss
[52, 24, 38, 25] in an alternated manner.

However, some recent studies have shown both empirically and theoretically
that these IW methods could fail to achieve better worst-case subpopulation
performance compared with ERM. Empirically, prior works [10, 58] recognize
that various IW methods tend to exacerbate overfitting, which leads to a
diminishing effect on stochastic gradient descent (SGD) over training epochs
especially when they are applied to over-parameterized neural networks (NNs).
Theoretically, previous studies prove that for over-parameterized neural
networks, reweighting algorithms do not improve over ERM because their
implicit biases are (almost) equivalent [73, 59, 68]. In addition, some prior
works also point out that using conventional regularization techniques such as
weight decay cannot significantly improve the performance of IW [58].

To this end, we introduce a novel technique called uncertainty-aware mixup
(UMix), by reweighting the mixed samples according to uncertainty within the
mini-batch while mitigating overfitting. Specifically, we employ the well-
known mixup technique to produce “mixed” augmented samples. Then we train the
model on these mixed samples to make sure it can always see “novel” samples
thus the effects of IW will not dissipate even at the end of the training
epoch. To enforce the model to perform fairly well on all subpopulations, we
further efficiently reweight the mixed samples according to uncertainty of the
original samples. The weighted mixup loss function is induced by combining the
weighted losses of the corresponding two original samples. At a high level,
this approach augments training samples in an uncertainty-aware manner, i.e.,
putting more focus on samples with higher prediction uncertainties that belong
to minority subpopulations with high probabilities. We also show UMix can
provide additional theoretical benefit which achieves a tighter generalization
bound than weighted ERM [41, 40, 72, 38]. The contributions of this paper are:

  * •

We propose a simple and practical approach called uncertainty-aware mixup
(UMix) to improve previous IW methods by reweighting the mixed samples, which
provides a new framework to mitigate overfitting in over-parameterized neural
networks.

  * •

Under the proposed framework, we provide theoretical analysis with insight
that UMix can achieve a tighter generalization bound than the weighted ERM.

  * •

We perform extensive experiments on a wide range of tasks, where the proposed
UMix achieves excellent performance in both group-oblivious and group-aware
settings.

Comparison with existing works. Here, we discuss the key differences between
UMix and other works. In contrast to most IW methods (e.g., CVaR-DRO [38] and
JTT [41]), UMix employs a mixup strategy to improve previous IW methods and
mitigate the model overfitting. Among these methods, JTT [41] and LISA [70]
are the two most related works to ours. Specifically, JTT provides a two-stage
optimization framework in which an additional network is used for building the
error set, and then JTT upweights samples in the error set in the following
training stage. Besides, LISA also modifies mixup for improving model
robustness against distribution shift. However, LISA intuitively mixes the
samples within the same subpopulation or same label thus it needs additional
subpopulation information. In contrast to them, UMix introduces sample weights
into the vanilla mixup strategy by quantitatively measuring the sample
uncertainties without subpopulation information. In addition, our work is
orthogonal to LISA, i.e., we can use our weight building strategy to improve
LISA’s performance. In practice, our method consistently outperforms previous
approaches that do not use subpopulation information and even achieves quite
competitive performance to those methods which leverage subpopulation
information. We also provide theoretical analysis to explain why UMix works
better than the weighted ERM [41, 40, 72, 38].

##  2 Related Work

###  2.1 Importance weighting

To improve the model robustness against subpopulation shift, importance
weighting (IW) is a classical yet effective technique by imposing static or
adaptive weight on each sample and then building weighted empirical loss.
Therefore each subpopulation group can have a comparable strength in the final
training objective. Specifically, there are typically two ways to achieve
importance reweighting, i.e., using static or adaptive importance weights.

Static methods. The naive reweighting approaches perform static reweighting
based on the distribution of training samples [61, 59, 13, 58, 11, 42]. Their
core motivation is to make different subpopulations have a comparable
contribution to the training objective by reweighting. Specifically, the most
intuitive way is to set the weight of each sample to be inversely proportional
to the number of samples in each subpopulation [61, 59, 58]. Besides, there
are some methods to obtain sample weights based on the effective number of
samples [13], subpopulation margins [11], and Bayesian networks [42].

Adaptive methods. In contrast to the above static methods, a more essential
way is to assign each individual sample an adaptive weight that can vary
according to training dynamics [66, 74, 47, 72, 35, 48, 40, 62].
Distributional robust optimization (DRO) is one of the most representative
methods in this line, which minimizes the loss over the worst-case
distribution in a neighborhood of the empirical training distribution. A
commonly-used dual form of DRO can be considered as a special case of
importance reweighting wherein the sampling weights are updated based on the
current loss [52, 24, 38, 25] in an alternated manner. For example, in the
group-aware setting (i.e., we know each sample belongs to which
subpopulation), GroupDRO [58] introduces an online optimization algorithm to
update the weights of each group. In the group-oblivious setting, [66, 35, 47,
48] model the problem as a (regularized) minimax game, where one player aims
to minimize the loss by optimizing the model parameters and another player
aims to maximize the loss by assigning weights to each sample.

###  2.2 Uncertainty quantification

The core of our method is based on the high-quality uncertainty quantification
of each sample. There are many approaches proposed for this goal. The
uncertainty of deep learning models includes epistemic (model) uncertainty and
aleatoric (data) uncertainty [30]. To obtain the epistemic uncertainty,
Bayesian neural networks (BNNs) [53, 45, 15, 30] have been proposed which
replace the deterministic weight parameters of model with distribution. Unlike
BNNs, ensemble-based methods obtain the epistemic uncertainty by training
multiple models and ensembling them [36, 22, 2, 26]. Aleatoric uncertainty
focuses on the inherent noise in the data, which usually is learned as a
function of the data [30, 37, 54]. Uncertainty quantification has been
successfully equipped in many fields such as multimodal learning [44, 20, 19],
multitask learning [31, 14], and reinforcement learning [29, 39]. Unlike
previous methods, our method focuses on estimating the epistemic uncertainty
of training samples with subpopulation shift and upweighting uncertain
samples, thereby improving the performance of minority subpopulations with
high uncertainty.

##  3 Method

In this section, we introduce technical details of UMix. The key idea of UMix
is to exploit uncertainty information to upweight mixed samples, and thus can
encourage the model to perform uniformly well on all subpopulations. We first
introduce the basic procedure of UMix and then present how to provide high-
quality uncertainty estimations which is the fundamental block of UMix.

###  3.1 Background

The necessary background and notations are provided here. Let the input and
label space be $\mathcal{X}$ and $\mathcal{Y}$ respectively. Given training
dataset $\mathcal{D}$ with $N$ training samples
$\\{(x_{i},y_{i})\\}_{i=1}^{N}$ i.i.d. sampled from a probability distribution
$P$. We consider the setting that the training distribution $P$ is a mixture
of $G$ predefined subpopulations, i.e., $P=\sum_{g=1}^{G}k_{g}P_{g}$, where
$k_{g}$ and $P_{g}$ denote the $g$-th subpopulation’s proportion and
distribution respectively. Our goal is to obtain a model
$f_{\theta}:\mathcal{X}\rightarrow\mathcal{Y}$ parameterized by
$\theta\in\Theta$ that performs well on all subpopulations.

The well-known empirical risk minimization (ERM) algorithm doesn’t consider
the subpopulations and minimizes the expected risk
$\mathbb{E}{[\ell(\theta,x_{i},y_{i})]}$, where $\ell$ denotes the loss
function. This leads to the model paying more attention to the majority
subpopulations in the training set and resulting in poor performance on the
minority subpopulations. For example, the ERM-based models may learn spurious
correlations that exist in majority subpopulations but not in minority
subpopulations [58]. The proposed method aims to learn a model that is robust
against subpopulation shift by importance weighting.

Previous works on improving subpopulation shift robustness investigate several
different settings, i.e., group-aware and group-oblivious [72, 41, 58]. Most
of the previous works have assumed that the group label is available during
training [58, 70]. This is called the group-aware setting. However, due to
some reasons, we may not have training group labels. For example, in many real
applications, it’s hard to extract group label information. Meanwhile, the
group label information may not be available due to privacy concerns. This
paper studies the group-oblivious setting, which cannot obtain group
information for each example at training time. This requires the model to
identify underperforming samples and then pay more attention to them during
training.

###  3.2 Importance-weighted mixup

UMix employs an aggressive data augmentation strategy called uncertainty-aware
mixup to mitigate overfitting. Specifically, vanilla mixup [75, 76] constructs
virtual training examples (i.e., mixed samples) by performing linear
interpolations between data/features and corresponding labels as:

$\widetilde{x}_{i,j}=\lambda
x_{i}+(1-\lambda)x_{j},\;\widetilde{y}_{i,j}=\lambda y_{i}+(1-\lambda)y_{j},$
(1)

where $(x_{i},y_{i}),{(x_{j},y_{j})}$ are two samples drawn at random from
empirical training distribution and $\lambda\in[0,1]$ is usually sampled from
a beta distribution. Then vanilla mixup optimizes the following loss function:

$\mathbb{E}_{\\{(x_{i},y_{i}),(x_{j},y_{j})\\}}[\ell(\theta,\widetilde{x}_{i,j},\widetilde{y}_{i,j})].$
(2)

When the cross entropy loss is employed, Eq. 2 can be rewritten as:

$\mathbb{E}_{\\{(x_{i},y_{i}),(x_{j},y_{j})\\}}[\lambda\ell(\theta,\widetilde{x}_{i,j},y_{i})+(1-\lambda)\ell(\theta,\widetilde{x}_{i,j},y_{j})].$
(3)

Eq. 3 can be seen as a linear combination (mixup) of
$\ell(\theta,\widetilde{x}_{i,j},y_{i})$ and
$\ell(\theta,\widetilde{x}_{i,j},y_{j})$. Unfortunately, since vanilla mixup
doesn’t consider the subpopulations with poor performance, it has been shown
experimentally to be non-robust against subpopulation shift [70]. To this end,
we introduce a simple yet effective method called UMix, which further employs
a weighted linear combination of the original loss based on Eq. 3 to encourage
the learned model to pay more attention to samples with poor performance.

In contrast to previous IW methods, the importance weights of UMix are used on
the mixed samples. To do this, we first estimate the uncertainty of each
sample and then use this quantity to construct the importance weight (i.e.,
the higher the uncertainty, the higher the weight, and vice versa). For the
$i$-th sample $x_{i}$, we denote its importance weight as $w_{i}$. Once we
obtain the importance weight, we can perform weighted linear combination of
$\ell(\theta,\widetilde{x}_{i,j},y_{i})$ and
$\ell(\theta,\widetilde{x}_{i,j},y_{j})$ by:

$\mathbb{E}_{\\{(x_{i},y_{i}),(x_{j},y_{j})\\}}[{\color[rgb]{.75,0,.25}{w}_{i}}\lambda\ell(\theta,\widetilde{x}_{i,j},y_{i})+{\color[rgb]{.75,0,.25}w_{j}}(1-\lambda)\ell(\theta,\widetilde{x}_{i,j},y_{j})],$
(4)

where ${w}_{i}$ and ${w}_{j}$ denote the importance weight of the $i$-th and
$j$-th samples respectively. In practice, to balance the UMix and normal
training, we set a hyperparameter $\sigma$ that denotes the probability to
apply UMix. The whole training pseudocode for UMix is shown in Algorithm 1.

Input: Training dataset $\mathcal{D}$ and the corresponding importance weights
$\mathbf{w}=[w_{1},\cdots,w_{N}]$, hyperparameter $\sigma$ to control the
probability of doing UMix, and parameter $\alpha$ of the beta distribution;

1 for  _each iteration_ do

2 Obtain training samples $(x_{i},y_{i})$, $(x_{j},y_{j})$ and the
corresponding weight $w_{i}$, $w_{j}$;

3 Sample $p\sim$ Uniform(0,1);

4 if $p<\sigma$ then Sample $\lambda\sim Beta(\alpha,\alpha)$; else
$\lambda=0$;

5 Obtain the mixed input $\widetilde{x}_{i,j}$ where
$\widetilde{x}_{i,j}=\lambda x_{i}+(1-\lambda)x_{j}$;

6 Obtain the loss of the model with
${\color[rgb]{.75,0,.25}{w}_{i}}\lambda\ell(\theta,\widetilde{x}_{i,j},y_{i})+{\color[rgb]{.75,0,.25}w_{j}}(1-\lambda)\ell(\theta,\widetilde{x}_{i,j},y_{j})$;

7 Update model parameters $\theta$ to minimize loss with an optimization
algorithm.

Algorithm 1 The training pseudocode of UMix.

###  3.3 Uncertainty-aware importance weights

Now we present how to obtain the uncertainty-aware training importance
weights. In the group-oblivious setting, the key to obtaining importance
weights is to find samples with high uncertainty. For example, DRO-based
algorithms construct the uncertainty set with the current loss [52, 24, 38,
25]. It has been shown experimentally that the uncertain samples found in this
way are constantly changing during training [41], resulting in these methods
not always upweighting the minority subpopulations. Therefore, we introduce a
sampling-based stable uncertainty estimation to better characterize the
subpopulation shift.

Given a well trained neural classifier
$f_{\theta}:\mathcal{X}\rightarrow\mathcal{Y}$ that could produce the
predicted class $\hat{f}_{\theta}(x)$, a simple way to obtain the uncertainty
of a sample is whether the sample is correctly classified. However, as pointed
out in previous work [36], a single model cannot accurately characterize the
sampling uncertainty. Therefore, we propose to obtain the uncertainty through
Bayesian sampling from the model posterior distribution
$p(\theta;\mathcal{D})$. Specifically, given a sample $(x_{i},y_{i})$, we
define the training uncertainty as:

$u_{i}=\int\kappa(y_{i},\hat{f_{\theta}}(x_{i}))p(\theta;\mathcal{D})d\theta,\text{where}\;\kappa(y_{i},\hat{f}_{\theta}(x_{i}))=\begin{cases}0,&\text{
if }y_{i}=\hat{f}_{\theta}(x_{i})\\\ 1,&\text{ if
}y_{i}\neq\hat{f}_{\theta}(x_{i})\end{cases}.$ (5)

Then, we can obtain an approximation of Eq. 5 with $T$ Monte Carlo samples as
$u_{i}\approx\frac{1}{T}\sum_{t=1}^{T}\kappa(y_{i},\hat{f}_{\theta_{t}}(x_{i}))$,
where $\theta_{t}\in\Theta$ can be obtained by minimizing the expected risk.

In practice, sampling $\\{\theta_{t}\\}_{t=1}^{T}$ from the posterior (i.e.,
$\theta_{t}\sim p(\theta;\mathcal{D})$) is computationally expensive and
sometimes even intractable since multiple training models need to be built or
extra approximation errors need to be introduced. Inspired by a recent
Bayesian learning paradigm named SWAG [46], we propose to employ the
information from the historical training trajectory to approximate the
sampling process. More specifically, we train a model with ERM and save the
prediction results $\hat{f}_{\theta_{t}}(x_{i})$ of each sample on each
iteration epoch $t$. Then, to avoid the influence of inaccurate predictions at
the beginning of training, we estimate uncertainty with predictions after
training $T_{s}-1$ epochs with:

$u_{i}\approx\frac{1}{T}\sum_{t=T_{s}}^{T_{s}+T}\kappa(y_{i},\hat{f}_{\theta_{t}}(x_{i})).$
(6)

We empirically show that the proposed approximation could obtain reliable
uncertainty in Sec. B.4 of the Appendix.

To obtain reasonable importance weights, we assume that the samples with high
uncertainty should be given a higher weight and vice versa. Therefore a
reasonable importance weight could be linearly positively related to the
corresponding uncertainty,

$w_{i}=\eta u_{i}+c,$ (7)

where $\eta\in\mathbb{R}_{+}$ is a hyperparameter and $c\in\mathbb{R}_{+}$ is
a constant that keeps the weight to be positive. In practice, we set $c$ to 1.
The whole process for obtaining training importance weights is shown in
Algorithm 2.

Input: Training dataset $\mathcal{D}$, sampling start epoch $T_{s}$, the
number of sampling $T$, and upweight hyperparameter $\eta$ ;

Output: The training importance weights $\mathbf{w}=[w_{1},\cdots,w_{n}]$;

1 for  _each iteration_ do

2 Train $f_{\theta}$ by minimizing the expected risk
$\mathbb{E}\\{\ell(\theta,x_{i},y_{i})\\}$;

3 Save the prediction results $\\{\hat{f}_{\theta_{t}}(x_{i})\\}_{i=1}^{N}$ of
the current epoch $t$;

4

5Obtain the uncertainty of each sample with
$u_{i}\approx\frac{1}{T}\sum_{t=T_{s}}^{T_{s}+T}\kappa(y_{i},\hat{f}_{\theta_{t}}(x_{i}))$;

Obtain the importance weight of each sample with $w_{i}=\eta u_{i}+c$.

Algorithm 2 The process for obtaining training importance weights.

Remark. Total uncertainty can be divided into epistemic and aleatoric
uncertainty [30]. In the proposed method, the samples are weighted only based
on epistemic uncertainty by sampling from the model on the training
trajectory, which can be seen as sampling from the model posterior in a more
efficient way. What’s more, we consider that the training samples do not
contain the inherent noise (aleatoric uncertainty) since it is usually
intractable to distinguish between noisy samples and minority samples from
data with subpopulation shifts.

Rethink why this estimation approach could work? Recent work has empirically
shown that compared with the hard-to-classify samples, the easy-to-classify
samples are learned earlier during training [18]. Meanwhile, the hard-to-
classify samples are also more likely to be forgotten by the neural networks
[64]. The frequency with which samples are correctly classified during
training can be used as supervision information in confidence calibration
[51]. Snapshot performs ensemble learning on several local minima models along
the optimization path [26]. The proposed method is also inspired by these
observations and algorithms. During training, samples from the minority
subpopulations are classified correctly less frequently, which corresponds to
higher training uncertainty. On the other hand, samples from the majority
subpopulations will have lower training uncertainty due to being classified
correctly more often. In Sec. B.5 of the Appendix, we show the accuracy of
different subpopulations during training to empirically validate our claim.
Meanwhile, we explain in detail why the uncertainty estimation based on
historical information is chosen in Sec. C of the Appendix.

##  4 Experiments

In this section, we conduct experiments on multiple datasets with
subpopulation shift to answer the following questions. Q1 Effectiveness (I).
In the group-oblivious setting, does the proposed method outperform other
algorithms? Q2 Effectiveness (II). How does UMIX perform without the group
labels in the validation set? Q3 Effectiveness (III). Although our method does
not use training group labels, does it perform better than the algorithms
using training group labels? Q4 Effectiveness (IV). Can UMix improve the model
robustness against domain shift where the training and test distributions have
different subpopulations. Q5 Qualitative analysis. Are the obtained
uncertainties of the training samples trustworthy? Q6 Ablation study. What is
the key factor of performance improvement in our method?

###  4.1 Setup

We briefly present the experimental setup here, including the experimental
datasets, evaluation metrics, model selection, and comparison methods. Please
refer to Sec. B in Appendix for more detailed setup.

Datasets. We perform experiments on three datasets with multiple
subpopulations, including Waterbirds [58], CelebA [43] and CivilComments [9].
We also validate UMix on domain shift scenario which is a more challenging
distribution shift problem since there are different subpopulations between
training and test data. Hence, we conduct experiments on a medical dataset
called Camelyon17 [5, 33] that consists of pathological images from five
different hospitals. The training data is drawn from three hospitals, while
the validation and test data are sampled from other hospitals.

Evaluation metrics. To be consistent with existing works [70, 33, 56], we
report the average accuracy of Camelyon17 over 10 different random seeds. On
other datasets, we repeat experiments over 3 times and report the average and
worst-case accuracy among all subpopulations. The trade-off between the
average and worst-case accuracy is a well-known challenge [21]. In this paper,
we lay emphasis on worst-case accuracy, which is more important than average
accuracy in some application scenarios. For example, in fairness-related
applications, we should pay more attention to the performance of the minority
groups to reduce the gap between the majority groups and ensure the fairness
of the machine learning decision system.

Model selection. Following prior works [41, 72], we assume the group labels of
validation samples are available and select the best model based on worst-case
accuracy among all subpopulations on the validation set. We also conduct model
selection based on the average accuracy to show the impact of validation group
label information in our method.

Comparisons in the group-oblivious setting. Here we list the baselines used in
the group-oblivious setting. (1) ERM trains the model using standard empirical
risk minimization. (2) Focal loss [40] downweights the well-classified
examples’ loss according to the current classification confidences. (3) DRO-
based methods including CVaR-DRO, $\chi^{2}$-DRO [38], CVaR-DORO and
$\chi^{2}$-DORO [72] minimize the loss over the worst-case distribution in a
neighborhood of the empirical training distribution. (4) JTT [41] constructs
an error set and upweights the samples in the error set to improve the worst-
case performance among all subpopulations.

Comparison in the group-aware setting. To better demonstrate the performance
of the proposed method, we compare our method with multiple methods that use
training group labels, including IRM [3], IB-IRM [1], V-REx [34], CORAL [63],
Group DRO [58], DomainMix [69], Fish [60], and LISA [70].

Mixup-based comparison methods. We compare our method with vanilla mixup and
in-group mixup, where vanilla mixup is performed on any pair of samples and
in-group mixup is performed on the samples with the same labels and from the
same subpopulations.

###  4.2 Experimental results

We present experimental results and discussions to answer the above-posed
questions.

Q1 Effectiveness (I). Since our algorithm does not need training group labels,
thus we conduct experiments to verify its superiority over current group-
oblivious algorithms. The experimental results are shown in Table 1 and we
have the following observations: (1) The proposed UMix achieves the best
worst-case accuracy on all three datasets. For example, for the CelebA
dataset, UMix achieves worst-case accuracy of 85.3%, while the second-best is
81.1%. (2) ERM consistently outperforms other methods in terms of average
accuracy. However, it typically comes with the lowest worst-case accuracy. The
underlying reason is that the dominance of the majority subpopulations during
training leads to poor performance of the minority subpopulations. (3) UMix
shows competitive average accuracy compared to other methods. For example, on
CelebA, UMix achieves the average accuracy of 90.1%, which outperforms all
other IW/DRO methods.

Q2 Effectiveness (II). We conduct the evaluation on the Waterbirds and CelebA
datasets without using the validation set group label information.
Specifically, after each training epoch, we evaluate the performance of the
current model on the validation set and save the model with the best average
accuracy. Finally, we test the performance of the saved model on the test set.
The experimental results are shown in Table 2. From the experimental results,
we can observe that when the validation set group information is not used, the
worst-case accuracy of our method drops a little while the average accuracy
improves a little.

Q3 Effectiveness (III). We further conduct comparisons with algorithms that
require training group labels. The comparison results are shown in Table 3.
According to the experimental results, it is observed that the performance
from our UMix without using group label is quite competitive compared with
these group-aware algorithms. Specifically, benefiting from the uncertainty-
aware mixup, UMix usually performs in the top three in terms of both average
and worst-case accuracy. For example, on WaterBirds, UMix achieves the best
average accuracy of 93.0% and the second-best worst-case accuracy of 90.0%.

Table 1: Comparison results with other methods in the group-oblivious setting. The best results are in bold and blue. Full results with standard deviation are in the Table 6 in Appendix. | Waterbirds | CelebA | CivilComments | Camelyon17
---|---|---|---|---
| Avg. | Worst | Avg. | Worst | Avg. | Worst | Avg.
ERM | 97.0% | 63.7% | 94.9% | 47.8% | 92.2% | 56.0% | 70.3%
Focal Loss [40] | 87.0% | 73.1% | 88.4% | 72.1% | 91.2% | 60.1% | 68.1%
CVaR-DRO [38] | 90.3% | 77.2% | 86.8% | 76.9% | 89.1% | 62.3% | 70.5%
CVaR-DORO [72] | 91.5% | 77.0% | 89.6% | 75.6% | 90.0% | 64.1% | 67.3%
$\chi^{2}$-DRO [38] | 88.8% | 74.0% | 87.7% | 78.4% | 89.4% | 64.2% | 68.0%
$\chi^{2}$-DORO [72] | 89.5% | 76.0% | 87.0% | 75.6% | 90.1% | 63.8% | 68.0%
JTT [41] | 93.6% | 86.0% | 88.0% | 81.1% | 90.7% | 67.4% | 69.1%
Ours | 93.0% | 90.0% | 90.1% | 85.3% | 90.6% |  70.1% | 75.1%
Table 2: Experimental results when the group labels in the validation set are available or not. Group labels in | Waterbirds | CelebA
---|---|---
validation set? | Average ACC | Worst-case ACC | Average ACC | Worst-case ACC
Yes | 93.00% | 90.00% | 90.10% | 85.30%
No | 93.60% | 88.90% | 90.40% | 84.60%
Table 3: Comparison results with the algorithms using training group labels (Our method is not dependent on this type of information). Results of baseline models are from [70]. The best three results are in bold brown or bold blue and the color indicates whether the training group labels are used. Full results with standard deviation are in the Table 7 in Appendix. |  Group labels | Waterbirds | CelebA | CivilComments |  Cam17
---|---|---|---|---|---
|  in train set? |  Avg. |  Worst | Avg. | Worst | Avg. | Worst | Avg.
IRM [3] |  Yes |  87.5% |  75.6% |  94.0% |  77.8% |  88.8% |  66.3% |  64.2%
IB-IRM [1] |  Yes |  88.5% |  76.5% |  93.6% |  85.0% |  89.1% |  65.3% |  68.9%
V-REx [34] |  Yes |  88.0% |  73.6% |  92.2% |  86.7% |  90.2% |  64.9% |  71.5%
CORAL [63] |  Yes |  90.3% |  79.8% |  93.8% |  76.9% |  88.7% |  65.6% |  59.5%
GroupDRO [58] |  Yes |  91.8% |  90.6% |  92.1% |  87.2% |  89.9% |  70.0% |  68.4%
DomainMix [69] |  Yes |  76.4% |  53.0% |  93.4% |  65.6% |  90.9% |  63.6% |  69.7%
Fish [60] |  Yes |  85.6% |  64.0% |  93.1% |  61.2% |  89.8% |  71.1% |  74.7%
LISA [70] |  Yes |  91.8% |  89.2% |  92.4% |  89.3% |  89.2% |  72.6% |  77.1%
Ours |  No |  93.0% |  90.0% |  90.1% |  85.3% |  90.6% |  70.1% |  75.1%

Q4 Effectiveness (IV). We conduct comparison experiments on Camelyon17 to
investigate the effectiveness of our algorithm under the domain shift
scenario. The experimental results are shown in the last column of Table 1 and
Table 3 respectively. In the group-oblivious setting, the proposed method
achieves the best average accuracy on Camelyon17 as shown in Table 1. For
example, UMix achieves the best average accuracy of 75.1% while the second is
70.3%. Meanwhile, in Table 3, benefiting from upweighting the mixed samples
with poor performance, our method achieves a quite competitive generalization
ability on Camelyon17 compared with other algorithms using training group
labels.

(a) Waterbirds

(b) CelebA

Figure 1: Visualization of the obtained uncertainty with kernel density
estimation on Waterbirds and CelebA datasets, where group size refers to the
sample number of the group.

Q5 Qualitative analysis. To intuitively investigate the rationality of the
estimated uncertainty, we visualize the density of the uncertainty for
different groups with kernel density estimation. As shown in Fig. 1, the
statistics of estimated uncertainty is basically correlated to the training
sample size of each group. For example, on Waterbirds and CelebA, the average
uncertainties of minority groups are much higher, while those of majority
groups are much lower.

Q6 Ablation study. Finally, we conduct the ablation study in comparison with
vanilla mixup and in-group mixup. The experimental results are shown in Table
4. Compared with ERM, vanilla mixup cannot significantly improve worst-case
accuracy. After using the group label, the in-group mixup slightly improves
the worst-case accuracy compared to ERM. The possible reason is that mixup-
based methods do not increase the influence of minority subpopulations in the
model objective function. Although our method does not use the group label of
the training samples, our method can still significantly improve the worst-
case accuracy.

Table 4: Comparison with ERM and mixup based methods. Results of baseline models are from [70]. The best results are in bold brown or bold blue and the color indicates whether the training group labels are used. Full results with standard deviation are in the Table 8 in Appendix. |  Group labels | Waterbirds | CelebA | CivilComments |  Cam17
---|---|---|---|---|---
|  in train set? |  Avg. |  Worst | Avg. | Worst | Avg. | Worst | Avg.
ERM |  No |  97.0% |  63.7% |  94.9% |  47.8% |  92.2% |  56.0% |  70.3%
vanilla mixup |  No |  81.0% |  56.2% |  95.8% |  46.4% |  90.8% |  67.2% |  71.2%
in-group mixup |  Yes |  88.7% |  68.0% |  95.2% |  58.3% |  90.8% |  69.2% |  75.5%
Ours |  No |  93.0% |  90.0% |  90.1% |  85.3% |  90.6% |  70.1% |  75.1%

##  5 Theory

In this section, we provide a theoretical understanding of the generalization
ability for UMix. At a high level, we prove that our method can achieve a
better generalization error bound than traditional IW methods without using
mixup. For simplicity, our analysis focuses on generalized linear model (GLM).
The roadmap of our analysis is to first approximate the mixup loss and then
study the generalization bound from a Rademacher complexity perspective. To
introduce the theoretical framework, we first present the basic settings.

Basic settings. Our analysis mainly focuses on GLM model classes whose loss
function $\ell$ follows $\ell(\theta,x,y)=A(\theta^{\top}x)-y\theta^{\top}x$,
where $x\in\mathbb{R}^{d}$ is the input , $\theta\in\mathbb{R}^{d}$ is the
parameter, $y\in\mathbb{R}$ is the label and $A(\cdot)$ is the log-partition
function.

Recall the setting of subpopulation shift, we assume that the population
distribution $P$ consists of $G$ different subpopulations with the $g$-th
subpopulation’s proportion being $k_{g}$ and the $g$-th subpopulation follows
the distribution $P_{g}$. Specifically, we have $P=\sum_{g=1}^{G}k_{g}P_{g}$.
Then we denote the covariance matrix for the $g$-th subpopulation as
$\Sigma_{X}^{g}=\mathbb{E}_{(x,y)\sim P_{g}}[xx^{\top}]$. For simplicity, we
consider the case where a shared weight $w_{g}$ is assigned to all samples
from the $g$-th subpopulation. The main goal of our theoretical analysis is to
characterize the generalization ability of the model learned using Algorithm
1. Formally, we focus on analyzing the upper bound of the weighted
generalization error defined as:

$\displaystyle\operatorname{GError}(\theta)=\mathbb{E}_{(x,y)\sim
P}[w(x,y)\ell(\theta,x,y)]-\frac{1}{N}\sum_{i=1}^{N}w(x_{i},y_{i})\ell(\theta,x_{i},y_{i}),$

where the function $w(x,y)$ is the weighted function to return the weight of
the subpopulation to which the sample $(x,y)$ belongs.

First of all, we present our main result in this section. The main theorem of
our analysis provides a subpopulation-heterogeneity dependent bound for the
above generalization error. This theorem is formally presented as:

######  Theorem 5.1.

Suppose $A(\cdot)$ is $L_{A}$-Lipschitz continuous, then there exists
constants $L,B>0$ such that for any $\theta$ satisfying
$\theta^{\top}\Sigma_{X}\theta\leq\gamma$, the following holds with a
probability of at least $1-\delta$,

$\displaystyle\operatorname{GError}(\theta)\leq 2L\cdot
L_{A}\cdot(\max\\{(\frac{\gamma(\delta/2)}{\rho})^{1/4},(\frac{\gamma(\delta/2)}{\rho})^{1/2}\\}\cdot\sqrt{{\color[rgb]{.75,0,.25}\frac{\operatorname{rank}(\Sigma_{X})}{n}}})+B\sqrt{\frac{\log(2/\delta)}{2n}},$

where $\gamma(\delta)$ is a constant dependent on $\delta$,
$\Sigma_{X}=\sum_{g=1}^{G}k_{g}w_{g}\Sigma_{X}^{g}$ and $\rho$ is some
constant related to the data distribution, which will be formally introduced
in Assumption 5.1.

We will show later that the output of our Algorithm 1 can satisfy the
constraint $\theta^{\top}\Sigma_{X}\theta\leq\gamma$ and thus Theorem 5.1 can
provide a theoretical understanding of our algorithm. In contrast to weighted
ERM, the bound improvement of UMix is on the red term which can partially
reflect the heterogeneity of the training subpopulations. Specifically, the
red term would become $\sqrt{d/n}$ in the weighted ERM setting (see more
detailed theoretical comparisons in Appendix). Thus our bound can be tighter
when the intrinsic dimension of data is small (i.e.,
$\text{rank}(\Sigma_{X})\ll d$).

The proof of Theorem 5.1 follows this roadmap: (1) We first show that the
model learned with UMix can fall into a specific hypothesis set
$\mathcal{W}_{\gamma}$. (2) We analyze the Rademacher complexity of the
hypothesis set and obtain its complexity upper bound (Lemma A.3). (3) Finally,
we can characterize the generalization bound by using complexity-based
learning theory [7] (Theorem 8). More details of the proof can be found in
Appendix.

As we discuss in Appendix, the weighted mixup can be seen as an approximation
of a regularization term
$\frac{C}{n}[\sum_{i=1}^{n}w_{i}A^{\prime\prime}(x_{i}^{\top}\theta)]\theta^{\top}\widehat{\Sigma}_{X}\theta$
for some constant $C$ compared with the non-mixup algorithm, which motivates
us to study the following hypothesis space

$\displaystyle\mathcal{W}_{\gamma}\coloneqq\\{x\rightarrow\theta^{\top}x,\text{such
that }\theta\text{ satisfying
}\mathbb{E}_{x,y}[w(x,y)A^{\prime\prime}(x^{\top}\theta)]\theta^{\top}\Sigma_{X}\theta\leq\gamma\\},$

for some constant $\gamma$.

To further derive the generalization bound, we also need the following
assumption, which is satisfied by general GLMs when $\theta$ has bounded
$\ell_{2}$ norm and it is adopted in, e.g., [4, 76].

######  Assumption 5.1 ($\rho$-retentive).

We say the distribution of $x$ is $\rho$-retentive for some $\rho\in(0,1/2]$
if for any non-zero vector $v\in\mathbb{R}^{d}$ and given the event that
$\theta\in\mathcal{W}_{\gamma}$ where the $\theta$ is output by our Algorithm
1, we have

$\displaystyle\mathbb{E}_{x}^{2}[A^{\prime\prime}(x^{\top}v)]\geq\rho\cdot\min\\{1,\mathbb{E}_{x}(v^{\top}x)^{2}\\}.$

Finally, we can derive the Rademacher complexity of the $\mathcal{W}_{\gamma}$
and the proof of Theorem 5.1 is obtained by combining Lemma A.3 and the
Theorem 8 of [7].

######  Lemma 5.1.

Assume that the distribution of $x_{i}$ is $\rho$-retentive, i.e., satisfies
the assumption 5.1. Then the empirical Rademacher complexity of
$\mathcal{W}_{r}$ satisfies

$\displaystyle
Rad(\mathcal{W}_{r},\mathcal{S})\leq\max\\{(\frac{\gamma(\delta)}{\rho})^{1/4},(\frac{\gamma(\delta)}{\rho})^{1/2}\\}\cdot\sqrt{\frac{rank(\Sigma_{X})}{n}},$

with probability at least $1-\delta$.

##  6 Conclusion

In this paper, we propose a novel method called UMix to improve the model
robustness against subpopulation shift. We propose a simple yet reliable
approach to estimate the sample uncertainties and integrate them into the
mixup strategy so that UMix can mitigate the overfitting thus improving over
prior IW methods. Our method consistently outperforms previous approaches on
commonly-used benchmarks. Furthermore, UMix also shows the theoretical
advantage that the learned model comes with subpopulation-heterogeneity
dependent generalization bound. In the future, how to leverage subpopulation
information to improve UMix can be a promising research direction.

## Acknowledgements

This work was supported in part by the National Key Research and Development
Program of China under Grant 2019YFB2101900, the National Natural Science
Foundation of China (61976151, 61925602, 61732011).

## References

  * [1] Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Jean-Christophe Gagnon-Audet, Yoshua Bengio, Ioannis Mitliagkas, and Irina Rish.  Invariance principle meets information bottleneck for out-of-distribution generalization.  In NeurIPS, 2021.
  * [2] Javier Antorán, James Urquhart Allingham, and José Miguel Hernández-Lobato.  Depth uncertainty in neural networks.  In NeurIPS, 2020.
  * [3] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz.  Invariant risk minimization.  arXiv preprint arXiv:1907.02893, 2019.
  * [4] Raman Arora, Peter Bartlett, Poorya Mianjy, and Nathan Srebro.  Dropout: Explicit forms and capacity control.  In ICML, 2021.
  * [5] Peter Bandi, Oscar Geessink, Quirine Manson, Marcory Van Dijk, Maschenka Balkenhol, Meyke Hermsen, Babak Ehteshami Bejnordi, Byungjae Lee, Kyunghyun Paeng, Aoxiao Zhong, et al.  From detection of individual metastases to classification of lymph node status at the patient level: the camelyon17 challenge.  IEEE transactions on medical imaging, 38(2):550–560, 2018.
  * [6] Solon Barocas and Andrew D Selbst.  Big data’s disparate impact.  Calif. L. Rev., 104:671, 2016.
  * [7] Peter L Bartlett and Shahar Mendelson.  Rademacher and gaussian complexities: Risk bounds and structural results.  Journal of Machine Learning Research, 3(Nov):463–482, 2002.
  * [8] Steffen Bickel, Michael Brückner, and Tobias Scheffer.  Discriminative learning for differing training and test distributions.  In ICML, 2007.
  * [9] Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman.  Nuanced metrics for measuring unintended bias with real data for text classification.  In WWW, 2019.
  * [10] Jonathon Byrd and Zachary Lipton.  What is the effect of importance weighting in deep learning?  In ICML, 2019.
  * [11] Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, and Tengyu Ma.  Learning imbalanced datasets with label-distribution-aware margin loss.  In NeurIPS, 2019.
  * [12] Yongqiang Chen, Yonggang Zhang, Yatao Bian, Han Yang, Kaili Ma, Binghui Xie, Tongliang Liu, Bo Han, and James Cheng.  Invariance principle meets out-of-distribution generalization on graphs.  arXiv preprint arXiv:2202.05441, 2022.
  * [13] Yin Cui, Menglin Jia, Tsung-Yi Lin, Yang Song, and Serge Belongie.  Class-balanced loss based on effective number of samples.  In CVPR, 2019.
  * [14] Didan Deng, Liang Wu, and Bertram E Shi.  Iterative distillation for better uncertainty estimates in multitask emotion recognition.  In ICCV, 2021.
  * [15] John Denker and Yann LeCun.  Transforming neural-net output levels to probability distributions.  In NeurIPS, 1990.
  * [16] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  In NAACL, 2019.
  * [17] Yarin Gal and Zoubin Ghahramani.  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In ICML, 2016.
  * [18] Yonatan Geifman, Guy Uziel, and Ran El-Yaniv.  Bias-reduced uncertainty estimation for deep neural classifiers.  In ICLR, 2019.
  * [19] Yu Geng, Zongbo Han, Changqing Zhang, and Qinghua Hu.  Uncertainty-aware multi-view representation learning.  In AAAI, 2021.
  * [20] Zongbo Han, Changqing Zhang, Huazhu Fu, and Joey Tianyi Zhou.  Trusted multi-view classification with dynamic evidential fusion.  IEEE TPAMI, 2022.
  * [21] Tatsunori Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang.  Fairness without demographics in repeated loss minimization.  In ICML, 2018.
  * [22] Marton Havasi, Rodolphe Jenatton, Stanislav Fort, Jeremiah Zhe Liu, Jasper Snoek, Balaji Lakshminarayanan, Andrew Mingbo Dai, and Dustin Tran.  Training independent subnetworks for robust prediction.  In ICLR, 2021.
  * [23] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In CVPR, 2016.
  * [24] Weihua Hu, Gang Niu, Issei Sato, and Masashi Sugiyama.  Does distributionally robust supervised learning give robust classifiers?  In ICML, 2018.
  * [25] Zhaolin Hu and L Jeff Hong.  Kullback-leibler divergence constrained distributionally robust optimization.  Available at Optimization Online, pages 1695–1724, 2013.
  * [26] Gao Huang, Yixuan Li, Geoff Pleiss, Zhuang Liu, John E Hopcroft, and Kilian Q Weinberger.  Snapshot ensembles: Train 1, get m for free.  In ICLR, 2017.
  * [27] Jiayuan Huang, Arthur Gretton, Karsten Borgwardt, Bernhard Schölkopf, and Alex Smola.  Correcting sample selection bias by unlabeled data.  In NeurIPS, 2006.
  * [28] Nathalie Japkowicz.  The class imbalance problem: Significance and strategies.  In IJCAI, 2000.
  * [29] Gabriel Kalweit and Joschka Boedecker.  Uncertainty-driven imagination for continuous deep reinforcement learning.  In Conference on Robot Learning, pages 195–206. PMLR, 2017.
  * [30] Alex Kendall and Yarin Gal.  What uncertainties do we need in bayesian deep learning for computer vision?  In NeurIPS, 2017.
  * [31] Alex Kendall, Yarin Gal, and Roberto Cipolla.  Multi-task learning using uncertainty to weigh losses for scene geometry and semantics.  In CVPR, 2018.
  * [32] Allison Koenecke, Andrew Nam, Emily Lake, Joe Nudell, Minnie Quartey, Zion Mengesha, Connor Toups, John R Rickford, Dan Jurafsky, and Sharad Goel.  Racial disparities in automated speech recognition.  Proceedings of the National Academy of Sciences, 117(14):7684–7689, 2020.
  * [33] Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al.  Wilds: A benchmark of in-the-wild distribution shifts.  In ICML, 2021.
  * [34] David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville.  Out-of-distribution generalization via risk extrapolation (rex).  In ICML, 2021.
  * [35] Preethi Lahoti, Alex Beutel, Jilin Chen, Kang Lee, Flavien Prost, Nithum Thain, Xuezhi Wang, and Ed Chi.  Fairness without demographics through adversarially reweighted learning.  In NeurIPS, 2020.
  * [36] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell.  Simple and scalable predictive uncertainty estimation using deep ensembles.  In NeurIPS, 2017.
  * [37] Quoc V Le, Alex J Smola, and Stéphane Canu.  Heteroscedastic gaussian process regression.  In ICML, 2005.
  * [38] Daniel Levy, Yair Carmon, John C Duchi, and Aaron Sidford.  Large-scale methods for distributionally robust optimization.  In NeurIPS, 2020.
  * [39] Kevin Li, Abhishek Gupta, Ashwin Reddy, Vitchyr H Pong, Aurick Zhou, Justin Yu, and Sergey Levine.  Mural: Meta-learning uncertainty-aware rewards for outcome-driven reinforcement learning.  In ICML, 2021.
  * [40] Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár.  Focal loss for dense object detection.  In ICCV, 2017.
  * [41] Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn.  Just train twice: Improving group robustness without training group information.  In ICML, 2021.
  * [42] Wei Liu and Sanjay Chawla.  Class confidence weighted knn algorithms for imbalanced data sets.  In Pacific-Asia conference on knowledge discovery and data mining, pages 345–356. Springer, 2011.
  * [43] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang.  Deep learning face attributes in the wild.  In ICCV, 2015.
  * [44] Huan Ma, Zongbo Han, Changqing Zhang, Huazhu Fu, Joey Tianyi Zhou, and Qinghua Hu.  Trustworthy multimodal regression with mixture of normal-inverse gamma distributions.  NeurIPS, 2021.
  * [45] David JC MacKay.  A practical bayesian framework for backpropagation networks.  Neural computation, 4(3):448–472, 1992.
  * [46] Wesley J Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson.  A simple baseline for bayesian uncertainty in deep learning.  In NeurIPS, 2019.
  * [47] Paul Michel, Tatsunori Hashimoto, and Graham Neubig.  Modeling the second player in distributionally robust optimization.  In ICLR, 2021.
  * [48] Paul Michel, Tatsunori Hashimoto, and Graham Neubig.  Distributionally robust models with parametric likelihood ratios.  In ICLR, 2022.
  * [49] Microsoft.  Neural Network Intelligence, 1 2021.
  * [50] John Stuart Mill.  Utilitarianism.  In Seven masterpieces of philosophy, pages 337–383. Routledge, 2016\.
  * [51] Jooyoung Moon, Jihyo Kim, Younghak Shin, and Sangheum Hwang.  Confidence-aware learning for deep neural networks.  In ICML, 2020.
  * [52] Hongseok Namkoong and John C Duchi.  Stochastic gradient methods for distributionally robust optimization with f-divergences.  In NeurIPS, 2016.
  * [53] Radford M Neal.  Bayesian learning for neural networks, volume 118.  Springer Science & Business Media, 2012.
  * [54] David A Nix and Andreas S Weigend.  Estimating the mean and variance of the target probability distribution.  In ICNN, 1994.
  * [55] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  In NeurIPS, 2019.
  * [56] Vihari Piratla, Praneeth Netrapalli, and Sunita Sarawagi.  Focus on the common good: Group distributional robustness follows.  In ICLR, 2022.
  * [57] John Rawls.  Justice as fairness: A restatement.  Harvard University Press, 2001.
  * [58] Shiori Sagawa*, Pang Wei Koh*, Tatsunori B. Hashimoto, and Percy Liang.  Distributionally robust neural networks.  In ICLR, 2020.
  * [59] Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang.  An investigation of why overparameterization exacerbates spurious correlations.  In ICML, 2020.
  * [60] Yuge Shi, Jeffrey Seely, Philip HS Torr, N Siddharth, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve.  Gradient matching for domain generalization.  arXiv preprint arXiv:2104.09937, 2021.
  * [61] Hidetoshi Shimodaira.  Improving predictive inference under covariate shift by weighting the log-likelihood function.  Journal of statistical planning and inference, 90(2):227–244, 2000\.
  * [62] Jun Shu, Qi Xie, Lixuan Yi, Qian Zhao, Sanping Zhou, Zongben Xu, and Deyu Meng.  Meta-weight-net: Learning an explicit mapping for sample weighting.  In NeurIPS, 2019.
  * [63] Baochen Sun and Kate Saenko.  Deep coral: Correlation alignment for deep domain adaptation.  In European conference on computer vision, pages 443–450. Springer, 2016.
  * [64] Mariya Toneva, Alessandro Sordoni, Remi Tachet des Combes, Adam Trischler, Yoshua Bengio, and Geoffrey J. Gordon.  An empirical study of example forgetting during deep neural network learning.  In ICLR, 2019.
  * [65] Vikas Verma, Alex Lamb, Christopher Beckham, Amir Najafi, Ioannis Mitliagkas, David Lopez-Paz, and Yoshua Bengio.  Manifold mixup: Better representations by interpolating hidden states.  In ICML, 2019.
  * [66] Junfeng Wen, Chun-Nam Yu, and Russell Greiner.  Robust learning under uncertain test distributions: Relating covariate shift to model misspecification.  In ICML, 2014.
  * [67] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al.  Huggingface’s transformers: State-of-the-art natural language processing.  arXiv preprint arXiv:1910.03771, 2019.
  * [68] Da Xu, Yuting Ye, and Chuanwei Ruan.  Understanding the role of importance weighting for deep learning.  In ICLR, 2021.
  * [69] Minghao Xu, Jian Zhang, Bingbing Ni, Teng Li, Chengjie Wang, Qi Tian, and Wenjun Zhang.  Adversarial domain adaptation with domain mixup.  In AAAI, 2020.
  * [70] Huaxiu Yao, Yu Wang, Sai Li, Linjun Zhang, Weixin Liang, James Zou, and Chelsea Finn.  Improving out-of-distribution robustness via selective augmentation.  In ICML, 2022.
  * [71] Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo.  Cutmix: Regularization strategy to train strong classifiers with localizable features.  In CVPR, 2019.
  * [72] Runtian Zhai, Chen Dan, Zico Kolter, and Pradeep Ravikumar.  Doro: Distributional and outlier robust optimization.  In ICML, 2021.
  * [73] Runtian Zhai, Chen Dan, Zico Kolter, and Pradeep Ravikumar.  Understanding why generalized reweighting does not improve over erm.  arXiv preprint arXiv:2201.12293, 2022.
  * [74] Runtian Zhai, Chen Dan, Arun Suggala, J Zico Kolter, and Pradeep Ravikumar.  Boosted cvar classification.  In NeurIPS, 2021.
  * [75] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz.  mixup: Beyond empirical risk minimization.  In ICLR, 2018.
  * [76] Linjun Zhang, Zhun Deng, Kenji Kawaguchi, Amirata Ghorbani, and James Zou.  How does mixup help with robustness and generalization?  In ICLR, 2021.

## Checklist

  1. 1.

For all authors…

    1. (a)

Do the main claims made in the abstract and introduction accurately reflect
the paper’s contributions and scope? [Yes]

    2. (b)

Did you describe the limitations of your work? [Yes] See Sec. D in Appendix.

    3. (c)

Did you discuss any potential negative societal impacts of your work? [Yes]
See Sec. D in Appendix.

    4. (d)

Have you read the ethics review guidelines and ensured that your paper
conforms to them? [Yes]

  2. 2.

If you are including theoretical results…

    1. (a)

Did you state the full set of assumptions of all theoretical results? [Yes]
See Sec. 5.

    2. (b)

Did you include complete proofs of all theoretical results? [Yes] See Sec. A
in Appendix.

  3. 3.

If you ran experiments…

    1. (a)

Did you include the code, data, and instructions needed to reproduce the main
experimental results (either in the supplemental material or as a URL)? [Yes]
Code has been released.

    2. (b)

Did you specify all the training details (e.g., data splits, hyperparameters,
how they were chosen)? [Yes] See Sec. B in Appendix.

    3. (c)

Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [Yes] See Sec. B in Appendix.

    4. (d)

Did you include the total amount of compute and the type of resources used
(e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Sec. B in
Appendix.

  4. 4.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets…

    1. (a)

If your work uses existing assets, did you cite the creators? [Yes]

    2. (b)

Did you mention the license of the assets? [Yes]

    3. (c)

Did you include any new assets either in the supplemental material or as a
URL? [No]

    4. (d)

Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? [No] The datasets used are all publicly available
datasets.

    5. (e)

Did you discuss whether the data you are using/curating contains personally
identifiable information or offensive content? [No] The datasets used are all
publicly available datasets.

  5. 5.

If you used crowdsourcing or conducted research with human subjects…

    1. (a)

Did you include the full text of instructions given to participants and
screenshots, if applicable? [N/A] We didn’t conduct research with human
subjects.

    2. (b)

Did you describe any potential participant risks, with links to Institutional
Review Board (IRB) approvals, if applicable? [N/A] We didn’t conduct research
with human subjects.

    3. (c)

Did you include the estimated hourly wage paid to participants and the total
amount spent on participant compensation? [N/A] We didn’t conduct research
with human subjects.

Appendix

.tocmtappendix mtchapternone mtappendixsubsection

###### Contents

  1. 1 Introduction
  2. 2 Related Work
    1. 2.1 Importance weighting
    2. 2.2 Uncertainty quantification
  3. 3 Method
    1. 3.1 Background
    2. 3.2 Importance-weighted mixup
    3. 3.3 Uncertainty-aware importance weights
  4. 4 Experiments
    1. 4.1 Setup
    2. 4.2 Experimental results
  5. 5 Theory
  6. 6 Conclusion
  7. A Proofs
  8. B Experimental details
    1. B.1 Backbone model
    2. B.2 Datasets details
    3. B.3 Implementation details
    4. B.4 Uncertainty quantification results on simulated dataset
    5. B.5 Training accuracy throughout training
    6. B.6 Additional results
  9. C Justification for choosing historical-based uncertainty score
  10. D Societal impact and limitations
    1. D.1 Societal impact
    2. D.2 Limitations and future works

##  Appendix A Proofs

In this appendix, we prove the Theorem 5.1 in Section 5. We consider the
following optimization objective, which is the expected version of our
weighted mixup loss (Equation 4).

$\displaystyle
L_{n}^{\text{mix}}(\theta,S)=\frac{1}{n^{2}}\sum^{n}_{i,j=1}\mathbb{E}_{\lambda\sim
D_{\lambda}}[\lambda
w_{i}l(\theta,\tilde{x}_{i,j},y_{i})+(1-\lambda)w_{j}l(\theta,\tilde{x}_{i,j},y_{j})],$

where the loss function we consider is
$l(\theta,x,y)=h(f_{\theta}(x))-yf_{\theta}(x)$ and $h(\cdot)$ and
$f_{\theta}(\cdot)$ for all $\theta\in\Theta$ are twice differentiable. We
compare it with the standard weighted loss function

$\displaystyle
L_{n}^{std}(\theta,S)=\frac{1}{n}\sum_{i=1}^{n}w_{i}[h(f_{\theta}(x_{i}))-y_{i}f_{\theta}(x_{i})].$

######  Lemma A.1.

The weighted mixup loss can be rewritten as

$\displaystyle
L_{n}^{mix}(\theta,S)=L_{n}^{std}(\theta,S)+\sum_{i=1}^{3}\mathcal{R}_{i}(\theta,S)+\mathbb{E}_{\lambda\sim\tilde{\mathcal{D}}_{\lambda}}\left[(1-\lambda)^{2}\varphi(1-\lambda)\right],$

where $\tilde{\mathcal{D}}_{\lambda}$ is a uniform mixture of two Beta
distributions, i.e.,
$\frac{\alpha}{\alpha+\beta}Beta(\alpha+1,\beta)+\frac{\beta}{\alpha+\beta}Beta(\beta+1,\alpha)$
and $\psi(\cdot)$ is some function with $\lim_{a\rightarrow 0}\psi(a)=0$.
Moreover,

$\displaystyle\mathcal{R}_{1}(\theta,S)$
$\displaystyle=\frac{\mathbb{E}_{\lambda\sim\tilde{\mathcal{D}}_{\lambda}}[1-\lambda]}{n}\sum_{i=1}^{n}w_{i}\left(h^{\prime}\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}\right)\nabla
f_{\theta}\left(x_{i}\right)^{\top}\mathbb{E}_{r_{x}\sim\mathcal{D}_{X}}\left[r_{x}-x_{i}\right]$
$\displaystyle\mathcal{R}_{2}(\theta,S)$
$\displaystyle=\frac{\mathbb{E}_{\lambda\sim\tilde{\mathcal{D}}_{\lambda}}\left[(1-\lambda)^{2}\right]}{2n}\sum_{i=1}^{n}w_{i}h^{\prime\prime}\left(f_{\theta}\left(x_{i}\right)\right)\nabla
f_{\theta}\left(x_{i}\right)^{\top}\mathbb{E}_{r_{x}\sim\mathcal{D}_{X}}\left[\left(r_{x}-x_{i}\right)\left(r_{x}-x_{i}\right)^{\top}\right]\nabla
f_{\theta}\left(x_{i}\right)$ $\displaystyle\mathcal{R}_{3}(\theta,S)$
$\displaystyle=\frac{\mathbb{E}_{\lambda\sim\tilde{\mathcal{D}}_{\lambda}}\left[(1-\lambda)^{2}\right]}{2n}\sum_{i=1}^{n}w_{i}\left(h^{\prime}\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}\right)\mathbb{E}_{r_{x}\sim\mathcal{D}_{X}}\left[\left(r_{x}-x_{i}\right)\nabla^{2}f_{\theta}\left(x_{i}\right)\left(r_{x}-x_{i}\right)^{\top}\right].$

###### Proof.

The corresponding mixup version is

$\displaystyle L_{n}^{\text{mix}}(\theta,S)$
$\displaystyle=\frac{1}{n^{2}}\mathbb{E}_{\lambda\sim
Beta(\alpha,\beta)}\sum_{i,j=1}^{n}[\lambda
w_{i}h(f_{\theta}(\tilde{x}_{i,j}(\lambda)))-\lambda w_{i}y_{i}$
$\displaystyle\qquad\qquad\qquad\qquad\qquad+(1-\lambda)w_{j}h(f_{\theta}(\tilde{x}_{i,j}(\lambda)))-(1-\lambda)w_{j}y_{j}]$
$\displaystyle=\frac{1}{n^{2}}\mathbb{E}_{\lambda\sim
Beta(\alpha,\beta)}\mathbb{E}_{B\sim
Bern(\lambda)}\sum_{i,j=1}^{n}[w_{i}B(h(f_{\theta}(\tilde{x}_{i,j}))-y_{i})$
$\displaystyle\qquad\qquad\qquad\qquad\qquad+w_{j}(1-B)(h(f_{\theta}(\tilde{x}_{i,j}))-y_{j})]$
$\displaystyle=\frac{1}{n^{2}}\sum_{i,j=1}^{n}\\{\frac{\alpha}{\alpha+\beta}\mathbb{E}_{\lambda\sim
Beta(\alpha+1,\beta)}w_{i}[h(f_{\theta}(\tilde{x}_{i,j}))-y_{i}]$
$\displaystyle\qquad\qquad\qquad\qquad\qquad+\frac{\beta}{\alpha+\beta}\mathbb{E}_{\lambda\sim
Beta(\alpha,\beta+1)}w_{j}[h(f_{\theta}(\tilde{x}_{i,j}))-y_{j}])\\}$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}w_{i}\mathbb{E}_{\lambda\sim\tilde{D}_{\lambda}}\mathbb{E}_{r_{x}\sim
D_{x}^{w}}h(f(\theta,\lambda x_{i}+(1-\lambda)r_{x}))-y_{i}f(\theta,\lambda
x_{i}+(1-\lambda)r_{x})$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}w_{i}\mathbb{E}_{\lambda\sim\tilde{D}_{x}}l_{\check{x}_{i},y_{i}}(\theta),$

where $D_{x}^{w}=\frac{1}{n}\sum_{i=1}^{n}w_{i}\delta_{i}$ and
$\check{x}_{i}=\lambda x_{i}+(1-\lambda)r_{x}$.

We let $\alpha=1-\lambda$ and
$\psi_{i}(\alpha)=l_{\check{x}_{i},y_{i}}(\theta)$. Then since we know
$\psi_{i}$ is twice-differential, we have

$\displaystyle
l_{\breve{x}_{i},y_{i}}(\theta)=\psi_{i}(\alpha)=\psi_{i}(0)+\psi_{i}^{\prime}(0)\alpha+\frac{1}{2}\psi_{i}^{\prime\prime}(0)\alpha^{2}+\alpha^{2}\varphi_{i}(\alpha).$

By the proof of Lemma 3.1 in [76] we know

$\displaystyle\psi_{i}^{\prime}(0)$
$\displaystyle=\left(h^{\prime}\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}\right)\nabla
f_{\theta}\left(x_{i}\right)^{\top}\left(r_{x}-x_{i}\right),$
$\displaystyle\psi_{i}^{\prime\prime}(0)$
$\displaystyle=h^{\prime\prime}\left(f_{\theta}\left(x_{i}\right)\right)\nabla
f_{\theta}\left(x_{i}\right)^{\top}\left(r_{x}-x_{i}\right)\left(r_{x}-x_{i}\right)^{\top}\nabla
f_{\theta}\left(x_{i}\right)$
$\displaystyle\quad+\left(h^{\prime}\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}\right)\left(r_{x}-x_{i}\right)^{\top}\nabla^{2}f_{\theta}\left(x_{i}\right)\left(r_{x}-x_{i}\right).$

∎

######  Lemma A.2.

Consider the centralized dataset, i.e., $\frac{1}{n}\sum_{i=1}^{n}x_{i}=0$, we
have

$\displaystyle\mathbb{E}_{\lambda\sim\tilde{\mathcal{D}}_{\lambda}}[L_{n}^{mix}(\theta,\tilde{S})]\approx
L_{n}^{std}(\theta,S)+\frac{1}{2n}[\sum_{i=1}^{n}w_{i}A^{\prime\prime}(x_{i}^{\top}\theta)]\mathbb{E}_{\lambda\sim\tilde{\mathcal{D}}_{\lambda}}(\frac{(1-\lambda)^{2}}{\lambda^{2}})\theta^{\top}\widehat{\Sigma}_{X}\theta,$

where $\widehat{\Sigma}_{X}=\frac{1}{n}\sum_{i=1}^{n}w_{i}x_{i}x_{i}^{\top}$,
and the expectation is taken with respect to the randomness of $\lambda$.

###### Proof.

For GLM, the prediction is invariant to the scaling of the training data and
thus we consider the re-scaled dataset
$\tilde{S}=\\{(\tilde{x}_{i},y_{i})\\}_{i=1}^{n}$ where
$\tilde{x}_{i}=\frac{1}{\lambda}(\lambda x_{i}+(1-\lambda)r_{x})$. For GLM the
mixed stadard loss function is

$\displaystyle
L_{n}^{std}(\theta,\tilde{S})=\frac{1}{n}\sum_{i=1}^{n}w_{i}l_{\check{x}_{i},y_{i}}(\theta)=\frac{1}{n}\sum_{i=1}^{n}-w_{i}(y_{i}\tilde{x}_{i}^{\top}\theta-A(\tilde{x}_{i}^{\top}\theta)).$

In the proof of Lemma 3.3 in [76], we know by taking expectation with respect
to the randomness of $\lambda$ and $r_{x}$ we have the following second-order
approximation for the GLM loss,

$\displaystyle\mathbb{E}[L_{n}^{std}(\theta,\tilde{S})]\approx
L_{n}^{std}(\theta,S)+\frac{1}{2n}[\sum_{i=1}^{n}w_{i}A^{\prime\prime}(x_{i}^{\top}\theta)]\mathbb{E}(\frac{(1-\lambda)^{2}}{\lambda^{2}})\theta^{\top}\widehat{\Sigma}_{X}\theta,$

where $\widehat{\Sigma}_{X}=\frac{1}{n}\sum_{i=1}^{n}w_{i}x_{i}x_{i}^{\top}$.
∎

######  Lemma A.3.

Assume that the distribution of $x_{i}$ is $\rho$-retentive, i.e., satisfies
the Assumption 5.1. Then the empirical Rademacher complexity of
$\mathcal{W}_{r}$ satisfies

$\displaystyle
Rad(\mathcal{W}_{r},\mathcal{S})\leq\max\\{(\frac{\gamma(\delta)}{\rho})^{1/4},(\frac{\gamma(\delta)}{\rho})^{1/2}\\}\cdot\sqrt{\frac{rank(\Sigma_{X})}{n}},$

with probability at least $1-\delta$ for some constant $\gamma(\delta)$ that
only depends on $\delta$.

###### Proof.

The proof is mainly based on [76]. By definition, given $n$ i.i.d. Rademacher
rv. $\xi_{1},\ldots,\xi_{n}$, the empirical Rademacher complexity is

$\operatorname{Rad}\left(\mathcal{W}_{\gamma},S\right)=\mathbb{E}_{\xi}\sup_{a(\theta)\cdot\theta^{\top}\Sigma_{X}\theta\leq\gamma}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}\theta^{\top}x_{i}$

Let
$\tilde{x}_{i}=\Sigma_{X}^{\dagger/2}x_{i},a(\theta)=\mathbb{E}_{x}\left[A^{\prime\prime}\left(x^{\top}\theta\right)\right]$
and $v=\Sigma_{X}^{1/2}\theta$, then $\rho$-retentiveness condition implies
$a(\theta)^{2}\geq\rho\cdot\min\left\\{1,\mathbb{E}_{x}\left(\theta^{\top}x\right)^{2}\right\\}\geq\rho\cdot\min\left\\{1,\theta^{\top}\Sigma_{X}\theta\right\\}$
and therefore $a(\theta)\cdot\theta^{\top}\Sigma_{X}\theta\leq\gamma$ implies
that
$\|v\|^{2}=\theta^{\top}\Sigma_{X}\theta\leq\max\left\\{\left(\frac{\gamma}{\rho}\right)^{1/2},\frac{\gamma}{\rho}\right\\}$.

As a result,

$\displaystyle\operatorname{Rad}\left(\mathcal{W}_{\gamma},S\right)$
$\displaystyle=\mathbb{E}_{\xi}\sup_{a(\theta)\cdot\theta^{\top}\Sigma_{X}\theta\leq\gamma}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}\theta^{\top}x_{i}$
$\displaystyle=\mathbb{E}_{\xi}\sup_{a(\theta)\cdot\theta^{\top}\Sigma_{X}\theta\leq\gamma}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}v^{\top}\tilde{x}_{i}$
$\displaystyle\leq\mathbb{E}_{\xi}\sup_{\|v\|^{2}\leq\left(\frac{\gamma}{\rho}\right)^{1/2}\vee\frac{\gamma}{\rho}}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}v^{\top}\tilde{x}_{i}$
$\displaystyle\leq\frac{1}{n}\cdot\left(\frac{\gamma}{\rho}\right)^{1/4}\vee\left(\frac{\gamma}{\rho}\right)^{1/2}\cdot\mathbb{E}_{\xi}\left\|\sum_{i=1}^{n}\xi_{i}\tilde{x}_{i}\right\|$
$\displaystyle\leq\frac{1}{n}\cdot\left(\frac{\gamma}{\rho}\right)^{1/4}\vee\left(\frac{\gamma}{\rho}\right)^{1/2}\cdot\sqrt{\mathbb{E}_{\xi}\left\|\sum_{i=1}^{n}\xi_{i}\tilde{x}_{i}\right\|^{2}}$
$\displaystyle\leq\frac{1}{n}\cdot\left(\frac{\gamma}{\rho}\right)^{1/4}\vee\left(\frac{\gamma}{\rho}\right)^{1/2}\cdot\sqrt{\sum_{i=1}^{n}\tilde{x}_{i}^{\top}\tilde{x}_{i}}$

Consequently,

$\displaystyle\operatorname{Rad}\left(\mathcal{W}_{\gamma},S\right)=\mathbb{E}_{S}\left[\operatorname{Rad}\left(\mathcal{W}_{\gamma},S\right)\right]$
$\displaystyle\leq\frac{1}{n}\cdot\left(\frac{\gamma}{\rho}\right)^{1/4}\vee\left(\frac{\gamma}{\rho}\right)^{1/2}\cdot\sqrt{\sum_{i=1}^{n}\mathbb{E}_{x_{i}}\left[\tilde{x}_{i}^{\top}\tilde{x}_{i}\right]}$
$\displaystyle\leq\frac{1}{\sqrt{n}}\cdot\left(\frac{\gamma}{\rho}\right)^{1/4}\vee\left(\frac{\gamma}{\rho}\right)^{1/2}\cdot\operatorname{rank}\left(\Sigma_{X}\right)$

Based on this bound on Rademacher complexity, Theorem 5.1 can be proved by
directly applying the Theorem 8 from [7]. ∎

##  Appendix B Experimental details

In this section, we present experimental setup in detail. Specifically, we
describe the backbone model for each dataset in Sec. B.1, the detailed
datasets description in Sec. B.2, the implementation details in Sec. B.3,
uncertainty quantification results on simulated dataset in Sec. B.4, training
accuracy of different subpopulations throughout training process in Sec. B.5
and additional results in Sec. B.6.

###  B.1 Backbone model

Within each dataset, we keep the same model architecture as in previous work
[70]: ResNet-50 [23] for Waterbirds and CelebA, DistilBERT [16] for
CivilComments, and DenseNet-121 for Camelyon17. For ResNet-50, we used the
PyTorch [55] implementation pre-trained with ImageNet. For DistilBERT, we
employ the HuggingFace [67] implementation and start from the pre-trained
weights. Same as previous work [70], for DenseNet-121 we employ the
implementation without pretraining.

###  B.2 Datasets details

We describe the datasets used in the experiments in detail and summarize the
datasets in Table 4.

  * •

WaterBirds [58]. The task of this dataset is to distinguish whether the bird
is a waterbird or a landbird. According to the background and label of an
image, this dataset has four predefined subpopulations, i.e., “landbirds on
land”, “landbirds on water”, “waterbirds on land“ , and “waterbirds on water”.
In the training set, the largest subpopulation is “landbirds on land” with
3,498 samples, while the smallest subpopulation is “landbirds on water” with
only 56 samples.

  * •

CelebA [43]. CelebA is a well-known large-scale face dataset. Same as previous
works [58, 41], we employ this dataset to predict the color of the human hair
as “blond” or “not blond”. There are four predefined subpopulations based on
gender and hair color, i.e., “dark hair, female”, “dark hair, male”, “blond
hair, female” and “blond hair, male” with 71,629, 66,874, 22,880, and 1,387
training samples respectively.

  * •

CivilComments [9]. For this dataset, the task is to classify whether an online
comment is toxic or not, where according to the demographic identities (e.g.,
Female, Male, and White) and labels, 16 overlapping subpopulations can be
defined. We use 269,038, 45,180, and 133,782 samples as training, validation,
and test datasets respectively.

  * •

Camelyon17 [5, 33]. Camelyon17 is a pathological image dataset with over 450,
000 lymph-node scans used to distinguish whether there is cancer tissue in a
patch. The training data is drawn from three hospitals, while the validation
and test data are sampled from other hospitals. However, due to the different
coloring methods, even the same hospital samples have different distributions.
Therefore, we cannot get reliable subpopulation labels of Camelyon17.

Table 4: Summary of the datasets used in the experiments. Datasets | Labels | Groups | Population type | Data type | Backbone model
---|---|---|---|---|---
Waterbirds | 2 | 2 | Label×Group | Image | ResNet-50
CelebA | 2 | 2 | Label×Group | Image | ResNet-50
CivilComments | 2 | 8 | Label×Group | Text | DistilBERT-uncased
Camelyon17 | 2 | 5 | Group | Image | DenseNet-121

###  B.3 Implementation details

In this section, we present the implementation details of all approaches. We
implement our method in the codestack released with the WILDS datasets [33].
For some comparative methods, including ERM, IRM [3], IB-IRM [1], V-REx [34],
CORAL [63], Group DRO [58], DomainMix [69], Fish [60], LISA [70], vanilla
mixup and in-group mixup, we directly use the results in previous work [70].
For JTT [41], on the Waterbirds and CelebA datasets, we directly report the
results in the paper, and on the CivilComments dataset, due to a different
backbone model being employed, we reimplement the algorithm for fairly
comparison. Same as the proposed method, we reimplement other methods in the
codestack released with the WILDs datasets. We employ vanilla mixup on
WaterBirds and Camelyon17 datasets. On CelebA and CivilComments datasets, we
employ cutmix [71] and manifoldmix [65] respectively. For all approaches, we
tune all hyperparameters with AutoML toolkit NNI [49] based on validation
performance. Then we run the experiment multiple times on a computer with 8
Tesla V100 GPUs with different seeds to obtain the average performance and
standard deviation. The selected hyperparameters for Algorithm 1 and Algorithm
2 are listed in Tabel 5.

Table 5: Hyperparameter settings for Algorithm 1 and Algorithm 2.

| WaterBirds | CelebA | CivilComments | Camelyon17
---|---|---|---|---
Learning rate | 1e-5 | 1e-4 | 5e-5 | 1e-5
Weight decay | 1 | 1e-4 | 1e-4 | 1e-2
Batch size | 64 | 128 | 128 | 32
Optimizer | SGD | SGD | AdamW | SGD
Hyperparameter $\alpha$ | 0.5 | 1.5 | 0.5 | 0.5
Hyperparameter $\sigma$ | 0.5 | 0.5 | 1 | 1
Maximum Epoch | 300 | 20 | 10 | 5

(a) Hyperparameter settings for Algorithm 1.

| WaterBirds | CelebA | CivilComments | Camelyon17
---|---|---|---|---
Learning rate | 1e-5 | 1e-5 | 1e-05 | 1e-3
Weight decay | 1 | 1e-1 | 1e-2 | 1e-2
Batch size | 64 | 128 | 128 | 32
Optimizer | SGD | SGD | AdamW | SGD
Start epoch $T_{s}$ | 50 | 0 | 0 | 0
Sampling epoch $T$ | 50 | 5 | 5 | 5
Hyperparameter $\eta$ | 80 | 50 | 3 | 5
(b) Hyperparameter settings for Algorithm 2.

###  B.4 Uncertainty quantification results on simulated dataset

We conduct a toy experiment to show the uncertainty quantification could work
well on the dataset with subpopulation shift. Specifically, we construct a
four moons dataset (i.e., a dataset with four subpopulations) as shown in Fig.
2. We compare our approximation (i.e., Eq. 6) with the following ensemble-
based approximation:

$u_{i}\approx\frac{1}{T}\sum_{t=1}^{T}\kappa(y_{i},\hat{f}_{\theta_{t}}(x_{i}))p(\theta_{t};\mathcal{D})d\theta.$
(8)

Specifically, we train $T$ models and then ensemble them. The quantification
results are shown in Fig. 3. We can observe that (1) the proposed historical-
based uncertainty quantification method could work well on the simulated
dataset; (2) compared with the ensemble-based method, the proposed method
could better characterize the subpopulation shift.

Figure 2: Simulated dataset with four different subpopulations. In the four
subpopulations, Group 0 and Group 2 have the same label and groups 1 and 3
have the same labels.

(a) Ours

(b) Ensemble

Figure 3: Visualization of the obtained uncertainty with kernel density
estimation on simulated dataset, where group size refers to the sample number
of the group.

###  B.5 Training accuracy throughout training

We present how the training accuracy change throughout training in Fig. 4 on
the CelebA and Waterbirds datasets to empirically show why the proposed
estimation approach could work. From the experimental results, we observe that
during training, easy groups with sufficient samples can be fitted well, and
vice versa. For example, on the CelebA dataset, Group 0 and Group 1 with about
72K and 67K training samples quickly achieved over 95% accuracy. The accuracy
rate on Group 2, which has about 23K training samples, increased more slowly
and finally reached around 84%. The accuracy on Group 3, which has only about
1K training samples, is the lowest. Meanwhile, On the Waterbirds dataset, the
samples of hard-to-classify group (e.g., Group 1) are also more likely to be
forgotten by the neural networks.

(a) CelebA

(b) Waterbirds

Figure 4: Visualization of the changing of training accuracy on different
groups of CelebA and Waterbirds datasets.

###  B.6 Additional results

In this section, we present the full results with standard deviation in Table
6, Table 7, and Table 8.

Table 6: Full comparison results with other methods in the group-oblivious setting where NA indicates the standard deviation in the original paper [41] is not available. The best results are in bold blue.  | Waterbirds | CelebA
---|---|---
| Avg. | Worst | Avg. | Worst
ERM | 97.0 ± 0.2% | 63.7 ± 1.9% | 94.9 ± 0.2% | 47.8 ± 3.7%
Focal Loss [40] | 87.0 ± 0.5% | 73.1 ± 1.0% | 88.4 ± 0.3% | 72.1 ± 3.8%
CVaR-DRO [38] | 90.3 ± 1.2% | 77.2 ± 2.2% | 86.8 ± 0.7% | 76.9 ± 3.1%
CVaR-DORO [72] | 91.5 ± 0.7% | 77.0 ± 2.8% | 89.6 ± 0.4% | 75.6 ± 4.2%
$\chi^{2}$-DRO [38] | 88.8 ± 1.5% | 74.0 ± 1.8% | 87.7 ± 0.3% | 78.4 ± 3.4%
$\chi^{2}$-DORO [72] | 89.5 ± 2.0% | 76.0 ± 3.1% | 87.0 ± 0.6% | 75.6 ± 3.4%
JTT [41] | 93.6 ± (NA)% | 86.0 ± (NA)% | 88.0 ± (NA)% | 81.1 ± (NA)%
Ours | 93.0 ± 0.5% | 90.0 ± 1.1% | 90.1 ± 0.4% | 85.3 ± 4.1%
| CivilComments | Camelyon17
| Avg. | Worst | Avg.
ERM | 92.2 ± 0.1% | 56.0 ± 3.6% | 70.3 ± 6.4%
Focal Loss [40] | 91.2 ± 0.5% | 60.1 ± 0.7% | 68.1 ± 4.8%
CVaR-DRO [38] | 89.1 ± 0.4% | 62.3 ± 0.7% | 70.5 ± 5.1%
CVaR-DORO [72] | 90.0 ± 0.4% | 64.1 ± 1.4% | 67.3 ± 7.2%
$\chi^{2}$-DRO [38] | 89.4 ± 0.7% | 64.2 ± 1.3% | 68.0 ± 6.7%
$\chi^{2}$-DORO [72] | 90.1 ± 0.5% | 63.8 ± 0.8% | 68.0 ± 7.5%
JTT [41] | 90.7 ± 0.3% | 67.4 ± 0.5% | 69.1 ± 6.4%
Ours | 90.6 ± 0.4% | 70.1 ± 0.9% | 75.1 ± 5.9%
Table 7: Full comparison results with the algorithms using training group labels (Our method does not depend on this type of information). Results of baseline models are from [70]. The best three results are in bold brown or bold blue and the color indicates whether the train group label is used.  | Group labels | Waterbirds | CelebA
---|---|---|---
| in train set? | Avg. | Worst | Avg. | Worst
IRM | Yes | 87.5 ± 0.7% | 75.6 ± 3.1% | 94.0 ± 0.4% | 77.8 ± 3.9%
IB-IRM | Yes | 88.5 ± 0.6% | 76.5 ± 1.2% | 93.6 ± 0.3% | 85.0 ± 1.8%
V-REx | Yes | 88.0 ± 1.0% | 73.6 ± 0.2% | 92.2 ± 0.1% | 86.7 ± 1.0%
CORAL | Yes | 90.3 ± 1.1% | 79.8 ± 1.8% | 93.8 ± 0.3% | 76.9 ± 3.6%
GroupDRO | Yes | 91.8 ± 0.3% | 90.6 ± 1.1% | 92.1 ± 0.4% | 87.2 ± 1.6%
DomainMix | Yes | 76.4 ± 0.3% | 53.0 ± 1.3% | 93.4 ± 0.1% | 65.6 ± 1.7%
Fish | Yes | 85.6 ± 0.4% | 64.0 ± 0.3% | 93.1 ± 0.3% | 61.2 ± 2.5%
LISA | Yes | 91.8 ± 0.3% | 89.2 ± 0.6% | 92.4 ± 0.4% | 89.3 ± 1.1%
Ours | No | 93.0 ± 0.5% | 90.0 ± 1.1% | 90.1 ± 0.4% | 85.3 ± 4.1%
| Group labels | CivilComments | Camelyon17
| in train set? | Avg. | Worst | Avg.
IRM | Yes | 88.8 ± 0.7% | 66.3 ± 2.1% | 64.2 ± 8.1%
IB-IRM | Yes | 89.1 ± 0.3% | 65.3 ± 1.5% | 68.9 ± 6.1%
V-REx | Yes | 90.2 ± 0.3% | 64.9 ± 1.2% | 71.5 ± 8.3%
CORAL | Yes | 88.7 ± 0.5% | 65.6 ± 1.3% | 59.5 ± 7.7%
GroupDRO | Yes | 89.9 ± 0.5% | 70.0 ± 2.0% | 68.4 ± 7.3%
DomainMix | Yes | 90.9 ± 0.4% | 63.6 ± 2.5% | 69.7 ± 5.5%
Fish | Yes | 89.8 ± 0.4% | 71.1 ± 0.4% | 74.7 ± 7.1%
LISA | Yes | 89.2 ± 0.9% | 72.6 ± 0.1% | 77.1 ± 6.5%
Ours | No | 90.6 ± 0.5% | 70.1 ± 0.9% | 75.1 ± 5.9%
Table 8: Full comparison with ERM and mixup based methods. Results of baseline models are from [70]. The best results are in bold brown or bold blue and the color indicates whether the train group label is used. | Group labels | Waterbirds | CelebA
---|---|---|---
| in train set? | Avg. | Worst | Avg. | Worst
ERM | No | 97.0 ± 0.2% | 63.7 ± 1.9% | 94.9 ± 0.2% | 47.8 ± 3.7%
vanilla mixup | No | 81.0 ± 0.2% | 56.2 ± 0.2% | 95.8 ± 0.0% | 46.4 ± 0.5%
in-group mixup | Yes | 88.7 ± 0.3% | 68.0 ± 0.4% | 95.2 ± 0.3% | 58.3 ± 0.9%
Ours | No | 93.0 ± 0.5% | 90.0 ± 1.1% | 90.1 ± 0.4% | 85.3 ± 4.1%
| Group labels | CivilComments | Camelyon17
| in train set? | Avg. | Worst | Avg.
ERM | No | 92.2 ± 0.1% | 56.0 ± 3.6% | 70.3 ± 6.4%
vanilla mixup | No | 90.8 ± 0.8% | 67.2 ± 1.2% | 71.2 ± 5.3%
in-group mixup | Yes | 90.8 ± 0.6% | 69.2 ± 0.8% | 75.5 ± 6.7%
Ours | No | 90.6 ± 0.5% | 70.1 ± 0.9% | 75.1 ± 5.9%

##  Appendix C Justification for choosing historical-based uncertainty score

We employ the information from the historical training trajectory to
approximate the sampling process because it is simple and effective in
practice. Empirically, in contrast to other typical uncertainty quantification
methods such as Bayesian learning or model ensemble [17, 36], our method can
significantly reduce the computational and memory-storage cost by employing
the information from the historical training trajectory, since Bayesian
learning or model ensemble needs to sample/save multiple DNN models and
performs inference computations on them. Meanwhile, our method has achieved
quite promising final accuracy in contrast to other methods. In summary, we
choose an uncertainty score that can achieve satisfactory performance while
being more memory and computationally efficient.

##  Appendix D Societal impact and limitations

###  D.1 Societal impact

Algorithmic fairness and justice are closely related to our work.
Philosophically, there are two different views on justice. Firstly, Jeremy
Bentham believes “the greatest good for the greatest number” can be seen as
justice [50]. ERM can be considered to inherit this spirit which pays more
attention to minimizing the majority subpopulation risks. Different from
Jeremy Bentham’s opinion, Rawlsian distributive justice [57] argues that we
should maximize the welfare of the worst-off group. The proposed method and
other IW-based methods can be seen as the practice of Rawlsian distributive
justice due to focusing more on the minority subpopulations. However, in
practice, the proposed method and other IW-based methods may sacrifice the
average accuracy. Therefore, the ones using the proposed method need to
carefully consider what fairness and justice are in a social context to decide
whether to sacrifice the average accuracy and improve the worst-case accuracy.

###  D.2 Limitations and future works

Even though the proposed method achieves excellent performance, it still has
some potential limitations. (1) Similar to other IW-based methods, the
proposed method may sacrifice the average accuracy. Therefore, it is also
important and valuable to conduct a theoretical analysis of this phenomenon
and explore novel ways to improve the worst-case accuracy of the model without
sacrificing the average accuracy in the future work. (2) Although our method
does not require training set group labels, how to leverage unreliable
subpopulation information (e.g., subpopulation labels are noise) to improve
UMix would be a promising research topic. For example, when the unreliable
subpopulation labels are available, UMix could be improved by equipping with
existing importance weighting methods. (3) Similar to the previous IW-based
methods, the label noise is also not considered in our method, which may lead
to over-focusing on noisy samples. Currently, it’s still a challenging open
problem to distinguish the minority samples from the mislabeled noise samples
in the data with subpopulation shift. (4) At the same time, this work only
considers subpopulation shifts on Euclidean data, hence it is also a promising
future direction to generalize IW-based methods to graph-structured data,
under the guidance of invariance principle on graphs, such as that of [12]. We
leave them as important future works.