capbtabboxtable[][]

# Reweighted Mixup for Subpopulation Shift

Zongbo Han, Zhipeng Liang, Fan Yang, Liu Liu, Lanqing Li, Yatao Bian, Peilin
Zhao,
Qinghua Hu, Bingzhe Wu, Changqing Zhang, Jianhua Yao  Z. Han, C. Zhang and Q.
Hu are with the College of Intelligence and Computing, Tianjin University,
Tianjin 300072, China (e-mail: {zongbo, zhangchangqing,
huqinghua}@tju.edu.cn).
Z. Liang is with Department of Industrial Engineering and Decision Analytics,
Hong Kong University of Science and Technology, Hongkong 999077, China
(e-mail: zliangao@connect.ust.hk).
F. Yang, Liu Liu, B. Wu, J. Yao, La. Li, Y. Bian and P. Zhao are with Tencent
AI Lab, Shenzhen 518054, China (e-mail: {fionafyang, leonliuliu, bingzhewu,
<EMAIL_ADDRESS>{lanqingli1993<EMAIL_ADDRESS>and
peilinzhao@hotmail.com).
This work is partially done by Zongbo Han and Zhipeng Liang during their
internship at Tencent AI Lab. Corresponding authors: Bingzhe Wu, Yao Jianhua
and Changqing Zhang.

###### Abstract

Subpopulation shift exists widely in many real-world applications, which
refers to the training and test distributions that contain the same
subpopulation groups but with different subpopulation proportions. Ignoring
subpopulation shifts may lead to significant performance degradation and
fairness concerns. Importance reweighting is a classical and effective way to
handle the subpopulation shift. However, recent studies have recognized that
most of these approaches fail to improve the performance especially when
applied to over-parameterized neural networks which are capable of fitting any
training samples. In this work, we propose a simple yet practical framework,
called reweighted mixup (RMix), to mitigate the overfitting issue in over-
parameterized models by conducting importance weighting on the “mixed”
samples. Benefiting from leveraging reweighting in mixup, RMix allows the
model to explore the vicinal space of minority samples more, thereby obtaining
more robust model against subpopulation shift. When the subpopulation
memberships are unknown, the training-trajectories-based uncertainty
estimation is equipped in the proposed RMix to flexibly characterize the
subpopulation distribution. We also provide insightful theoretical analysis to
verify that RMix achieves better generalization bounds over prior works.
Further, we conduct extensive empirical studies across a wide range of tasks
to validate the effectiveness of the proposed method.

###### Index Terms:

Importance Reweighting, Mixup, Subpopulation Shift, Uncertainty.

##  1 Introduction

Distribution shift is a common challenge for empirical risk minimization
(ERM), which refers to that the data used for training and test have different
distributions [1, 2, 3]. A specific form of distribution shift is
subpopulation shift wherein the training and test distributions consist of the
same subpopulation groups but differ in subpopulation frequencies [4, 5]. Many
practical research problems (e.g., fairness of machine learning and class
imbalance) can all be considered as a special case of subpopulation shift [6,
7, 8, 9]. For example, when we conduct training on a dataset with biased
demographic subpopulations, the obtained model may be unfair to the minority
subpopulations [7, 6, 8]. Therefore the essential goal of fair machine
learning is to mitigate the performance gap between different subpopulations
in the test datasets.

Plenty of approaches have been proposed for solving the subpopulation shift
problem. Among these approaches, importance weighting (IW) is a classical yet
effective technique by imposing static or dynamic weight 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) [1, 10, 11, 12, 13, 14], such as class-imbalanced
learning approaches [11, 13, 14]. Alternatively, a more flexible way is to
assign sample-wise importance weights to individual samples adaptively
according to training dynamics [15, 16, 17, 18, 19, 20, 21, 22].
Distributional robust optimization (DRO) is one of the most representative
methods in this line, which minimizes loss over the worst-case distribution in
a neighborhood of the empirical training distribution. A commonly used dual
form of DRO can be regarded as a special case of importance reweighting
wherein the importance weights of training samples are updated based on the
current loss [23, 24, 25, 26] 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 [27, 12] recognize
that various IW methods tend to exacerbate overfitting, which leads to a
diminishing effect on stochastic gradient descent over training epochs
especially when they are applied to over-parameterized neural networks.
Theoretically, previous studies prove that for over-parameterized neural
networks, reweighting algorithms do not improve over ERM because their
implicit biases are (almost) equivalent [28, 10, 29]. In addition, prior work
also points out that using conventional regularization techniques such as
weight decay cannot significantly improve the performance of IW [12].
Therefore, in this paper, we propose to use more exploratory importance
reweighting equipped with mixup strategies to eliminate model overfitting.

To this end, we introduce a novel technique called reweighted mixup (RMix), by
reweighting the mixed samples according to the subpopulation memberships
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
the subpopulation memberships of the original samples. When the subpopulation
memberships are unknown, uncertainty-based importance weights are proposed by
assigning importance weights according to the uncertainty of training samples
to better explore the minority subpopulation samples. The reweighted 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 importance-aware manner, i.e., putting more focus on samples that belong
to minority subpopulations when constructing “mixed” samples and corresponding
losses. We also show RMix can provide an additional theoretical benefit that
achieves a tighter generalization bound than weighted ERM [30, 21, 18, 25].

Compared with the conference version [31], we significantly improve our
previous work in the following aspects: (1) further reweighting analysis of
the proposed RMix by providing a more intuitive comparison with previous
mixup-based methods (Sec. 4.2); (2) extension of the proposed method in the
group-aware setting which enables the model to use subpopulation memberships
information (Sec. 4.3.1); (3) more theoretical extension to demonstrate that
the proposed method is compatible with both vanilla mixup [32] and CutMix-
based mixup[33], and can achieve a tighter generalization error bound (Sec.
5); (4) more experiments and discussions to verify the effectiveness of the
proposed method in both group-aware and group-oblivious settings (Sec. 6).
Overall, the proposed RMix supported by theoretical guarantees significantly
improves previous methods against subpopulation shift. The contributions of
this paper are summarized as follows:

  * (1)

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

  * (2)

Under the proposed framework, we provide two strategies to assign the
importance weights making the proposed method work well in both group-
oblivious and group-aware settings. Especially for group-oblivious setting, we
propose an uncertainty-based importance weighting strategy, which could
accurately characterize the subpopulation in training.

  * (3)

We provide an insightful theoretical analysis that RMix can achieve a tighter
generalization bound than previous importance weighting methods. Moreover, a
qualitative analysis with strong intuition is provided to illustrate the
reweighting effect of the proposed method over the previous mixup.

  * (4)

We conduct extensive experiments on a wide range of tasks, and the proposed
RMix achieves state-of-the-art performance in both group-oblivious and group-
aware settings.

##  2 Related Work

###  2.1 Importance weighting for subpopulation shift

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 [1, 10, 11, 12, 13, 14]. 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 [1, 10, 12]. Besides, there are
some methods to obtain sample weights based on the effective number of samples
[11], subpopulation margins [13], and Bayesian networks [14].

Dynamic 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 [15, 16, 17, 18, 19, 20, 21, 22]. 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 importance
weight is updated based on the current loss [23, 24, 25, 26] in an alternated
manner. For example, in the group-aware setting (i.e., we know each sample
belongs to which subpopulation), GroupDRO [12] introduces an online
optimization algorithm to update the weights of each group. In the group-
oblivious setting, existing methods mostly [15, 19, 17, 20] 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 Mixup-based methods

Mixup training [32] has been both theoretically [34, 35] and experimentally
[32, 33, 36] shown that it could improve the generalization and robustness of
the model. Specifically, mixup could extend the training distribution by
conducting a convex combination between any pair of training samples, which is
equivalent to minimizing the vicinal risk of the training distribution [37].

Many mixup variants are proposed to improve the vanilla mixup by changing how
the training samples are combined. Typically, CutMix [33] cuts and pastes the
original training image to avoid the blurring and noise of the generated
images caused by vanilla mixup using linear interpolation. Puzzlemix [38]
further improves CutMix by selecting cut and pasted image patches via saliency
information, thus preserving key information of samples before mixing.
Adversarial mixup resynthesis [39] introduces adversarial learning to train a
mixing function that could generate more realistic mixed samples. Manifold
mixup [36] obtains neural networks with smoother decision boundaries by mixing
samples in representation space. To explore the representation space more,
patch up [40] mixes and exchanges blocks of any training samples in the
feature map space generated by the convolutional neural network.

Mixup has also been extensively studied to improve the generalization of the
model [41, 42, 43, 44, 45]. Typically, existing methods [45, 43, 45] try to
mix the samples to synthesize novel domains, thereby improving the
generalization. Besides, by selecting mixed sample pairs during training, LISA
[42] and cmix [44] improve the generalization of the model in classification
and regression respectively. Different from the above methods, we do not focus
on how to mixup the training samples, but improve the subpopulation shift
robustness of the model by re-weighting the loss of the mixed samples.
Besides, our work is orthogonal to the previous methods, i.e., we can use our
weight-building strategy to improve the performance of previous mixup-based
methods.

###  2.3 Uncertainty quantification

When the subpopulation memberships are unknown, the core of our method is
based on the high-quality uncertainty quantification of each sample. There are
many approaches proposed to quantify the uncertainty. Specifically, the
uncertainty of deep learning models includes epistemic (model) uncertainty and
aleatoric (data) uncertainty [46]. To obtain the epistemic uncertainty,
Bayesian neural networks (BNNs) [47, 48, 49, 46] have been proposed which
replace the deterministic weight parameters of the model with the
distribution. Unlike BNNs, ensemble-based methods obtain the epistemic
uncertainty by training multiple models and ensembling them [50, 51, 52, 53].
Aleatoric uncertainty focuses on the inherent noise in the data, which usually
is learned as a function of the data [46, 54, 55]. Uncertainty quantification
has been successfully equipped in many fields such as multimodal learning [56,
57, 58], multitask learning [59, 60], and reinforcement learning [61, 62].
Unlike previous methods, our method focuses on estimating the epistemic
uncertainty of training samples with subpopulation shifts and then reweighting
uncertain samples, thereby improving the performance of minority
subpopulations with high uncertainty.

##  3 Preliminary

In this section, the necessary background and notations are clarified. Let the
input and label space be $\mathcal{X}$ and $\mathcal{Y}$ respectively. Given a
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 (also known as groups), 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
($\sum_{g=1}^{G}k_{g}=1$). 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 following expected loss

${\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\mathbb{E}_{(x,y)\sim
P}[\ell(\theta,x,y)],}$ (1)

where $\ell:\Theta\times\mathcal{X}\times\mathcal{Y}\rightarrow\mathbb{R}_{+}$
denotes the loss function. For training data contains multiple subpopulations,
the expected loss of ERM can be written as

$\sum_{g=1}^{G}k_{g}\mathbb{E}_{(x,y)\sim P_{g}}[\ell(\theta,x,y)].$ (2)

This may lead to the model paying more attention to the majority
subpopulations (i.e., subpopulations with a larger proportion) in the training
set and resulting in poor performance on the minority subpopulations (i.e.,
subpopulations with a smaller proportion). For example, the ERM-based models
may learn spurious correlation between the majority subpopulations and labels
but this spurious correlation does not hold in the minority
subpopulations[12].

In this paper, we focus on learning a model that is robust against
subpopulation shift by importance reweighting, i.e., we can still achieve
better performance even on the worst-case subpopulation among all the
subpopulations. Specifically, importance reweighting tries to assign an
importance weight to each sample, therefore each subpopulation could
contribute comparably to the final loss, e.g.,

$\mathbb{E}_{(x,y)\sim P}[w(x,y)\ell(\theta,x,y)],$ (3)

where the function $w(x,y)$ is the weighted function to return the weight of
the sample $(x,y)$. This article hopes to improve the importance reweighting
algorithm by equipping it with mixup.

Setting of subpopulation shift. Previous works on improving subpopulation
shift robustness investigate two different settings, i.e., group-aware and
group-oblivious [18, 30, 12]. Most of the previous works assume that the group
labels are available during training [12, 42], i.e., which subpopulation the
training sample belongs to is known. This is called group-aware setting.
However, we may not have training group labels in many real applications [8].
Meanwhile, the group label information may not be available due to privacy
concerns [19, 8]. This paper studies both group-aware and group-oblivious
settings under a unified reweighting mixup framework. On the group-aware
setting, we set weights for different subpopulations according to the group
information. Besides, on the group-oblivious setting, we 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.

##  4 Method

This section presents the technical details of RMix. The proposed method can
promote the prediction fairness on all subpopulations by reweighting the mixed
samples. Intuitively, by reweighting the mixed samples instead of the original
training samples, the deep neural networks especially overparameterized neural
networks can always encounter “novel” samples, which avoids the model from
memorizing the minority subpopulations rather than learning the intrinsic
correlation between inputs and labels. From another perspective, RMix assigns
greater weights to the minority subpopulation samples when creating mixed
samples, so the model can better explore the vicinity space of the minority
samples, thereby making the model perform well on all subpopulations
uniformly. Last but not least, RMix is a general reweighting framework that
can work well in both group-oblivious or group-aware settings. When the group
label is known, we can simply assign a negative-correlated importance weight
to each group based on their sample size. When the group label is unknown, we
use uncertainty information to assign importance weights to training samples.

We first introduce the basic procedure of RMix in Sec. 4.1, and then provide a
deeper analysis of RMix to intuitively show the effects of importance weights
in Sec. 4.2. Finally, we present how RMix assigns importance weights in both
group-aware and group-oblivious settings in Sec. 4.3.

###  4.1 Importance-weighted mixup

It has been both theoretically and experimentally shown that the effect of
importance reweighting could fail to achieve better worst-case subpopulation
performance especially when over-parameterized neural networks are employed
[27, 12, 28, 10, 29]. Specifically, deep neural networks tend to memorize
minority samples with higher weights in the training set rather than learn the
correlation between minority samples and labels, resulting in poor
generalization on minority samples at test time. To this end, RMix employs an
aggressive data augmentation strategy called importance reweighted mixup to
mitigate overfitting, which enables the model to see “new” samples during the
training by reweighting the mixed samples, avoiding the model simply
memorizing minority samples instead of learning for the intrinsical
correlation between inputs and labels.

Formally, vanilla mixup [32] constructs virtual training examples (i.e., mixed
samples) by performing linear interpolations between data/features and
corresponding labels

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

where $(x_{i},y_{i}),{(x_{j},y_{j})}$ are two samples drawn randomly from
empirical training distribution and $\lambda\in[0,1]$ is usually sampled from
a beta distribution $Beta(\alpha,\alpha)$. In computer vision tasks, another
popular mixup variant is CutMix [33], which constructs virtual training
examples by performing cutting and pasting in the original samples, i.e.,

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

where $M(\lambda)$ is a binary mask randomly chosen covering $\lambda$
proportion of the input, and $\odot$ represents the element-wise product. Note
that the underlying assumption of vanilla mixup and CutMix is the mixture of
original inputs which leads to the same linear interpolation of the
corresponding labels.

After constructing the virtual training examples, mixup optimizes the
following loss function

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

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

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

Eq. 7 can be regarded as a linear combination (mixup) of
$\ell(\theta,\widetilde{x}_{i,j},y_{i})$ and
$\ell(\theta,\widetilde{x}_{i,j},y_{j})$[32]. Unfortunately, since the
previous mixup-based methods do not consider the subpopulations with poor
performance, it has been shown experimentally to be non-robust against
subpopulation shift [42]. To this end, we introduce a simple yet effective
method called RMix, which further employs a reweighted linear combination of
the original loss based on Eq. 7 to encourage the learned model to pay more
attention to subpopulations with poor performance.

In contrast to the previous importance reweighting paradigm in Eq. 3, the
importance weights of RMix are used on the mixed samples. For the $i$-th
sample $x_{i}$, we denote its importance weight as $w_{i}$. Once we obtain the
importance weight of each sample, we can perform weighted linear combination
of $\ell(\theta,\widetilde{x}_{i,j},y_{i})$ and
$\ell(\theta,\widetilde{x}_{i,j},y_{j})$ with the following rule

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

where ${w}_{i}$ and ${w}_{j}$ denote the importance weight of the $i$-th and
$j$-th samples respectively. $\widetilde{x}_{i,j}$ can be obtained with
vanilla mixup (i.e., Eq. 4), CutMix (i.e., Eq. 5) or other mixup variants.
Therefore, the proposed method can cooperate with various mixup variants to
adapt to different downstream tasks. In practice, to balance the RMix and
normal training, we set a hyperparameter $\sigma$ that denotes the probability
to apply RMix. The overall training pseudocode for RMix 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 RMix, and parameter $\alpha$ of the beta distribution;

for  _each iteration_ do

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

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

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

Obtain the mixed input $\widetilde{x}_{i,j}$ with Eq. 4 or Eq. 5;

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

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

Algorithm 1 The training pseudocode of RMix.

###  4.2 Rethinking reweighting effects in RMix

In this section, the improvement of the proposed RMix over mixup is analyzed
by further decomposing Eq. 8, which can provide a stronger intuition and
theoretical analysis of why the proposed method works. Compared with the
mixup, the proposed RMix could achieve better worst-case subpopulation
performance by exploring the vicinal space of the minority samples more.

Formally, to analyze the effect of the importance reweighting, we convert Eq.
8 to the form of original mixup loss (i.e., Eq. 6). Then we have Proposition
4.1, which provides an intuitive comparison of RMix with original mixup loss.

######  Proposition 4.1.

When the cross-entropy loss is employed, the loss of RMix (i.e., Eq. 8) can be
rewritten as

$\displaystyle\mathbb{E}[\overline{w}_{i,j}\ell(\theta,\widetilde{x}_{i,j},\overline{y}_{i,j})],$
(9)

where

$\displaystyle\overline{w}_{i,j}=w_{i}\lambda+w_{j}(1-\lambda)\text{\ \ \
and},$
$\displaystyle\overline{y}_{i,j}=\frac{w_{i}}{\overline{w}_{i,j}}\lambda
y_{i}+\frac{w_{j}}{\overline{w}_{i,j}}(1-\lambda)y_{j}.$ (10)

$\overline{w}_{i,j}$ can be seen as a constant related with constants
$\lambda,w_{i}$, and $w_{j}$.

###### Proof.

(Proof of Proposition 4.1). For $K$-class classification, given a well trained
neural classifier $f_{\theta}:\mathcal{X}\rightarrow\mathcal{Y}$ that could
produce the predicted probability of each class
$f_{\theta}(x)=[f_{\theta}^{1}(x),\cdots,f_{\theta}^{K}(x)]$ and
$y=[y^{1},\cdots,y^{K}]$ denotes the one-hot vector form of label $y$. Then
when the cross-entropy loss is employed, Eq. 8 can be written as

$\displaystyle\mathbb{E}\\{-\sum_{k=1}^{K}[w_{i}\lambda
y_{i}^{k}+w_{j}(1-\lambda)y_{j}^{k}]\log
f_{\theta}^{k}(\widetilde{x}_{i,j})\\}$ $\displaystyle=$
$\displaystyle\mathbb{E}\\{-\overline{w}_{i,j}\sum_{k=1}^{K}[\frac{w_{i}\lambda}{\overline{w}_{i,j}}y_{i}^{k}+\frac{w_{j}(1-\lambda)}{\overline{w}_{i,j}}y_{j}^{k}]\log
f_{\theta}^{k}(\widetilde{x}_{i,j})\\}$ $\displaystyle=$
$\displaystyle\mathbb{E}[-\overline{w}_{i,j}\sum_{k=1}^{K}\overline{y}_{i,j}\log
f_{\theta}^{k}(\widetilde{x}_{i,j})]=\mathbb{E}[\overline{w}_{i,j}\ell(\theta,\widetilde{x}_{i,j},\overline{y}_{i,j})].$

∎

TABLE I: The improvements of RMix over mixup. | mixup | RMix
---|---|---
$y$ | $\lambda y_{i}+(1-\lambda)y_{j}$ | ${\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\frac{w_{i}}{\overline{w}_{i,j}}}\lambda y_{i}+{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\frac{w_{j}}{\overline{w}_{i,j}}}(1-\lambda)y_{j}$
loss | $\ell(\theta,\widetilde{x}_{i,j},\widetilde{y}_{i,j})$ | ${\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\overline{w}_{i,j}}\ell(\theta,\widetilde{x}_{i,j},\overline{y}_{i,j})$

Now based on the Proposition 4.1, we give a further comparison between RMix
and the previous mixup where the results are shown in Tab. I. Compared with
previous mixup, the effect of importance reweighting in RMix is reflected in
two aspects:

(1) Reweighted label interpolation. RMix introduces importance weights when
constructing the labels of virtual samples, which allows more attention on the
labels of minority samples. Specifically, given two training samples
$(x_{i},y_{i}),(x_{j},y_{j})$, unlike previous mixup, RMix conducts importance
weighted label interpolation (i.e., Eq. 10). Then, when samples from the
majority and minority subpopulations are mixed to construct virtual training
examples, the labels of the minority subpopulations (usually with larger
importance weights) could dominate the label mixing. Consequently, the labels
of the minority subpopulations will be the main components of the mixed
labels.

(2) Reweighted mixed loss. Different from the loss of mixup, RMix pays more
attention to the virtual samples mixed by minority subpopulations.
Specifically, in the mixup loss, all mixed samples have the same importance
weight while the samples from the minority subpopulations could have larger
weights in the RMix loss. Thus the minority subpopulations contribute more to
the loss function by increasing their weights.

All the above reweightings are conducted adpatively within the proposed
framework by setting the weights in Eq. 8 and we will elaborate it in the
following section.

###  4.3 Obtaining importance weights

In this section, we present how to obtain the importance weights for training
samples in Eq. 8 in the group-aware and group-oblivious settings, especially
the group-oblivious setting which is more challenging due to the lack of group
information for training samples. Specifically, we first introduce how to
obtain the importance weights of each training sample according to their group
information in the group-aware setting. Then, we mainly focus on how to obtain
reasonable importance weights in the group-oblivious setting. In general, we
propose an uncertainty estimation method to find samples with high uncertainty
and then set importance weights according to the uncertainty of training
samples.

####  4.3.1 Group-aware importance weights

In the group-aware setting, the key to improving the performance of the
minority subpopulations is to reweight the training samples according to the
group information of different samples, which means that we usually need to
assign larger weights to minority subpopulations to promote their learning
while giving relatively small weights to majority subpopulations. Since the
proposed RMix is a general framework that can cooperate with any reweighting
methods, to improve the efficiency and better verify the effectiveness of
RMix, we employ a simple and effective way to determine the importance weights
of different subpopulations.

In general, we leverage the sample size of each subpopulation to assign
importance weights, where the sample size is closely related to the learning
difficulty of different subpopulations. Specifically, similar to previous
works [63, 64, 65], instead of directly setting the weight of each sample to
be inversely proportional to the sample size of each subpopulation, we
empirically set the importance weight $w_{i}$ of each sample as a function of
the square root of its training group size $\sqrt{k_{g}N}$, where $k_{g}$ and
$N$ denote the $g$-th subpopulation’s proportion and training samples
respectively. Formally, for the $i$-th sample belonging to subpopulation $g$,
the corresponding importance weight $w_{i}$ is set as
$w_{i}:=\exp(C/\sqrt{k_{g}N})$, where $C\in\mathbb{R}_{+}$ is model capacity
constant and can be seen as a hyperparameter. In fact, the scaling
$\sqrt{k_{g}N}$ of different subpopulations could reflect how minority
subpopulations are more prone to overfitting than majority groups due to the
general size dependence of model-complexity-based generalization bounds [12,
13], which can be intuitively understood as how easily a subpopulation can be
learned by the neural network, i.e., the more samples in the subpopulation,
the lower the learning difficulty. Therefore, setting the importance weights
as a function of $\sqrt{k_{g}N}$ could better balance the contribution of
different subpopulations during training.

####  4.3.2 Group-oblivious importance weights

In the group-oblivious setting, finding samples with high uncertainty during
training is the key to assigning importance weights. For instance, in DRO-
based algorithms, the uncertainty set is constructed with the current loss
[23, 24, 25, 26]. However, the uncertain samples found in this way may
constantly change during training, because the loss is affected by the
reweighting effect and cannot accurately identify minority samples[30],
resulting in these methods not always upweighting the minority subpopulations.
To this end, a sampling-based stable uncertainty estimation method is
introduced to better characterize the subpopulation information of each
sample.

Formally, given a well-trained neural classifier
$f_{\theta}:\mathcal{X}\rightarrow\mathcal{Y}$ producing the predicted class
$\hat{f}_{\theta}(x)$, a simple way to obtain the uncertainty $u$ of a sample
$x$ is whether the sample is correctly classified. However, as pointed out in
previous work [50], 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 $u_{i}$ as

$u_{i}=\int\kappa(y_{i},\hat{f_{\theta}}(x_{i}))p(\theta;\mathcal{D})d\theta,$
(11)

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}.$ (12)

Then, we can obtain an approximation of Eq. 11 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}))$.

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
[50] or extra approximation errors need to be introduced [66]. We propose to
employ the information from the historical training trajectory to approximate
the sampling process. More specifically, we train a model with simple ERM and
record the prediction result $\hat{f}_{\theta_{t}}(x_{i})$ of each sample on
each iteration epoch $t$. Then, to reduce the influence of inaccurate
prediction at the beginning of training, we estimate uncertainty with
prediction after training $T_{s}-1$ epochs

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

We empirically show that the proposed approximation could obtain more reliable
uncertainty for data with subpopulation shift in Sec. A.1 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 and positively related to the
corresponding uncertainty with the following rule

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

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
and promising performance is achieved. The whole process for obtaining
training importance weights is summarized 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}]$;

for  _each iteration_ do

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

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

Obtain 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. (1) Epistemic and aleatoric uncertainty. The overall uncertainty could
be divided into epistemic and aleatoric uncertainty [46]. In our method,
samples are weighted only based on epistemic uncertainty by sampling from the
model on the training trajectory, which can be regarded as sampling from the
model posterior in a more efficient way. Note that, we do not consider
inherent noise (aleatoric uncertainty) since it is usually intractable to
distinguish between noisy samples and minority samples from data under
subpopulation shifts. (2) 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 [67]. Meanwhile,
the hard-to-classify samples are also more likely to be forgotten by the
neural networks [68]. The frequency of which samples are correctly classified
during training can be used in confidence calibration [69] and ensemble [53].
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 uncertainty. On the
other hand, samples from the majority subpopulations will have lower
uncertainty due to being classified correctly more frequently. In Sec. A.2 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. A.3
of the Appendix.

##  5 Theory

In this section, we provide a theoretical understanding of the generalization
ability for RMix in a unified framework that the RMix can achieve a better
generalization error bound than traditional IW methods without using vanilla
mixup or CutMix-based methods. For simplicity, our analysis focuses on
generalized linear model (GLM). The roadmap of our analysis is to first
formalize the vanilla mixup and CutMix in a unified framework and then
approximate the mixup loss. Finally we 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}$
$\displaystyle[w(x,y)\ell(\theta,x,y)]$
$\displaystyle-\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 reformalize the vanilla mixup and CutMix in a unified
framework briefly and more formal definitions are given in the appendix,

$\displaystyle\tilde{x}_{i,j}$ $\displaystyle=M(\lambda)\odot
x_{i}+(1-M(\lambda))\odot x_{j},$ (15) $\displaystyle\tilde{y}_{i,j}$
$\displaystyle=\lambda y_{i}+(1-\lambda)y_{j},$

where $\lambda$ is the ratio parameter between samples drawn from
$\mathcal{D}_{\lambda}$. $\odot$ means a component-wise multiplication in
vector or matrix. $M(\lambda)$ is a random variable conditioned on $\lambda$
that indicates how we mix the input, e.g., by linear interpolation [32] or by
a pixel mask [33]. For simplicity, we denote $M(\lambda)$ by $M$ below. Under
this template, two most popular mixup methods, vanilla mixup [32] and CutMix
[33], for $i$-th and $j$-th samples with $\lambda$ drawn from
$\mathcal{D}_{\lambda}$, can be rewritten as follow

$\displaystyle\tilde{x}_{i,j}^{\text{(mixup)}}=\lambda
x_{i}+(1-\lambda)x_{j},$ (16)
$\displaystyle\tilde{x}_{i,j}^{\text{(CutMix)}}=M\odot
x_{i}+\left(1-M\right)\odot x_{j}.$ (17)

i.e., the vanilla mixup is a special case of Equation 15 by putting
$M=\lambda\overrightarrow{1}$ and the CutMix is also a special case where $M$
is a binary mask that indicates the location of the cropped box region with a
relative area.

Now 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 exist constants
$L,B>0$ such that for any $\theta$ satisfying $\theta\in\mathcal{W}_{\gamma}$,
the following holds with a probability of at least $1-\delta$,

$\displaystyle\operatorname{GError}(\theta)$ $\displaystyle\leq 2L\cdot
L_{A}(\frac{1}{\sqrt{n}}(\gamma/\rho)^{1/4}(\sqrt{\operatorname{tr}((\Sigma_{X}^{(M)})^{\dagger}\Sigma_{X})}+\operatorname{rank}\left(\Sigma_{X}\right)))$
$\displaystyle\quad+B\sqrt{\frac{\log(2/\delta)}{2n}},$

where $\gamma$ is a constant,
$\Sigma_{X}=\sum_{g=1}^{G}k_{g}w_{g}\Sigma_{X}^{g}$,
$\Sigma_{X}^{(M)}=\mathbb{E}[(1-M)\Sigma_{X}(1-M)^{\top}]$, and $\rho$ is some
constant related to the data distribution, which will be formally introduced
in Assumption 5.1. Moreover,
$\mathcal{W}_{\gamma}=\\{\theta:\mathbb{E}_{x}A^{{}^{\prime\prime}}(\theta^{\top}x)\cdot\theta^{\top}(\Sigma_{X}^{(M)}+x\operatorname{Var}(M)x^{\top})\theta\leq\gamma\\}$.

We will show later that the output of our Algorithm 1 would yield the solution
$\theta$ falls into the set $\mathcal{W}_{\gamma}$ and thus Theorem 5.1 can
provide a theoretical understanding of our algorithm. In contrast to weighted
ERM, the bound improvement of RMix is on the term,
$\sqrt{\operatorname{tr}((\Sigma_{X}^{(M)})^{\dagger}\Sigma_{X})}+\operatorname{rank}\left(\Sigma_{X}\right)$,
which can partially reflect the heterogeneity of the training subpopulations.
Specifically, the term would become $\sqrt{d}$ 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.,
$\sqrt{\operatorname{tr}((\Sigma_{X}^{(M)})^{\dagger}\Sigma_{X})}+\operatorname{rank}\left(\Sigma_{X}\right)\ll
d$). Such an improvement can be achieved by tuning proper weights $w_{g}$ for
different subpopulations.

The proof of Theorem 5.1 follows this roadmap: (1) We first show that the
model learned with RMix 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 B.3). (3) Finally,
we can characterize the generalization bound by using complexity-based
learning theory [70] (Theorem 8). More details of the proof can be found in
Appendix.

As we discuss in the Appendix, the weighted mixup can be seen as an
approximation of a regularization term
$\frac{C}{2n}[\sum_{i=1}^{n}w_{i}A^{\prime\prime}(x_{i}^{\top}\theta)]\theta^{\top}(\mathbb{E}(1-M)\widehat{\Sigma}_{X}(1-M)^{\top}+x_{i}\operatorname{Var}(M)x_{i}^{\top})\theta$
for some constant $C$, compared with the non-mixup algorithm, which motivates
us to study the following hypothesis space $\mathcal{W}_{\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., [71, 34].

######  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\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 B.3 and the
Theorem 8 of [70].

######  Lemma 5.2.

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

$\displaystyle Rad(\mathcal{W}_{r})\leq$
$\displaystyle\frac{1}{\sqrt{n}}(\gamma/\rho)^{1/4}\left(\sqrt{\operatorname{tr}\left(\left(\Sigma_{X}^{(M)}\right)^{\dagger}\Sigma_{X}\right)}+\operatorname{rank}\left(\Sigma_{X}\right)\right).$

TABLE II: We conduct experiments on diverse tasks and here we list the summary of the datasets used in the experiments. For each dataset, we present the number of classes, the predefined subpopulation’s number, population type, input data type, the number of dataset instances, and the lowest minority subpopulation’s proportion in turn. As can be seen from the table, the minimal minority subpopulation proportion in the training set is only about 2% at most, which could significantly reduce the performance of minority subpopulations. Datasets | #Classes | #Subpopulations | Population type | Data type | #Size | Minimal $k_{g}$
---|---|---|---|---|---|---
Waterbirds | 2 | 4 | Background×Label | Synthetic Image | 11,788 | 1.17%
CelebA | 2 | 4 | Gender×Label | Face Image | 202,599 | 0.85%
CivilComments | 2 | 16 | Demographic×Label | Online Comment | 448,000 | 0.37%
FMoW | 62 | 5 | Geographical Regions | Satellite Image | 141,696 | 2.06%
Camelyon17 | 2 | 30 | Tissue Slides | Pathological Image | 455,954 | 0.47%

##  6 Experiments

In this section, we conduct multiple experiments to comprehensively evaluate
the proposed algorithm. We first introduce the experimental details. Then we
present the experimental results and discussions in the group-aware and group-
oblivious settings respectively. Finally, we conduct an ablation study to show
the key factor of performance improvement in our method.

###  6.1 Experimental Setup

The experimental setup is presented in detail in this section. Specifically,
we present the used datasets details, evaluation metrics, model selection
strategy, comparison methods, and other implementation details.

Datasets. We perform experiments on multiple benchmark datasets with
subpopulation shifts, including Waterbirds [12], CelebA [72], CivilComments
[73], Camelyon17[74, 75], and FMoW [76, 75]. We summarize the datasets in
Table II and the details of these datasets are as follows:

  * •

WaterBirds. 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 lowest subpopulation is “landbirds on water” with only 56
samples.

  * •

CelebA. CelebA is a well-known large-scale face dataset. Same as previous
works [12, 30], 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. For this dataset, the task is to classify whether an online
comment is toxic or not. 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. In the group-aware setting, if one sample
belongs to multiple subpopulations, we use the least subpopulation (with the
smallest training number) as the group label for that example to obtain the
importance weights.

  * •

FMoW. FMoW is a satellite image dataset and the task is to distinguish the
building or land use categories of each image. The training data was taken in
five different geographical regions (Africa, the Americas, Oceania, Asia, and
Europe). Since the distribution of satellite images in different geographic
regions is different due to different development statuses and the training
number of different geographical regions is different, we directly use the
geographical region information as the subpopulation label.

  * •

Camelyon17. 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 30 tissue slides from three different
hospitals, while the validation and test data are sampled from another 10
tissue slides from other hospitals respectively. Due to the different coloring
methods, even the same hospital samples have different distributions.
Therefore, we use the tissue slides information as the subpopulation label in
the group-aware setting.

Evaluation metrics. To be same as previous works [42, 63] and WILDs
leaderboard[75], we report the average accuracy over 10 different random seeds
on the Camelyon17 dataset. On other datasets, we repeat experiments over 3
times and report the average and worst-case accuracy among all subpopulations.
Besides, to better show the performance gap between different subpopulations,
we also present the gap between the average accuracy and worst-case accuracy.
Note that, the trade-off between the average and worst-case accuracy is a
well-known challenge [8, 77] and in this paper we emphasize worst-case
accuracy, which is more important than the 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.

TABLE III: Comparison results in the group aware setting with current state-of-the-art algorithms. For better presentation, the best results are in bold and blue. Meanwhile, the second and third results are in bold and brown. _Note that the proposed method focuses on improving the model’s worst-case accuracy to improve model subpopulation shift robustness rather than the average accuracy._ The proposed method consistently outperforms all comparison methods in worst-case accuracy on all datasets. Datasets/ | Waterbirds | CelebA
---|---|---
ACC(%) |  Avg.$(\uparrow)$ |  Worst$(\uparrow)$ |  Gap$(\downarrow)$ |  Avg.$(\uparrow)$ |  Worst$(\uparrow)$ |  Gap$(\downarrow)$
ERM |  97.0(0.2) |  63.7(1.9) |  33.3 |  94.9(0.2) |  47.8(3.7) |  47.1
IRM |  87.5(0.7) |  75.6(3.1) |  11.9 |  94.0(0.4) |  77.8(3.9) |  16.2
IB-IRM |  88.5(0.6) |  76.5(1.2) |  12.0 |  93.6(0.3) |  85.0(1.8) |  8.6
CORAL |  90.3(1.1) |  79.8(1.8) |  10.5 |  93.8(0.3) |  76.9(3.6) |  16.9
Fish |  85.6(0.4) |  64.0(0.3) |  21.6 |  93.1(0.3) |  61.2(2.5) |  31.9
GroupDRO |  91.8(0.3) |  90.6(1.1) |  1.2 |  92.1(0.4) |  87.2(1.6) |  4.9
V-REX |  88.0(1.0) |  73.6(0.2) |  14.4 |  92.2(0.1) |  86.7(1.0) |  5.5
CGD |  91.3(0.6) |  88.9(0.8) |  2.4 |  92.5(0.2) |  90.0(0.8) |  2.5
DomainMix |  76.4(0.3) |  53.0(1.3) |  23.4 |  93.4(0.1) |  65.6(1.7) |  27.8
LISA |  91.8(0.3) |  89.2(0.6) |  2.6 |  92.4(0.4) |  89.3(1.1) |  3.1
Ours |  93.5(0.2) |  91.6(0.2) |  1.9 |  91.3(0.5) |  90.6(0.6) |  0.7
Datasets/ | CivilComments | FMoW |  Camely.
ACC(%) |  Avg.$(\uparrow)$ |  Worst$(\uparrow)$ |  Gap$(\downarrow)$ |  Avg.$(\uparrow)$ |  Worst$(\uparrow)$ |  Gap$(\downarrow)$ |  Avg.$(\uparrow)$
ERM |  92.2(0.1) |  56.0(3.6) | 36.2 |  53.0(0.6) |  32.3(1.3) | 20.7 |  70.3(6.4)
IRM |  88.8(0.7) |  66.3(2.1) | 22.5 |  50.8(0.1) |  30.0(1.4) | 20.8 |  64.2(8.1)
IB-IRM |  89.1(0.3) |  65.3(1.5) | 23.8 |  49.5(0.5) |  28.4(0.9) | 21.1 |  68.9(6.1)
CORAL |  88.7(0.5) |  65.6(1.3) | 23.1 |  50.5(0.4) |  31.7(1.2) | 18.8 |  59.5(7.7)
Fish |  89.8(0.4) |  71.1(0.4) | 18.7 |  51.8(0.3) |  34.6(0.2) | 17.2 |  74.7(7.1)
GroupDRO |  89.9(0.5) |  70.0(2.0) | 19.9 |  52.1(0.5) |  30.8(0.8) | 21.3 |  68.4(7.3)
V-REX |  90.2(0.3) |  64.9(1.2) | 25.3 |  48.0(0.6) |  27.2(0.8) | 20.8 |  71.5(8.3)
CGD |  89.6(0.4) |  69.1(1.9) | 20.5 |  50.6(1.4) |  32.0(2.3) | 18.6 |  69.4(7.8)
DomainMix |  90.9(0.4) |  63.6(2.5) | 27.3 |  51.6(0.2) |  34.2(0.8) | 17.4 |  69.7(5.5)
LISA |  89.2(0.9) |  72.6(0.1) | 16.6 |  52.8(0.9) |  35.5(0.7) | 17.3 |  77.1(6.5)
Ours |  90.6(0.4) |  73.3(0.3) | 17.3 |  49.5(1.5) |  38.2(1.7) | 11.3 |  77.0(5.5)

Model selection. Following prior works [30, 18], in the both group-aware and
group-oblivious settings, we assume the group information of validation
samples are available and select the best model based on worst-case accuracy
among all subpopulations on the validation set. In the group-oblivious
setting, we also conduct model selection based on the average accuracy to show
the impact of validation group label information in our method.

Comparison methods in the group-aware setting. In order to demonstrate the
performance of the proposed method, we compare our method with multiple
methods that use training group labels during training. The details are as
follows:

  * •

Invariant risk minimization and other closely related methods, including, IRM
[3] tries to learn the invariant correlations between different training
distributions; IB-IRM [78] improves IRM with information bottleneck; CORAL
[79] aims to align the correlations of deep neural networks layer activations
by a nonlinear transformation, and Fish [80] learns deep neural network with
invariant gradient direction for different populations which could promote
invariant predictions.

  * •

Group distributionally robust optimization and the follow-up algorithms,
including, GroupDRO [12] is an online optimization algorithm that dynamically
assigns importance weights during training to improve the worst-case
performance across all the subpopulations; V-REx [81] is an extension of
distributionally robust optimization by conducting robust optimization over a
perturbation set of extrapolated domain; CGD [63] assigns importance weights
to each subpopulation so that the model could be trained on the direction
leading to the largest decrease in average training loss.

  * •

Mixup-based methods. We also compare our method with DomainMix [43] and LISA
[42], where DomainMix introduces adversarial training to mix the samples to
synthesize novel examples while LISA mixes the samples within the same
subpopulation or same label.

Further, we also compare the proposed method with empirical risk minimization
(ERM).

Comparison methods in the group-oblivious setting. In the group-oblivious
setting, we compare the proposed method with multiple methods that do not
require training group information, including

  * •

ERM trains the model using standard empirical risk minimization;

  * •

Focal loss [21] downweights the well-classified examples’ loss according to
the current classification confidences.

  * •

CVaRDro and $\chi^{2}$Dro [25] minimize the loss over the worst-case
distribution in a neighborhood of the empirical training distribution.

  * •

CVaRDoro and $\chi^{2}$Doro [18] are outlier robust version of CVarDro and
$\chi^{2}$Dro respectively, which dynamically remove the outlier examples
during training.

  * •

JTT [30] is a two-stage reweighting algorithm to improve the subpopulation
shift robustness, which first constructs an error set and then upweights the
samples in the error set.

TABLE IV: Comparison results with other methods in the group-oblivious setting. We report the mean and corresponding standard deviations in parentheses for multiple random seeds. The best results are in bold and blue. Asterisks($*$) indicate that standard deviations were not provided in the original paper. _Note that the proposed method focuses on improving the model’s worst-case accuracy to improve model subpopulation shift robustness rather than the average accuracy._ It can be seen that the proposed method achieves the best performance on all datasets. Datasets/ | Waterbirds | CelebA | CivilComments | Camely.
---|---|---|---|---
ACC(%) | Avg.(↑) | Worst(↑) | Gap(↓) | Avg.(↑) | Worst(↑) | Gap(↓) | Avg.(↑) | Worst(↑) | Gap(↓) | Avg.(↑)
ERM | 97.0(0.2) | 63.7(1.9) | 33.3 | 94.9(0.2) | 47.8(3.7) | 47.1 | 92.2(0.1) | 56.0(3.6) | 36.2 | 70.3(6.4)
Focal Loss | 87.0(0.5) | 73.1(1.0) | 13.9 | 88.4(0.3) | 72.1(3.8) | 16.3 | 91.2(0.5) | 60.1(0.7) | 31.1 | 68.1(4.8)
CVaRDro | 90.3(1.2) | 77.2(2.2) | 13.1 | 86.8(0.7) | 76.9(3.1) | 9.9 | 89.1(0.4) | 62.3(0.7) | 26.8 | 70.5(5.1)
CVaRDoro | 91.5(0.7) | 77.0(2.8) | 14.5 | 89.6(0.4) | 75.6(4.2) | 14.0 | 90.0(0.4) | 64.1(1.4) | 25.9 | 67.3(7.2)
$\chi^{2}$Dro | 88.8(1.5) | 74.0(1.8) | 14.8 | 87.7(0.3) | 78.4(3.4) | 9.3 | 89.4(0.7) | 64.2(1.3) | 25.2 | 68.0(6.7)
$\chi^{2}$Doro | 89.5(2.0) | 76.0(3.1) | 13.5 | 87.0(0.6) | 75.6(3.4) | 11.4 | 90.1(0.5) | 63.8(0.8) | 26.3 | 68.0(7.5)
JTT | 93.6(∗) | 86.0(∗) | 7.6 | 88.0(∗) | 81.1(∗) | 6.9 | 90.7(0.3) | 67.4(0.5) | 23.3 | 69.1(6.4)
Ours | 93.0(0.5) | 90.0(1.1) | 3.0 | 90.1(0.4) | 85.3(4.1) | 4.8 | 90.6(0.4) | 70.1(0.9) | 20.5 | 75.1(5.9)
TABLE V: Experimental results when the group labels in the validation set are available or not. We report the average accuracy and worst-case accuracy of the models, while reporting the variation in performance of the model with and without group labels on the validation set.  Group labels in | Waterbirds | CelebA
---|---|---
validation set? | Average ${(\uparrow)}$ | Worst-case ${(\uparrow)}$ | Gap ${(\downarrow)}$ | Average (↑) | Worst-case (↑) | Gap ${(\downarrow)}$
Yes | 93.0(0.5) | 90.0(1.1) | 3.0 | 90.1(0.4) | 85.3(4.1) | 4.8
No | 93.6(0.5) | 88.9(0.8) | 4.7 | 90.4(0.4) | 84.6(0.8) | 5.8
Variations | +0.6 | -1.1 | +1.7 | +0.3 | -0.7 | +1.0

Comparison methods in the ablation study. 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 (IGMix) is performed on the samples with the same
labels and from the same subpopulations. Meanwhile, we also compare the
proposed method with the naive importance-weighting (IW) method.

Implementation details. Within each dataset, we keep the same model
architecture as in previous work [42]: ResNet-50 pretrained on ImageNet [82]
for Waterbirds and CelebA, DistilBERT [83] for CivilComments, DenseNet-121
pretrained on ImageNet[84] for FMoW, and DenseNet-121 without pretraining for
Camelyon17. For DistilBERT, we employ the HuggingFace [85] implementation and
start from the pre-trained weights. For the proposed method, we implement it
on the codestack released with the WILDS datasets [75]. For most comparison
methods, we directly use the results in previous work [42], original papers,
or reported in the WILDs leaderboard. For the rest of the comparison methods,
we reimplement them according to the original papers. We employ vanilla mixup
on WaterBirds and Camelyon17 datasets, CutMix [33] on CelebA and FMoW
datasets, and manifoldmix [36] on CivilComments dataset. For all approaches,
we tune all hyperparameters with AutoML toolkit NNI [86] based on validation
performance.

###  6.2 Group-aware experimental results

We conduct extensive comparison experiments in the group-ware setting on five
datasets to verify the robustness of the proposed model against subpopulation
shifts. The experimental results are reported in Tab. III. From the
experimental results, we have the following observations.

  * •

(1) The proposed method consistently outperforms all comparison methods in
terms of worst-case accuracy on all the datasets. For example, on the
Waterbirds dataset, the proposed method achieves the worst-case accuracy of
91.6%, which outperforms all the other comparison methods. Meanwhile, on the
FMoW dataset, the proposed method outperforms the second method by 2.7% in
terms of worst-case accuracy.

  * •

(2) ERM can achieve the best average accuracy on almost all used datasets, but
its corresponding worst-case accuracy performance is much lower than the
proposed method and comparison methods. The reason for this phenomenon is that
the ERM algorithm may over-focus on the majority subpopulation’s samples and
learn the spurious correlations, therefore improving the average performance
of the model based on the spurious correlation. However, due to over-reliance
on spurious correlations, the performance of ERM will degrade significantly in
the minority subpopulations. While our method can significantly improve the
performance of the minority subpopulations by introducing importance-weighted
mixup. For example, on the Waterbirds dataset and CelebA dataset, the worst-
case accuracy of the proposed method is 27.9% and 42.8% higher than ERM,
respectively.

  * •

(3) Compared with other comparison methods, the proposed method is also
competitive in terms of average accuracy. For example, on the Waterbirds and
CivilComments dataset, the average accuracy of the proposed method is in the
top three.

  * •

(4) The proposed method could significantly reduce the gap between the average
and worst-case accuracy, which is important for the fairness of the model.
Specifically, on the CelebA and FMoW datasets, the proposed method has the
lowest average and worst-case accuracy gap. On the Waterbirds and
CivilComments, the proposed method has the second-lowest average and worst-
case accuracy gap. Moreover, ERM usually has the largest average and worst-
case accuracy gap, which significantly affects the application of ERM to
fairness-sensitive tasks.

  * •

(5) By performing an importance-reweighted mixup and exploring the minority
subpopulations, the proposed method can also enhance the robustness of the
model to out-of-distribution data. Specifically, the proposed method achieves
the second-best performance i.e., 77.0% average accuracy, on the Camelyon17
dataset.

(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 in the training dataset. It can be seen from the experimental results that the proposed uncertainty estimation method can better characterize the subpopulation distribution of the training set and find hard samples belonging to the minority subpopulations.  TABLE VI: Comparison results with ERM, importance weighting, and mixup-based methods. From the experimental results, it can be seen that the proposed method significantly improves the worst-case accuracy of the model compared with mixup and importance weighting algorithms. Datasets/ | Experimental | Waterbirds | CelebA | CivilComments
---|---|---|---|---
ACC(%) | Setting | Avg.(↑) | Worst(↑) | Gap(↓) | Avg.(↑) | Worst(↑) | Gap(↓) | Avg.(↑) | Worst(↑) | Gap(↓)
ERM | oblivious | 97.0(0.2) | 63.7(1.9) | 33.3 | 94.9(0.2) | 47.8(3.7) | 47.1 | 92.2(0.1) | 56.0(3.6) | 36.2
mixup | oblivious | 81.0(0.2) | 56.2(0.2) | 24.8 | 95.8(0.0) | 46.4(0.5) | 49.4 | 90.8(0.8) | 67.2(1.2) | 23.6
IGMix | aware | 88.7(0.3) | 68.0(0.4) | 20.7 | 95.2(0.3) | 58.3(0.9) | 36.9 | 90.8(0.6) | 69.2(0.8) | 21.6
IW | aware | 95.1(0.3) | 88.0(1.3) | 7.1 | 92.9(0.2) | 83.3(2.8) | 9.6 | 89.8(0.5) | 69.2(0.9) | 20.6
Ours | oblivious | 93.0(0.5) | 90.0(1.1) | 3.0 | 90.1(0.4) | 85.3(4.1) | 4.8 | 90.6(0.4) | 70.1(0.9) | 20.5
Ours | aware | 93.5(0.5) | 91.6(0.2) | 1.9 | 91.3(0.5) | 90.6(0.6) | 0.7 | 90.6(0.4) | 73.3(0.3) | 17.3

###  6.3 Group-oblivious experimental results

In this section, we conduct various experiments in the group-oblivious setting
on multiple datasets with subpopulation shifts 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 Qualitative
analysis. Are the obtained uncertainties of the training samples trustworthy?

Q1 Effectiveness (I). When the group information in the training set is not
available, the proposed method could estimate the uncertainty of each training
sample and then construct sample-wise importance weights. Therefore, our
algorithm could work even when the training group information is not
available. We conduct experiments to verify its superiority over current
group-oblivious algorithms. The experimental results are shown in Table IV and
we have the following observations.

  * •

(1) Benefiting from exploring the hard samples with importance reweighted
mixup, the proposed RMix achieves the best worst-case accuracy on Waterbirds,
CelebA, and CivilComments datasets. For example, for the Waterbirds dataset,
RMix achieves worst-case accuracy of 90.0%, while the second-best is 86.0%.
Moreover, on the CelebA dataset, the proposed method outperforms the second-
best by 4.2%.

  * •

(2) As previously pointed out, ERM consistently outperforms other methods in
terms of average accuracy. However, it typically comes with the lowest worst-
case accuracy and highest average and worst-case accuracy gap.

  * •

(3) Compared to other methods, RMix still shows competitive average accuracy.
For example, on CelebA, RMix achieves the average accuracy of 90.1%, which
outperforms all other IW/DRO methods.

  * •

(4) The proposed method consistently achieves the lowest average and worst-
case accuracy gap. Specifically, the proposed method achieves 3.0%, 4.8%, and
20.5% accuracy gaps on Waterbirds, CelebA, and CivilComments datasets
respectively, while the second-best is 7.6%, 6.9%, and 23.3% respectively.

  * •

(5) Benefiting from exploring difficult samples, on the Camelyon17 dataset,
the proposed method achieves the best average accuracy. Specifically, RMix
achieves the best average accuracy of 75.1% while the second is 70.3%. This
shows that the proposed method can improve the generalization of the model
better than other reweighting algorithms.

Q2 Effectiveness (II). We further evaluate the Waterbirds and CelebA datasets
without using the validation set group label information and the experimental
results are shown in Table V. 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 or best worst-case accuracy. Finally,
we test the performance of the saved model on the test set. 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. Specifically, the average
accuracy on the Waterbirds and CelebA datasets increased by 0.6% and 0.3%,
respectively, while the worst-case accuracy decreased by 1.1% and 0.7%. It is
worth noting that the gap between the average and worst-case accuracy will
also increase slightly when the validation group label is not available.

Q3 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 are 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.

###  6.4 Ablation study

In this section, we conduct the ablation study to show the proposed method
could achieve better performance than mixup-based and importance weighting.
The experimental results are reported in Table LABEL:tab:ablation. From the
results, we have the following observations.

  * •

(1) The proposed method could slightly improve the model’s performance against
subpopulation shifts when the group information is available during training.

  * •

(2) Due to the underutilization of the minority subpopulation‘s samples, mixup
cannot improve the model’s robustness to subpopulation shifts significantly,
which also has been both experimentally and theoretically illustrated in
previous work [42].

  * •

(3) Compared with the mixup, the proposed method could achieve better
performance in terms of both average and worst-case accuracy on all the
datasets. For example, in the group-oblivious setting, the proposed method
outperforms mixup by 33.8%, 38.9%, and 2.9% in terms of worst-case accuracy on
Waterbirds, CelebA, and CivilComments datasets. Meanwhile, in the group-aware
setting, the proposed method also outperforms the in-group mixup.

  * •

(4) By equipping mixup with importance weighting, the proposed method
significantly improves the model robustness against subpopulation shifts.
Specifically, which is also consistent with our conclusions stated in the
Theorem 5.1, i.e., the proposed method has a tighter generalization bound than
the importance weighting method.

##  7 Conclusion

In this paper, we aim to propose a method to improve the robustness against
subpopulation shifts, which is critical for fair machine learning. To this
end, we proposed a simple yet effective approach called RMix to integrate
importance weights into the well-known mixup strategy so that RMix can
mitigate the overfitting during training thus improving prior IW methods. We
provide a strong theoretical intuition to show why the proposed method can
improve performance. To obtain reliable importance weights, we propose two
different strategies to adapt to the group-obvious and group-aware settings of
subpopulation shift. Specifically, in the group-aware setting, we propose a
method based on subpopulation size to set the weight of samples. Moreover, in
the group-oblivious setting, we propose methods to estimate the uncertainty of
training samples and set weights accordingly. We also provide insightful
theoretical analysis to show why the proposed method works better than vanilla
mixup and CutMix from the Rademacher complexity perspective. Specifically,
RMix shows the theoretical advantage that the learned model comes with a
subpopulation-heterogeneity dependent generalization bound. Extensive
experiments are conducted and the results show that the proposed method
consistently outperforms previous approaches on commonly-used benchmarks. In
the future, we will think about how to employ RMix in large-scale pre-trained
models to improve the fairness of the large-scale model.

## References

  * [1] H. Shimodaira, “Improving predictive inference under covariate shift by weighting the log-likelihood function,” _Journal of statistical planning and inference_ , vol. 90, no. 2, pp. 227–244, 2000.
  * [2] J. Huang, A. Gretton, K. Borgwardt, B. Schölkopf, and A. Smola, “Correcting sample selection bias by unlabeled data,” in _NeurIPS_ , 2006\.
  * [3] M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz, “Invariant risk minimization,” _arXiv preprint arXiv:1907.02893_ , 2019.
  * [4] S. Barocas and A. D. Selbst, “Big data’s disparate impact,” _Calif. L. Rev._ , vol. 104, p. 671, 2016.
  * [5] S. Bickel, M. Brückner, and T. Scheffer, “Discriminative learning for differing training and test distributions,” in _ICML_ , 2007.
  * [6] D. Mincu and S. Roy, “Developing robust benchmarks for driving forward ai innovation in healthcare,” _Nature Machine Intelligence_ , vol. 4, no. 11, pp. 916–921, 2022.
  * [7] A. Koenecke, A. Nam, E. Lake, J. Nudell, M. Quartey, Z. Mengesha, C. Toups, J. R. Rickford, D. Jurafsky, and S. Goel, “Racial disparities in automated speech recognition,” _Proceedings of the National Academy of Sciences_ , vol. 117, no. 14, pp. 7684–7689, 2020.
  * [8] T. Hashimoto, M. Srivastava, H. Namkoong, and P. Liang, “Fairness without demographics in repeated loss minimization,” in _ICML_ , 2018.
  * [9] N. Japkowicz, “The class imbalance problem: Significance and strategies,” in _IJCAI_ , 2000.
  * [10] S. Sagawa, A. Raghunathan, P. W. Koh, and P. Liang, “An investigation of why overparameterization exacerbates spurious correlations,” in _ICML_ , 2020\.
  * [11] Y. Cui, M. Jia, T.-Y. Lin, Y. Song, and S. Belongie, “Class-balanced loss based on effective number of samples,” in _CVPR_ , 2019.
  * [12] S. Sagawa*, P. W. Koh*, T. B. Hashimoto, and P. Liang, “Distributionally robust neural networks,” in _ICLR_ , 2020.
  * [13] K. Cao, C. Wei, A. Gaidon, N. Arechiga, and T. Ma, “Learning imbalanced datasets with label-distribution-aware margin loss,” in _NeurIPS_ , 2019\.
  * [14] W. Liu and S. Chawla, “Class confidence weighted knn algorithms for imbalanced data sets,” in _Pacific-Asia conference on knowledge discovery and data mining_. Springer, 2011, pp. 345–356.
  * [15] J. Wen, C.-N. Yu, and R. Greiner, “Robust learning under uncertain test distributions: Relating covariate shift to model misspecification,” in _ICML_ , 2014.
  * [16] R. Zhai, C. Dan, A. Suggala, J. Z. Kolter, and P. Ravikumar, “Boosted cvar classification,” in _NeurIPS_ , 2021.
  * [17] P. Michel, T. Hashimoto, and G. Neubig, “Modeling the second player in distributionally robust optimization,” in _ICLR_ , 2021.
  * [18] R. Zhai, C. Dan, Z. Kolter, and P. Ravikumar, “Doro: Distributional and outlier robust optimization,” in _ICML_ , 2021.
  * [19] P. Lahoti, A. Beutel, J. Chen, K. Lee, F. Prost, N. Thain, X. Wang, and E. Chi, “Fairness without demographics through adversarially reweighted learning,” in _NeurIPS_ , 2020.
  * [20] P. Michel, T. Hashimoto, and G. Neubig, “Distributionally robust models with parametric likelihood ratios,” in _ICLR_ , 2022.
  * [21] T.-Y. Lin, P. Goyal, R. Girshick, K. He, and P. Dollár, “Focal loss for dense object detection,” in _ICCV_ , 2017.
  * [22] J. Shu, Q. Xie, L. Yi, Q. Zhao, S. Zhou, Z. Xu, and D. Meng, “Meta-weight-net: Learning an explicit mapping for sample weighting,” in _NeurIPS_ , 2019.
  * [23] H. Namkoong and J. C. Duchi, “Stochastic gradient methods for distributionally robust optimization with f-divergences,” in _NeurIPS_ , 2016.
  * [24] W. Hu, G. Niu, I. Sato, and M. Sugiyama, “Does distributionally robust supervised learning give robust classifiers?” in _ICML_ , 2018.
  * [25] D. Levy, Y. Carmon, J. C. Duchi, and A. Sidford, “Large-scale methods for distributionally robust optimization,” in _NeurIPS_ , 2020.
  * [26] Z. Hu and L. J. Hong, “Kullback-leibler divergence constrained distributionally robust optimization,” _Available at Optimization Online_ , pp. 1695–1724, 2013.
  * [27] J. Byrd and Z. Lipton, “What is the effect of importance weighting in deep learning?” in _ICML_ , 2019.
  * [28] R. Zhai, C. Dan, Z. Kolter, and P. Ravikumar, “Understanding why generalized reweighting does not improve over erm,” _arXiv preprint arXiv:2201.12293_ , 2022.
  * [29] D. Xu, Y. Ye, and C. Ruan, “Understanding the role of importance weighting for deep learning,” in _ICLR_ , 2021.
  * [30] E. Z. Liu, B. Haghgoo, A. S. Chen, A. Raghunathan, P. W. Koh, S. Sagawa, P. Liang, and C. Finn, “Just train twice: Improving group robustness without training group information,” in _ICML_ , 2021.
  * [31] Z. Han, Z. Liang, F. Yang, L. Liu, L. Li, Y. Bian, P. Zhao, B. Wu, C. Zhang, and J. Yao, “UMIX: Improving importance weighting for subpopulation shift via uncertainty-aware mixup,” in _Advances in Neural Information Processing Systems_ , 2022.
  * [32] H. Zhang, M. Cisse, Y. N. Dauphin, and D. Lopez-Paz, “mixup: Beyond empirical risk minimization,” in _ICLR_ , 2018.
  * [33] S. Yun, D. Han, S. J. Oh, S. Chun, J. Choe, and Y. Yoo, “Cutmix: Regularization strategy to train strong classifiers with localizable features,” in _CVPR_ , 2019.
  * [34] L. Zhang, Z. Deng, K. Kawaguchi, A. Ghorbani, and J. Zou, “How does mixup help with robustness and generalization?” in _ICLR_ , 2021.
  * [35] L. Carratino, M. Ciss, R. Jenatton, and J.-P. Vert, “On mixup regularization,” _Journal of Machine Learning Research_ , vol. 23, no. 325, pp. 1–31, 2022.
  * [36] V. Verma, A. Lamb, C. Beckham, A. Najafi, I. Mitliagkas, D. Lopez-Paz, and Y. Bengio, “Manifold mixup: Better representations by interpolating hidden states,” in _ICML_ , 2019.
  * [37] O. Chapelle, J. Weston, L. Bottou, and V. Vapnik, “Vicinal risk minimization,” _Advances in neural information processing systems_ , vol. 13, 2000.
  * [38] J.-H. Kim, W. Choo, and H. O. Song, “Puzzle mix: Exploiting saliency and local statistics for optimal mixup,” in _International Conference on Machine Learning_. PMLR, 2020, pp. 5275–5285.
  * [39] C. Beckham, S. Honari, V. Verma, A. M. Lamb, F. Ghadiri, R. D. Hjelm, Y. Bengio, and C. Pal, “On adversarial mixup resynthesis,” _Advances in neural information processing systems_ , vol. 32, 2019.
  * [40] M. Faramarzi, M. Amini, A. Badrinaaraayanan, V. Verma, and S. Chandar, “Patchup: A regularization technique for convolutional neural networks,” _arXiv preprint arXiv:2006.07794_ , 2020.
  * [41] K. Zhou, Y. Yang, Y. Qiao, and T. Xiang, “Domain generalization with mixstyle,” in _International Conference on Learning Representations_ , 2020\.
  * [42] H. Yao, Y. Wang, S. Li, L. Zhang, W. Liang, J. Zou, and C. Finn, “Improving out-of-distribution robustness via selective augmentation,” in _ICML_ , 2022\.
  * [43] M. Xu, J. Zhang, B. Ni, T. Li, C. Wang, Q. Tian, and W. Zhang, “Adversarial domain adaptation with domain mixup,” in _AAAI_ , 2020.
  * [44] H. Yao, Y. Wang, L. Zhang, J. Zou, and C. Finn, “C-mixup: Improving generalization in regression,” in _Advances in Neural Information Processing Systems_.
  * [45] M. Hong, J. Choi, and G. Kim, “Stylemix: Separating content and style for enhanced data augmentation,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2021, pp. 14 862–14 870.
  * [46] A. Kendall and Y. Gal, “What uncertainties do we need in bayesian deep learning for computer vision?” in _NeurIPS_ , 2017.
  * [47] R. M. Neal, _Bayesian learning for neural networks_. Springer Science & Business Media, 2012, vol. 118.
  * [48] D. J. MacKay, “A practical bayesian framework for backpropagation networks,” _Neural computation_ , vol. 4, no. 3, pp. 448–472, 1992.
  * [49] J. Denker and Y. LeCun, “Transforming neural-net output levels to probability distributions,” in _NeurIPS_ , 1990.
  * [50] B. Lakshminarayanan, A. Pritzel, and C. Blundell, “Simple and scalable predictive uncertainty estimation using deep ensembles,” in _NeurIPS_ , 2017\.
  * [51] M. Havasi, R. Jenatton, S. Fort, J. Z. Liu, J. Snoek, B. Lakshminarayanan, A. M. Dai, and D. Tran, “Training independent subnetworks for robust prediction,” in _ICLR_ , 2021.
  * [52] J. Antorán, J. U. Allingham, and J. M. Hernández-Lobato, “Depth uncertainty in neural networks,” in _NeurIPS_ , 2020.
  * [53] G. Huang, Y. Li, G. Pleiss, Z. Liu, J. E. Hopcroft, and K. Q. Weinberger, “Snapshot ensembles: Train 1, get m for free,” in _ICLR_ , 2017.
  * [54] Q. V. Le, A. J. Smola, and S. Canu, “Heteroscedastic gaussian process regression,” in _ICML_ , 2005.
  * [55] D. A. Nix and A. S. Weigend, “Estimating the mean and variance of the target probability distribution,” in _ICNN_ , 1994.
  * [56] H. Ma, Z. Han, C. Zhang, H. Fu, J. T. Zhou, and Q. Hu, “Trustworthy multimodal regression with mixture of normal-inverse gamma distributions,” _NeurIPS_ , 2021.
  * [57] Z. Han, C. Zhang, H. Fu, and J. T. Zhou, “Trusted multi-view classification with dynamic evidential fusion,” _IEEE TPAMI_ , 2022.
  * [58] Y. Geng, Z. Han, C. Zhang, and Q. Hu, “Uncertainty-aware multi-view representation learning,” in _AAAI_ , 2021.
  * [59] A. Kendall, Y. Gal, and R. Cipolla, “Multi-task learning using uncertainty to weigh losses for scene geometry and semantics,” in _CVPR_ , 2018.
  * [60] D. Deng, L. Wu, and B. E. Shi, “Iterative distillation for better uncertainty estimates in multitask emotion recognition,” in _ICCV_ , 2021.
  * [61] G. Kalweit and J. Boedecker, “Uncertainty-driven imagination for continuous deep reinforcement learning,” in _Conference on Robot Learning_. PMLR, 2017, pp. 195–206.
  * [62] K. Li, A. Gupta, A. Reddy, V. H. Pong, A. Zhou, J. Yu, and S. Levine, “Mural: Meta-learning uncertainty-aware rewards for outcome-driven reinforcement learning,” in _ICML_ , 2021.
  * [63] V. Piratla, P. Netrapalli, and S. Sarawagi, “Focus on the common good: Group distributional robustness follows,” in _ICLR_ , 2022.
  * [64] D. Mahajan, R. Girshick, V. Ramanathan, K. He, M. Paluri, Y. Li, A. Bharambe, and L. van der Maaten, “Exploring the limits of weakly supervised pretraining,” in _ECCV_ , 2018.
  * [65] T. Mikolov, I. Sutskever, K. Chen, G. S. Corrado, and J. Dean, “Distributed representations of words and phrases and their compositionality,” _ArXiv_ , vol. abs/1310.4546, 2013.
  * [66] H. Ritter, A. Botev, and D. Barber, “A scalable laplace approximation for neural networks,” in _6th International Conference on Learning Representations, ICLR 2018-Conference Track Proceedings_ , vol. 6. International Conference on Representation Learning, 2018.
  * [67] Y. Geifman, G. Uziel, and R. El-Yaniv, “Bias-reduced uncertainty estimation for deep neural classifiers,” in _ICLR_ , 2019.
  * [68] M. Toneva, A. Sordoni, R. T. des Combes, A. Trischler, Y. Bengio, and G. J. Gordon, “An empirical study of example forgetting during deep neural network learning,” in _ICLR_ , 2019.
  * [69] J. Moon, J. Kim, Y. Shin, and S. Hwang, “Confidence-aware learning for deep neural networks,” in _ICML_ , 2020.
  * [70] P. L. Bartlett and S. Mendelson, “Rademacher and gaussian complexities: Risk bounds and structural results,” _Journal of Machine Learning Research_ , vol. 3, no. Nov, pp. 463–482, 2002.
  * [71] R. Arora, P. Bartlett, P. Mianjy, and N. Srebro, “Dropout: Explicit forms and capacity control,” in _ICML_ , 2021.
  * [72] Z. Liu, P. Luo, X. Wang, and X. Tang, “Deep learning face attributes in the wild,” in _ICCV_ , 2015.
  * [73] D. Borkan, L. Dixon, J. Sorensen, N. Thain, and L. Vasserman, “Nuanced metrics for measuring unintended bias with real data for text classification,” in _WWW_ , 2019.
  * [74] P. Bandi, O. Geessink, Q. Manson, M. Van Dijk, M. Balkenhol, M. Hermsen, B. E. Bejnordi, B. Lee, K. Paeng, A. 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_ , vol. 38, no. 2, pp. 550–560, 2018.
  * [75] P. W. Koh, S. Sagawa, H. Marklund, S. M. Xie, M. Zhang, A. Balsubramani, W. Hu, M. Yasunaga, R. L. Phillips, I. Gao _et al._ , “Wilds: A benchmark of in-the-wild distribution shifts,” in _ICML_ , 2021.
  * [76] G. Christie, N. Fendley, J. Wilson, and R. Mukherjee, “Functional map of the world,” in _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , 2018, pp. 6172–6180.
  * [77] A. Kumar, T. Ma, P. Liang, and A. Raghunathan, “Calibrated ensembles can mitigate accuracy tradeoffs under distribution shift,” in _The 38th Conference on Uncertainty in Artificial Intelligence_ , 2022.
  * [78] K. Ahuja, E. Caballero, D. Zhang, J.-C. Gagnon-Audet, Y. Bengio, I. Mitliagkas, and I. Rish, “Invariance principle meets information bottleneck for out-of-distribution generalization,” in _NeurIPS_ , 2021.
  * [79] B. Sun and K. Saenko, “Deep coral: Correlation alignment for deep domain adaptation,” in _European conference on computer vision_. Springer, 2016, pp. 443–450.
  * [80] Y. Shi, J. Seely, P. H. Torr, N. Siddharth, A. Hannun, N. Usunier, and G. Synnaeve, “Gradient matching for domain generalization,” _arXiv preprint arXiv:2104.09937_ , 2021.
  * [81] D. Krueger, E. Caballero, J.-H. Jacobsen, A. Zhang, J. Binas, D. Zhang, R. Le Priol, and A. Courville, “Out-of-distribution generalization via risk extrapolation (rex),” in _ICML_ , 2021.
  * [82] K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” in _CVPR_ , 2016.
  * [83] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “Bert: Pre-training of deep bidirectional transformers for language understanding,” in _NAACL_ , 2019\.
  * [84] G. Huang, Z. Liu, G. Pleiss, L. Van Der Maaten, and K. Q. Weinberger, “Convolutional networks with dense connectivity,” _IEEE transactions on pattern analysis and machine intelligence_ , vol. 44, no. 12, pp. 8704–8716, 2019.
  * [85] T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz _et al._ , “Huggingface’s transformers: State-of-the-art natural language processing,” _arXiv preprint arXiv:1910.03771_ , 2019.
  * [86] Microsoft, “Neural Network Intelligence,” 1 2021. [Online]. Available: https://github.com/microsoft/nni
  * [87] Y. Gal and Z. Ghahramani, “Dropout as a bayesian approximation: Representing model uncertainty in deep learning,” in _ICML_ , 2016.
  * [88] C. Park, S. Yun, and S. Chun, “A unified analysis of mixed sample data augmentation: A loss function perspective,” in _Advances in Neural Information Processing Systems_ , A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, Eds., 2022.

##  Appendix A More details about uncertainty-based importance weights

In this section, we present More details about uncertainty-based importance
weights. Specifically, we describe uncertainty quantification results on
simulated dataset in Sec. A.1, training accuracy of different subpopulations
throughout training process in Sec. A.2 and Justification for choosing
historical-based uncertainty score in Sec. A.3.

###  A.1 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(a). We compare our approximation (i.e., Eq. 13) 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.$
(18)

Specifically, we train $T$ models and then ensemble them. The quantification
results are shown in Fig. 2(b)(c). 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.

(a) Simulated dataset

(b) Ours

(c) Ensemble

Figure 2: (a) 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. (b)(c) Visualization of the obtained uncertainty with
kernel density estimation on simulated dataset, where group size refers to the
sample number of the group.

###  A.2 Training accuracy throughout training

We present how the training accuracy change throughout training in Fig. 3 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 3: Visualization of the changing of training accuracy on different
groups of CelebA and Waterbirds datasets.

###  A.3 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 [87, 50], 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 B Proof

In this appendix, we prove the Theorem 5.1 in Section 5 using the template
from the proof in [88] and incoporate our RMix. We define a training dataset
as $D=\\{z_{i}=(x_{i},y_{i})\\}_{i=1}^{m}$, randomly sampled from a
distribution $\mathcal{P}_{z}$. Then for randomly selected two data samples,
$z_{i}$ and $z_{j}$, $\tilde{z}_{i,j}^{\textsc{RMix}}$, is synthesized as
follows

$\displaystyle\tilde{z}_{i,j}^{\textsc{RMix}}(\lambda)=(\tilde{x}_{i,j}^{\textsc{RMix}}(\lambda),\tilde{y}_{i,j}^{\textsc{RMix}}(\lambda)),$

where

$\displaystyle\tilde{x}_{i,j}^{\textsc{RMix}}(\lambda)$
$\displaystyle=M(\lambda)\odot x_{i}+(1-M(\lambda))\odot x_{j},$ (19)
$\displaystyle\tilde{y}_{i,j}^{\textsc{RMix}}(\lambda)$ $\displaystyle=\lambda
y_{i}+(1-\lambda)y_{j},$

and $\lambda$ is the ratio parameter between samples drawn from
$\mathcal{D}_{\lambda}$. $\odot$ means a component-wise multiplication in
vector or matrix. $M(\lambda)$ is a random variable conditioned on $\lambda$
that indicates how we mix the input, e.g., by linear interpolation [32] or by
a pixel mask [33].

Under this template, two most popular Mixup methods, Mixup [32] and CutMix
[33], for $i$-th and $j$-th samples with $\lambda$ drawn from
$\mathcal{D}_{\lambda}$, can be rewriten as follow

$\displaystyle\text{ Mixup: }\quad\tilde{z}_{i,j}^{(\text{mixup
})}(\lambda)=\left(\tilde{x}_{i,j}^{(\text{mixup
})}(\lambda),\tilde{y}_{i,j}^{(\text{mixup })}(\lambda)\right)$ (20) where
$\displaystyle\tilde{x}_{i,j}^{\text{(mixup })}(\lambda)=\lambda
x_{i}+(1-\lambda)x_{j}\quad\text{ and }\quad\tilde{y}_{i,j}^{(\text{mixup
})}(\lambda)=\lambda y_{i}+(1-\lambda)y_{j}.$

$\displaystyle\text{ CutMix:
}\quad\tilde{z}_{i,j}^{\text{(cutmix)}}(\lambda)=\left(\tilde{x}_{i,j}^{\text{(cutmix
})}(M),\tilde{y}_{i,j}^{\text{(cutmix)}}(\lambda)\right)$ (21) where
$\displaystyle\tilde{x}_{i,j}^{\text{(cutmix)}}\left({M},1-{M}\right)={M}\odot
x_{i}+\left(1-{M}\right)\odot x_{j}\quad\text{ and
}\quad\tilde{y}_{i,j}^{\text{(cutmix) }}(\lambda)=\lambda
y_{i}+(1-\lambda)y_{j},$

i.e., the Mixup is a special case of Equation 19 by putting
$M(\lambda)=\lambda\overrightarrow{1}$ and the CutMixup is also a special case
where ${M}$ is a binary mask that indicates the location of the cropped box
region with a relative area.

We consider the following optimization objective, which is the expected
version of our weighted mixup loss (Equation 8).

$\displaystyle
L_{n}^{\textsc{RMix}}(\theta,S)=\frac{1}{n^{2}}\sum^{n}_{i,j=1}\mathbb{E}_{\lambda\sim
D_{\lambda}}\mathbb{E}_{M}[l(\theta,\tilde{z}_{i,j}(\lambda))],$

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}(\theta,S)=\frac{1}{n}\sum_{i=1}^{n}w_{i}[h(f_{\theta}(x_{i}))-y_{i}f_{\theta}(x_{i})].$

######  Lemma B.1.

The weighted mixup loss can be rewritten as

$\displaystyle
L_{n}^{\textsc{RMix}}(\theta,S)=L_{n}(\theta,S)+\sum_{i=1}^{3}\mathcal{R}_{i}(\theta,S)+\mathbb{E}_{\lambda\sim\tilde{\mathcal{D}}_{\lambda}}\mathbb{E}_{M}\left[(1-M)^{\top}\varphi(1-M)(1-M)\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)=\frac{1}{n}\sum_{i=1}^{n}\left(y_{i}-h^{\prime}\left(f_{\theta}\left(x_{i}\right)\right)\right)\left(\nabla
f_{\theta}\left(x_{i}\right)^{\top}x_{i}\right)\mathbb{E}_{\lambda\sim\tilde{D}_{\lambda}}(1-\lambda),$
(22)
$\displaystyle\mathcal{R}_{2}(\theta,S)=\frac{1}{2n}\sum_{i=1}^{n}h^{\prime\prime}\left(f_{\theta}\left(x_{i}\right)\right)\mathbb{E}_{\lambda\sim\tilde{D}_{\lambda}}\mathbf{G}\left(\mathcal{D}_{X},x_{i},f,M\right),$
$\displaystyle\mathcal{R}_{3}(\theta,S)=\frac{1}{2n}\sum_{i=1}^{n}\left(h^{\prime}\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}\right)\mathbb{E}_{\lambda\sim\tilde{D}_{\lambda}}\mathbf{H}\left(\mathcal{D}_{X},x_{i},f,M\right),$

and

$\displaystyle\mathbf{G}\left(\mathcal{D}_{X},x_{i},f,M\right)=\mathbb{E}_{M}(1-M)^{\top}\mathbb{E}_{r_{x}\sim\mathcal{D}_{X}}\left(\nabla
f\left(x_{i}\right)\odot\left(r_{x}-x_{i}\right)\left(\nabla
f\left(x_{i}\right)\odot\left(r_{x}-x_{i}\right)\right)^{\top}\right)(1-M)$
(23) $\displaystyle=\sum_{j,k\in\text{ coord
}}a_{jk}\partial_{j}f_{\theta}\left(x_{i}\right)\partial_{k}f_{\theta}\left(x_{i}\right)\left(\mathbb{E}_{r_{x}\sim\mathcal{D}_{X}}\left[r_{xj}r_{xk}\right]+x_{ij}x_{ik}\right),$
$\displaystyle\mathbf{H}\left(\mathcal{D}_{X},x_{i},f,M\right)=\mathbb{E}_{r_{x}\sim\mathcal{D}_{X}}\mathbb{E}_{M}(1-M)^{\top}\left(\nabla^{2}f_{\theta}\left(x_{i}\right)\odot\left(\left(r_{x}-x_{i}\right)\left(r_{x}-x_{i}\right)^{\top}\right)\right)(1-M)$
$\displaystyle=\sum_{j,k\in\text{ coord
}}a_{jk}\left(\mathbb{E}_{r_{x}\sim\mathcal{D}_{X}}\left[r_{xj}r_{xk}\partial_{jk}^{2}f_{\theta}\left(x_{i}\right)\right]+x_{ij}x_{ik}\partial_{jk}^{2}f_{\theta}\left(x_{i}\right)\right).$

###### Proof.

Due to the assumption of the loss function, we can rewrite the empirical loss
for the non-augmented population as

$L_{n}(\theta)=\frac{1}{n}\sum_{i=1}^{n}w_{i}l\left(\theta,z_{i}\right)=\frac{1}{n}\sum_{i=1}^{n}w_{i}\left[h\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}f_{\theta}\left(x_{i}\right)\right].$

Similarly, we can rewrite the RMix loss as

$\displaystyle L_{n}^{\textsc{RMix}}(\theta)$
$\displaystyle=\frac{1}{n^{2}}\sum_{i,j=1}^{n}\mathbb{E}_{\lambda\sim\operatorname{Beta}(\alpha,\beta)}\mathbb{E}_{M}l\left(\theta,\tilde{z}_{i,j}^{\textsc{RMix}}(M)\right)$
$\displaystyle=\frac{1}{n^{2}}\sum_{i,j=1}^{n}\mathbb{E}_{\lambda\sim\operatorname{Beta}(\alpha,\beta)}\mathbb{E}_{M}\left[h\left(f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)-\tilde{y}_{i,j}^{\textsc{RMix}}(M)f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right].$

Putting the definition of $\tilde{z}_{i,j}^{\textsc{RMix}}(M)$ to the equation
above, we have

$\displaystyle L_{n}^{\textsc{RMix}}(\theta)$
$\displaystyle=\frac{1}{n^{2}}\sum_{i,j=1}^{n}\left(\mathbb{E}_{\lambda\sim\operatorname{Beta}(\alpha,\beta)}\mathbb{E}_{M}\left(\lambda
w_{i}\left(h\left(f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)-y_{i}f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)\right.\right.$
$\displaystyle\left.\left.+(1-\lambda)w_{j}\left(h\left(f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)-y_{j}f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)\right)\right)$
$\displaystyle=\frac{1}{n^{2}}\sum_{i,j=1}^{n}\left(\mathbb{E}_{\lambda\sim\operatorname{Beta}(\alpha,\beta)}\mathbb{E}_{B\sim\operatorname{Bin}(\lambda)}\mathbb{E}_{M}\left(Bw_{i}\left(h\left(f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)-y_{i}f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)\right.\right.$
$\displaystyle\left.\left.+(1-B)w_{j}\left(h\left(f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)-y_{j}f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)\right)\right)$

Note that $\lambda\sim\operatorname{Beta}(\alpha,\beta)$ and
$B\mid\lambda\sim\operatorname{Bin}(\lambda)$. By conjugacy, we can write the
joint distribution of $(\lambda,B)$ as

$B\sim\operatorname{Bin}\left(\frac{\alpha}{\alpha+\beta}\right),\quad\lambda\mid
B\sim\operatorname{Beta}(\alpha+B,\beta+1-B).$

Therefore, we have

$\displaystyle L_{n}^{\textsc{RMix}}(\theta)$
$\displaystyle=\frac{1}{n^{2}}\sum_{i,j=1}^{n}\left(\mathbb{E}_{\lambda\sim\operatorname{Beta}(\alpha+1,\beta)}\mathbb{E}_{M}\frac{\alpha}{\alpha+\beta}w_{i}\left(h\left(f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)-y_{i}f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)\right.$
$\displaystyle\left.+\mathbb{E}_{\lambda\sim\operatorname{Beta}(\alpha,\beta+1)}\mathbb{E}_{M}\frac{\beta}{\alpha+\beta}w_{j}\left(h\left(f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)-y_{j}f_{\theta}\left(\tilde{x}_{i,j}^{\textsc{RMix}}(M)\right)\right)\right)$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}w_{i}\mathbb{E}_{\lambda\sim\tilde{\mathcal{D}}(\lambda)}\mathbb{E}_{r_{x}\sim\mathcal{D}_{x}}\mathbb{E}_{M}\left[h\left(f_{\theta}\left(w_{i}M\odot
x_{i}+(1-M)\odot r_{x}\right)\right)-y_{i}f_{\theta}\left(w_{i}M\odot
x_{i}+(1-M)\odot r_{x}\right)\right]$ (24)
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}w_{i}\mathbb{E}_{\lambda\sim\tilde{\mathcal{D}(\lambda)}}\mathbb{E}_{r_{x}\sim\mathcal{D}_{x}}\mathbb{E}_{M}l\left(\theta,\hat{z}_{i}\right),$
(25)

where $\mathcal{D}_{x}$ is the empirical distribution induced by training
samples and their corresponding weights, and $\hat{z}_{i}=\left(w_{i}M\odot
x_{i}+(1-M)\odot r_{x},y_{i}\right)$.

Let $N=1-M$. By defining
$\phi_{i}(N)=h\left(f_{\theta}\left(w_{i}x_{i}+N\odot\left(r_{x}-x_{i}\right)\right)\right)-y_{i}f_{\theta}\left(w_{i}x_{i}+N\odot\left(r_{x}-x_{i}\right)\right)$
and applying Taylor expansion, we have

$\phi_{i}(N)=\phi_{i}(0)+\nabla_{N}\phi_{i}(0)^{\top}N+\frac{1}{2}N^{\top}\nabla_{N}^{2}\phi_{i}(0)N+N^{\top}\varphi(N)N,$
(26)

where $\lim_{N\rightarrow 0}\varphi(N)=0$. Firstly, we calculate $\phi_{i}(0)$
by

$\phi_{i}(0)=h\left(f_{\theta}\left(w_{i}x_{i}\right)\right)-y_{i}f_{\theta}\left(w_{i}x_{i}\right).$
(27)

Second, we calculate $\nabla_{N}\phi_{i}(0)$ by

$\frac{\partial\phi_{i}(N)}{\partial
N_{k}}=\left(h^{\prime}\left(f_{\theta}\left(w_{i}x_{i}+N\odot\left(r_{x}-x_{i}\right)\right)\right)-y_{i}\right)\frac{\partial
f_{\theta}}{\partial
x_{ik}}\left(w_{i}x_{i}+N\odot\left(r_{x}-x_{i}\right)\right)\left(r_{xk}-x_{ik}\right),$

where we denote $N_{k}$ as the $k$ th element of $N,x_{ik}$ as the $k$ th
element of $x_{i}$, and $r_{xk}$ as the $k$ th element of $r_{x}$. Therefore,
we have

$\displaystyle\nabla_{N}\phi_{i}(0)^{\top}N$
$\displaystyle=\left(h^{\prime}\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}\right)\sum_{k}\left(\frac{\partial
f_{\theta}}{\partial
x_{ik}}\left(w_{i}x_{i}\right)\left(r_{xk}-x_{ik}\right)\right)N_{k}$ (28)
$\displaystyle=\left(h^{\prime}\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}\right)\left(\nabla
f\odot\left(r_{x}-x_{i}\right)\right)\cdot N.$

Finally, we calculate $\nabla_{N}^{2}\varphi_{i}(\overrightarrow{0})^{\top}$
by

$\displaystyle\frac{\partial^{2}\phi_{k}(N)}{\partial N_{k}\partial N_{j}}=$
$\displaystyle\frac{\partial}{\partial
N_{j}}\left(\left(h^{\prime}\left(f_{\theta}\left(w_{i}x_{i}+N\odot\left(r_{x}-x_{i}\right)\right)\right)-y_{i}\right)\frac{\partial
f_{\theta}}{\partial
x_{ik}}\left(w_{i}x_{i}+N\odot\left(r_{x}-x_{i}\right)\right)\left(r_{xk}-x_{ik}\right)\right)$
$\displaystyle=$ $\displaystyle
h^{\prime\prime}\left(f_{\theta}\left(w_{i}x_{i}+N\odot\left(r_{x}-x_{i}\right)\right)\right)$
$\displaystyle\times\frac{\partial f_{\theta}}{\partial
x_{ik}}\left(w_{i}x_{i}+N\odot\left(r_{x}-x_{i}\right)\right)\left(r_{xk}-x_{ik}\right)\frac{\partial
f_{\theta}}{\partial
x_{ij}}\left(w_{i}x_{i}+N\odot\left(r_{x}-x_{i}\right)\right)\left(r_{xj}-x_{ij}\right)$
$\displaystyle+$
$\displaystyle\left(h^{\prime}\left(f_{\theta}\left(w_{i}x_{i}+N\odot\left(r_{x}-x_{i}\right)\right)\right)-y_{i}\right)$
$\displaystyle\times\frac{\partial^{2}f_{\theta}}{\partial x_{ik}\partial
x_{ij}}\left(w_{i}x_{i}+N\odot\left(r_{x}-x_{i}\right)\right)\left(r_{xk}-x_{ik}\right)\left(r_{xj}-x_{ij}\right).$

Therefore, we have

$\displaystyle\frac{1}{2}N^{\top}\nabla_{N}^{2}\phi_{i}(0)N$
$\displaystyle=\frac{1}{2}h^{\prime\prime}\left(f_{\theta}\left(x_{i}\right)\right)\sum_{k,j}\left(\frac{\partial
f_{\theta}}{\partial
x_{ik}}\left(x_{i}\right)\left(r_{xk}-x_{ik}\right)\frac{\partial
f_{\theta}}{\partial
x_{ij}}\left(x_{i}\right)\left(r_{xj}-x_{ij}\right)N_{k}N_{j}\right)$
$\displaystyle\quad+\frac{1}{2}\left(h^{\prime}\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}\right)\sum_{k,j}\frac{\partial^{2}f_{\theta}}{\partial
x_{ik}\partial
x_{ij}}\left(x_{i}\right)\left(r_{xk}-x_{ik}\right)\left(r_{xj}-x_{ij}\right)N_{k}N_{j}$
$\displaystyle=\frac{1}{2}h^{\prime\prime}\left(f_{\theta}\left(x_{i}\right)\right)N^{\top}\left(\left(\nabla
f\odot\left(r_{x}-x_{i}\right)\right)\left(\nabla
f\odot\left(r_{x}-x_{i}\right)\right)^{\top}\right)N$
$\displaystyle\quad+\frac{1}{2}\left(h^{\prime}\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}\right)N^{\top}\left(\nabla^{2}f_{\theta}\left(x_{i}\right)\odot\left(\left(r_{x}-x_{i}\right)\left(r_{x}-x_{i}\right)^{\top}\right)\right)N.$

Applying 27-B to 26,

$\displaystyle\phi_{i}(N)$
$\displaystyle=\left(h\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}f_{\theta}\left(x_{i}\right)\right)+\left(h^{\prime}\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}\right)\left(\nabla
f\odot\left(r_{x}-x_{i}\right)\right)\cdot N$ (29)
$\displaystyle+\frac{1}{2}h^{\prime\prime}\left(f_{\theta}\left(x_{i}\right)\right)N^{\top}\left(\left(\nabla
f\odot\left(r_{x}-x_{i}\right)\right)\left(\nabla
f\odot\left(r_{x}-x_{i}\right)\right)^{\top}\right)N$
$\displaystyle+\frac{1}{2}\left(h^{\prime}\left(f_{\theta}\left(x_{i}\right)\right)-y_{i}\right)N^{\top}\left(\nabla^{2}f_{\theta}\left(x_{i}\right)\odot\left(\left(r_{x}-x_{i}\right)\left(r_{x}-x_{i}\right)^{\top}\right)\right)N+N^{\top}\varphi(N)N$

Plugging 29 to 24, we conclude

$\displaystyle L_{n}^{\textsc{RMix}}(\theta)$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}_{\lambda\sim\tilde{\mathcal{D}}(\lambda)}\mathbb{E}_{r_{x}\sim\mathcal{D}_{x}}\mathbb{E}_{M}\phi(1-M)$
$\displaystyle=L_{n}(\theta)+\mathcal{R}_{1}(\theta)+\mathcal{R}_{2}(\theta)+\mathcal{R}_{3}(\theta)+\mathbb{E}_{\lambda\sim\tilde{\mathcal{D}}(\lambda)}\mathbb{E}_{M}\left[(1-M)^{\top}\varphi(1-M)M\right],$

where

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

∎

######  Lemma B.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}(\theta,S)+\frac{C}{2n}[\sum_{i=1}^{n}w_{i}A^{\prime\prime}(x_{i}^{\top}\theta)]\theta^{\top}(\mathbb{E}(1-M)\widehat{\Sigma}_{X}(1-M)^{\top}+x_{i}\operatorname{Var}(M)x_{i}^{\top})\theta,$

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

###### Proof.

For GLM, using Equation 25, since the prediction of the GLM model is invariant
to the scaling of the training data, we consider the dataset
$\hat{D}=\left\\{\hat{z}_{i}\right\\}_{i=1}^{n}$ with
$\hat{x_{i}}=1\oslash\bar{M}\odot\left(w_{i}M\odot x_{i}+(1-M)\odot
r_{x}\right)$ where $\bar{M}=\mathbb{E}M$. Then, the loss function is

$L_{m}^{\textsc{RMix}}=\frac{1}{n}\mathbb{E}_{\lambda}\mathbb{E}_{r_{x}}\mathbb{E}_{M}\sum_{i=1}^{n}w_{i}l\left(\theta,\tilde{z}_{i}\right)=\frac{1}{n}\mathbb{E}_{\xi}\sum_{i=1}^{n}w_{i}\left(A\left(\hat{x}_{i}^{\top}\theta\right)-y_{i}\hat{x}_{i}^{\top}\theta\right),$

where $\xi$ denotes the randomness of $\lambda,r_{x}$, and $M$. By the second
approximation of $A(\cdot)$, we can express
$A\left(\hat{x}_{i}^{\top}\theta\right)$ as

$A\left(\hat{x}_{i}^{\top}\theta\right)=A\left(x_{i}^{\top}\theta\right)+A^{\prime}\left(x_{i}^{\top}\theta\right)\left(\hat{x_{i}}-x_{i}\right)^{\top}\theta+\frac{1}{2}A^{\prime\prime}\left(x_{i}^{\top}\theta\right)\theta^{\top}\left(\hat{x_{i}}-x_{i}\right)\left(\hat{x_{i}}-x_{i}\right)^{\top}\theta$

to approximate the loss function. Therefore, we have

$\displaystyle\tilde{L}_{m}^{\textsc{RMix}}$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}w_{i}A\left(x_{i}^{\top}\theta\right)+\frac{1}{n}\mathbb{E}_{\xi}\sum_{i=1}^{n}w_{i}\left(A^{\prime}\left(x_{i}^{\top}\theta\right)\left(\hat{x}_{i}-x_{i}\right)^{\top}\theta+\frac{1}{2}A^{\prime\prime}\left(x_{i}^{\top}\theta\right)\theta^{\top}\left(\hat{x}_{i}-x_{i}\right)\left(\hat{x}_{i}-x_{i}\right)^{\top}\theta\right)$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}w_{i}A\left(x_{i}^{\top}\theta\right)+\frac{1}{n}\sum_{i=1}^{n}w_{i}\left(\frac{1}{2}A^{\prime\prime}\left(x_{i}^{\top}\theta\right)\theta^{\top}\operatorname{Var}_{\xi}\left(\hat{x_{i}}\right)\theta\right),$

where $\tilde{L}_{m}^{\textsc{RMix}}$ denotes the approximate loss of
$L_{m}^{\textsc{RMix}}$ since $\mathbb{E}_{\xi}r_{x}=0$ and
$\mathbb{E}_{\xi}\hat{x}_{i}=x_{i}$. For calculating
$\operatorname{Var}_{\xi}\left(\hat{x}_{i}\right)$, we use the law of total
variance. We have

$\displaystyle\operatorname{Var}_{\xi}\left(\tilde{x_{i}}\right)$
$\displaystyle=\left(\frac{1}{\bar{M}}\frac{1}{\bar{M}}^{\top}\right)\odot\operatorname{Var}_{\xi}\left(M\odot
x_{i}+(1-M)\odot r_{x}\right)$
$\displaystyle=\left(\frac{1}{\bar{M}}\frac{1}{\bar{M}}^{\top}\right)\odot\left(\mathbb{E}\left(\operatorname{Var}\left(M\odot
x_{i}+(1-M)\odot
r_{x}\mid\lambda,M\right)+\operatorname{Var}\left(\mathbb{E}\left(M\odot
x_{i}+(1-M)\odot r_{x}\mid\lambda,M\right)\right)\right.\right.$
$\displaystyle\left.=\left(\frac{1}{\bar{M}}\frac{1}{\bar{M}}^{\top}\right)\odot\left(\mathbb{E}(1-M)\hat{\Sigma}_{X}(1-M)^{\top}+x_{i}\operatorname{Var}(M)x_{i}^{\top}\right)\right)$
$\displaystyle=\frac{1}{\bar{\lambda}^{2}}\left(\mathbb{E}(1-M)\hat{\Sigma}_{X}(1-M)^{\top}+x_{i}\operatorname{Var}(M)x_{i}^{\top}\right)$

where $\hat{\Sigma}_{X}=\frac{1}{n}\sum_{i=1}^{n}w_{i}x_{i}x_{i}^{\top}$ with
some notational ambiguity that
$\frac{1}{n}:=\overrightarrow{1}\oslash\bar{M}$. In our setting
$\bar{M}=\bar{\lambda}\overrightarrow{1}$ where
$\bar{\lambda}=\mathbb{E}_{\lambda\sim\tilde{\mathcal{D}}_{\lambda}}[\lambda]$.
Now we think the related dual problem:

$\displaystyle\mathcal{W}_{\gamma}=\left\\{x\rightarrow\theta^{\top}x,\text{
such that
}\theta:\left(\mathbb{E}_{x}A^{\prime\prime}\left(\theta^{\top}x\right)\right)\cdot\left(\theta^{\top}\left(\mathbb{E}(1-M)\Sigma_{X}(1-M)^{\top}\right)\theta+\theta^{\top}\left(\left(x\operatorname{Var}(M)x^{\top}\right)\right)\theta\right)\leq\gamma\right\\}.$

∎

######  Lemma B.3.

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

$\displaystyle
Rad(\mathcal{W}_{r})\leq\frac{1}{\sqrt{n}}(\gamma/\rho)^{1/4}\left(\sqrt{\operatorname{tr}\left(\left(\Sigma_{X}^{(M)}\right)^{\dagger}\Sigma_{X}\right)}+\operatorname{rank}\left(\Sigma_{X}\right)\right).$

###### Proof.

Define $\Sigma_{X}^{(M)}=\mathbb{E}[(1-M)\Sigma_{X}(1-M)^{\top}]$. The proof
is mainly based on [34]. By definition, given $n$ i.i.d. Rademacher rv.
$\xi_{1},\ldots,\xi_{n}$, the empirical Rademacher complexity is

$\displaystyle\operatorname{Rad}\left(\mathcal{W}_{\gamma},S\right)$
$\displaystyle=\mathbb{E}_{\xi}\sup_{a(\theta)\cdot\theta^{\top}(\Sigma_{X}^{(M)}+x\operatorname{Var}(M)x^{\top})\theta\leq\gamma}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}\theta^{\top}x_{i}$
$\displaystyle\leq\mathbb{E}_{\xi}\sup_{(\mathbb{E}_{x}A^{\prime\prime}(\theta^{\top}x))\cdot\theta^{\top}\Sigma_{X}^{(M)}\theta\leq\gamma}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}\theta^{\top}x_{i}$
$\displaystyle\leq\mathbb{E}_{\xi}\sup_{\mathbb{E}_{x}A^{\prime\prime}(\theta^{\top}x))\cdot\theta^{\top}\Sigma_{X}^{(M)}\theta\leq\gamma}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}\theta^{\top}x_{i}$
$\displaystyle\leq\mathbb{E}_{\xi}\sup_{(\theta^{\top}\Sigma_{X}\theta)\cdot\theta^{\top}\Sigma_{X}^{(M)}\theta\leq\gamma}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}\theta^{\top}x_{i}$
$\displaystyle\leq\mathbb{E}_{\xi_{i}}\left(\sup_{\theta\top\Sigma_{X}\theta\leq\sqrt{\gamma/\rho}}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}\theta^{\top}x_{i}+\mathbb{E}_{\xi_{i}}\sup_{\theta\top\Sigma_{X}^{(M)}\theta\leq\sqrt{\gamma/\rho}}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}\theta^{\top}x_{i}\right),$

where the third inequality is from Assumption 5.1.

For the first part of the RHS, define
$\tilde{x}_{i}=\Sigma_{X}^{\dagger/2}x_{i}$ and $v=\Sigma_{X}^{1/2}\theta$.
Then, we have

$\displaystyle\mathbb{E}_{\xi_{i}}\sup_{\theta\top\Sigma_{X}\theta\leq\sqrt{\gamma/\rho}}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}\theta^{\top}x_{i}=\mathbb{E}_{\xi_{i}}\sup_{\|v\|^{2}\leq\sqrt{\gamma/\rho}}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}v^{\top}\tilde{x}_{i}$
$\displaystyle\leq\frac{1}{n}(\gamma/\rho)^{1/4}\mathbb{E}_{\xi_{i}}\left\|\sum_{i=1}^{n}\xi_{i}\tilde{x}_{i}\right\|\leq\frac{1}{n}(\gamma/\rho)^{1/4}\sqrt{\mathbb{E}_{\xi_{i}}\left\|\sum_{i=1}^{n}\xi_{i}\tilde{x}_{i}\right\|^{2}}$
$\displaystyle=\frac{1}{n}(\gamma/\rho)^{1/4}\sqrt{\sum_{i=1}^{n}\tilde{x}_{i}^{\top}\tilde{x}_{i}}.$

Similarly, by defining
$\check{x}_{i}=\left(\Sigma_{X}^{(M)}\right)^{\dagger/2}x_{i}$ and
$v=\left(\Sigma_{X}^{(M)}\right)^{1/2}\theta$,

$\displaystyle\mathbb{E}_{\xi_{i}}\sup_{\theta\top\Sigma_{X}^{(M)}\theta\leq\sqrt{\gamma/\rho}}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}\theta^{\top}x_{i}$
$\displaystyle=\mathbb{E}_{\xi_{i}}\sup_{\|v\|^{2}\leq\sqrt{\gamma/\rho}}\frac{1}{n}\sum_{i=1}^{n}\xi_{i}v^{\top}\check{x}_{i}$
$\displaystyle\leq\frac{1}{n}(\gamma/\rho)^{1/4}\mathbb{E}_{\xi_{i}}\left\|\sum_{i=1}^{n}\xi_{i}\check{x}_{i}\right\|\leq\frac{1}{n}(\gamma/\rho)^{1/4}\sqrt{\mathbb{E}_{\xi_{i}}\left\|\sum_{i=1}^{n}\xi_{i}\check{x}_{i}\right\|^{2}}$
$\displaystyle=\frac{1}{n}(\gamma/\rho)^{1/4}\sqrt{\sum_{i=1}^{n}\check{x}_{i}^{\top}\check{x}_{i}}.$

Therefore,

$\displaystyle\operatorname{Rad}\left(\mathcal{W}_{\gamma}\right)$
$\displaystyle=\mathbb{E}\left[\operatorname{Rad}\left(\mathcal{W}_{\gamma},n\right)\right]\leq\frac{1}{n}(\gamma/\rho)^{1/4}\left(\sqrt{\sum_{i=1}^{n}\mathbb{E}_{x}\tilde{x}_{i}^{\top}\tilde{x}_{i}}+\sqrt{\sum_{i=1}^{n}\mathbb{E}_{x}\check{x}_{i}^{\top}\check{x}_{i}}\right)$
$\displaystyle\leq\frac{1}{\sqrt{n}}(\gamma/\rho)^{1/4}\left(\sqrt{\operatorname{tr}\left(\left(\Sigma_{X}^{(M)}\right)^{\dagger}\Sigma_{X}\right)}+\operatorname{rank}\left(\Sigma_{X}\right)\right).$

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