UMIX: Improving Importance Weighting for Subpopulation Shift via Uncertainty-Aware Mixup
Authors: Zongbo Han
Zhipeng Liang
Fan Yang
Liu Liu
Lanqing Li
Yatao Bian
Peilin Zhao
Bingzhe Wu
Changqing Zhang
Jianhua Yao
Venue: NeurIPS 2022 Accept (2022)
Keywords: subpopulation shift
uncertainty
mixup

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

UMIX: Improving Importance Weighting for Subpopulation Shift via Uncertainty-Awarr Mixup
Zongbo Han Zhipeng Liang2* Fan Yang Liu", Lanqing Li Yatao Bian
Peilin Zhao?, Bingzhe Wu Changqing Zhang' Jianhua Yao
"College of Intelligence and Computing, Tianjin University.
2 Hong Kong University of Science and Technology, "Tencent AI Lab
Abstract
Subpopulation shift widely exists in many real-world machine learning applications, referring to the training and test distributions containing the same subpopulation groups but varying in subpopulation frequencies. Importance reweighting is a normal way to handle the subpopulation shift issue by imposing constant or adap- tive sampling weights on each sample the training dataset. However, some recent studies have recognized that most of these approaches fail to improve the performance over empirical risk minimization especially when applied over- parameterized neural networks. In this work, we propose simple yet practical framework, called uncertainty-awer awaaee mixup (UMIx), mitigate the overfitting issue in over- perraaamtteeized models by reweighting the "mixed"" samples according the sample uncertainty. The ceaninng-trajettriinn based uncertainty estimation is equipped in the proposed UMIX for each sample to flexibly characterize the subpopulation distribution. We also provide insightful theoretical analysis to verify that UMIX achieves better generalization bounds over prior works. Further, we conduct extensive empirical studies across a wide range of tasks to validate the effectiveness of our method both qualitatively and quantitatively. Code is available at this URL
1 Introduction
Empirical risk minimization (ERM) typically faces challenges from distribution shift, which refers to the difference between training and test distributions R
27 61 One common type of distribution shift subpopulation shift wherein the training and test distributions consist of the same subpopulation groups but differ in subpopulation frequencies 6 J. Many practical research problems (e.g., fairness of machine learning and class imbalance) can be considered as a special case of subpopulation shift For example, in the setting of fair machine learning, dataset with biased demographic subpopulations and test on an unbiased test dataset 241 32].
21 28 32]
train the model training
Therefore the essential goal of fair machine learning is mitigate the subpopulation shift between training and test datasets.
Many approaches have been proposed for solving this problem. Among these approaches, importance weighting (IW) is a classical yet effective technique by imposing static or adaptive weights on each sample when building weighted empirical loss. Therefore each subpopulation group contributes comparably to the final training objective. Specifically, there are normally two ways to achieve importance reweighting. Early works propose to reweight the sample inverse proportionally to the
subpopulation frequencies (i.c., static weights) [ 13 42 58 59 611, such as class-imbalanced
Equal contribution. & Supported by 2021 Tencent Rhino Bird Research Elite Training Program. Work donc during an internship at Tencent Al Lab. i Corresponding authors: (bingzhesu, haaeeeuu,,iinnaaaagee iianuayao) antnaaya))@tenceetttcom zhangchangginget, cn.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

learning approaches [1 13 42
Alternatively, a more flexible way is to reweight individual samples
adaptively according to training dynamics [35] optimization (DRO) is one of the most representative methods in this line, which minimizes the loss over the worst-case distribution in a neighhorhood of the empirical training distribution. A commonly used dual form of DRO can be seen as a special case of importance reweighting wherein the sampling in an alternated manner.
66 72 74. Distributional robust
weights are updated based on the current loss [24 25 38 52]
However, some recent studies have shown both empirically and theoretically that these IW methods could fail to achieve better worst-case subpopulation performance compared with ERM. Empirically, prior works [10, (58] recognize that various IW methods tend to exacerbate overfitting, which leads to a diminishing effect on stochastic gradient descent (SGD) over training epochs especially when they are applied to over weer-prameterized neural networks (NNs). Theoretically, previous studies prove that for over-pere per-parameterized neural networks, reweighting algorithms do not improve over ERM because their implicit biases are (almost) equivalent [59, 68, 73]. In addition, some prior works also point out that using conventional regularization techniques such as weight decay cannot significantly improve the performance of IW (58).
To this end, we introduce a novel technique called uncertainty-aware mixup (UMIX), by reweighting the mixed samples according to uncertainty within the mini-batch while mitigating overfitting. Specifically, we employ the well-known mixup technique to produce "mixed" augmented samples. Then we train the model on these mixed samples to make sure it can always see "novel" samples thus the effects of IW will not dissipate even at the end of the training epoch. To enforce the model to perform fairly well on all subpopulations, we further efficiently reweight the mixed samples according to uncertainty of the original samples. The weighted mixup loss function is induced by combining the weighted losses of the corresponding two original samples. At high level, this approach augments training samples in an uncertainty-awarr manner, i.e., putting more focus on samples with higher prediction uncertainties that belong to minority subpopulations with high probabilities. We also show UMIX can provide additional theoretical benefit which achieves a tighter generalization bound than weighted ERM
38 40
72..
The contributions of this paper are:
We propose a simple and practical approach called uncertainty-awwe aware mixup (UMIx) to improve previous IW methods by reweighting the mixed samples, which provides a new framework to mitigate overfitting in aver-parameterized neural networks.
Under the proposed framework, we provide theoretical analysis with insight that UMIX can achieve a tighter generalization bound than the weighted ERM.
We perform extensive experiments on wide range of tasks, where the proposed UMIX achieves excellent performance in both group-oblivious and group-aware settings.
Comparison with existing works. Here, we discuss the key differences between UMIX and other works. In contrast to most IW methods (e.g., CVaR-DRO (38) and JTT [41). UMIX employs mixup strategy to improve previous IW methods and mitigate the model overfitting. Among these methods, JTT [41] and LISA (70) are the two most related works to ours. Specifically, JTT provides a two-stage optimization framework in which an additional network is used for building the error set, and then JTT upweights samples in the error set in the following training stage. Besides, LISA also modifies mixup for improving model robustness against distribution shift. However, LISA intuitively mixes the samples within the same subpopulation same label thus it needs additional subpopulation information. In contrast to them, UMIX introduces sample weights into the vanilla mixup strategy by quantitatively measuring the sample uncertainties without subpopulation information. In addition, our work is orthogonal to LISA, i.e., we can use our weight building strategy to improve LISA's performance. In practice, our method consistently outperforms previous approaches that do not use subpopulation information and even achieves quite competitive performance to those methods which leverage subpopulation information. We also provide theoretical analysis to explain why UMIX
works better than the weighted ERM (3 40
72).
Related Work
2.1 Importance weighting
To improve the model robustness against subpopulation shift, importance weighting (IW) is a classical yet effective technique by imposing static or adaptive weight on each sample and then building

---

weighted empirical loss. Therefore each subpopulation group can have a comparable strength in the final training objective. Specifically, there are typically two ways to achieve importance reweighting. i.e., using static or adaptive importance weights.
Static methods. The naive reweighting approaches perform static reweighting based the dis- tribution of training samples [T] 13
61]. Their core motivation is to make different subpopulations have a comparable contribution 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
42
58
59
[58 59 61] Besides, there are some methods to obtain sample and Bayesian weights based on the effective number of samples [3], subpopulation margins
networks 442]
Adaptive methods. In contrast to the above static methods, a more essential way is to assign each
individual sample an adaptive weight that can vary according training dynamics R5 40 47 48
62 66 72 [4). Distributional robust optimization (DRO) is one of the most representative methods in this line, which minimizes the loss over the worst-case distribution in neighborhood of the empirical training distribution. A commonly-used dual form of DRO can be considered as a special case of importance reweighting wherein the sampling weights are updated based on the current loss [24] [2 38, [2] in an alternated manner. For example, in the group-aware setting (i.e., we know each sample belongs to which subpopulation), GroupDRO [58] introduces an online optimization algorithm to update the weights of each group. In the group-oblivion setting. [35] 47] 48, 66] model the problem as (regularized) minimax game, where one player aims to minimize the loss by optimizing the model parameters and another player aims to maximize the loss by assigning weights to each sample.
2.2 Uncertainty quantification
The core of our method is based on the high-quality uncertainty quantification of each sample. There many approaches proposed for this goal. The uncertainty of deep learning models includes epistemic (model) uncertainty and aleatoric (data) uncertainty (30] To obtain the epistemic uncer- tainty, Bayesian neural networks (BNNs) 15
30
45 53 have been proposed which replace the
deterministic weight parameters of model with distribution. Unlike BNNs, ensemble-basee methods obtain the epistemic uncertainty by training multiple models and ensembling them 2 22 26 36).
Alcatoric uncertainty focuses on the inherent noise in the data, which usually is learned as a function of the data (30 37 54). Uncertainty quantification has been successfully equipped in many fields
such as multimodal learning 19 20 44, multitask learning 14 31. and reinforcement learning 29 39. Unlike previous methods, our method focuses on estimating the epistemic uncertainty of training samples with subpopulation shift and upweighting uncertain samples, thereby improving the performance of minority subpopulations with high uncertainty.
3 Method
In this section, we introduce technical details of UMIX. The key idea of UMIX is to exploit uncertainty information upweight mixed samples, and thus can encourage the model to perform uniformly well on all subpopulations. We first introduce the basic procedure of UMIX and then present how to provide high-quality uncertainty estimations which is the fundamental block of UMIX.
3.1 Background
The necessary background and notations are provided here. Let the input and label space be X' and y respectively. Given training dataset D with N training samples {2,, y,) 1 i.i.d. sampled from probability distribution P. We consider the setting that the training distribution predefined subpopulations, i.e., P= - kg Pg. where kg and P, denote the g-th subpopulation's proportion and distribution respectively. Our goal is to obtain a model fo X + y parameterized by 0 E θ that performs well on all subpopulations.
a mixture of G
The well-known empirical risk mininization (ERM) algorithm doesn't it consider the subpopulations and minimizes the expected risk B[[(0, (1; v:3]. where & denotes the loss function. This leads to the model paying more attention to the majority subpopulations in the training set and resulting in poor performance on the minority subpopulations. For example, the ERM-based models may learn
3

---

spurious correlations that exist in majority subpopulations but not in minority subpopulations [58]. The proposed method aims to learn a model that is robust against weighting.
subpopulattiii subpopulation shift by importance
Previous works on improving subpopulation shift robustness investigate several different settings, i.e., group-aware and group-oblivious [41 58 72 Most of the previous works have assumed that the group label available during training [58 70 This is called the group-aware setting, However, due to some reasons, we may not have training group labels. For example, in many real applications, it's hard to extract group label information. Meanwhile, the group label information may not be available due to privacy concerns. This paper studies the -roupioblivious setting, which cannot obtain group information for each example training time. This requires the model identify underperforning samples and then pay more attention to them during training.
3.2 Importance- mixup
UMIX employs an aggressive data augmentation strategy called uncertainty- mixup to mitigate overfitting- Specifically, vanilla mixup 77 76 constructs virtual training examples (i.e., mixed samples) by performing linear interpolations between data/features and corresponding labels as:
ز Ar; + (1 - A)I,, Tiej Azy + (1 - A)y5>
(1)
where (I). (1y 81) are two samples drawn at random from empirical training distribution and [0, usually sampled from beta distribution. Then vanilla mixup optimizes the following loss function:
(c ay) [e(O, Fagy F(g))-
(2)
When the cross entropy loss is employed, Eq. 2 can be rewritten as:
E ((rraa))) [Ae Ta.j, v6) + (1
(3)
Eq. 3 can be seen as a linear combination (mixup) of E(O, F..j+ Vs) and W.., mj). Unfortunately. since vanilla mixup doesn't i consider the subpopulations with poor performance, it has been shown experimentally be non-robust against subpopulation shift [70]. To this end, we introduce simple yet effective method called UMIX, which further employs a weighted linear combination of the original loss based on Eq. to encourage the learned model to pay more attention to samples with poor performance.
In contrast to previous IW methods, the importance weights of UMix are used on the mixed samples. To do this, we first estimate the uncertainty of each sample and then use this quantity to construct the importance weight (i.c., the higher the uncertainty, the higher the weight, and vice versa). For the i-th sample F,, we denote its importance weight as Wig Once we obtain the importance weight, we can perform weighted linear combination of E(0, Fe.y yi) and E(O, F1,5. U;) by:
E ((( [w; (0, Fegs ve) + wy - A)( b;)5
where wi and w, denote the importance weight of the i-th and j-th samples respectively. In practice, to balance the UMIX and normal training, we set a hyperparameter that denotes the probability to apply UMIX. The whole training pseudocode for UMIX is shown in Algorithm I
Algorithm 1:
The training pscudocode of UMIX.
Input:
Training dataset D and the corresponding importance weights w [wi,... ... wNn),
hyperparameter a to control the probability of doing UMIX, and parameter a of the beta distribution;
1 for each iteration do
Obtain training samples (I;; y/). (2.5, Uz) and the corresponding weight w; wy;
3 Sample p Oni))rm(0.1);
if < then Sample X Beta a); else X ==0 0;
5 Obtain the mixed input ã., where Tigj Ao; + (1 A)æy;
Obtain the loss of the model with we Xe(O, Fi.j. Mh) + wy X)e(0, Te,, yj);
Update model parameters 0 to minimize loss with an optimization algorithm.

---

3.3 Uncertainty-awart importance weights
Now we present how to obtain the uncertainty-awe aware training importance weights. In the group- oblivious setting, the key to obtaining importance weights is to find samples with high uncertainty. For example, DRO-based algorithms construct the uncertainty set with the current loss 25 [52]
has been shown experimentally that the uncertain samples found in this way are constantly changing during training [11, resulting these methods not always upweighting the minority subpopulations. Therefore, we introduce sampling-bane based stable uncertainty estimation to better characterize the subpopulation shift.
Given a well trained neural classifier fo X + y that could produce the predicted class fz(x). simple way to obtain the uncertainty of sample is whether the sample is correctly classified. However, as pointed out previous work [36], single model cannot accurately characterize the sampling uncertainty. Therefore, we propose to obtain the uncertainty through Bayesian sampling from the model posterior distribution p(0; D). Specifically, given a sample (r;, yt), we define the training uncertainty as:
if = y fo( u, = f k(y (o())))p(0;; D)d0, where k(y,, fo(xx)) = 1, if + fo (x))
O,
(5)
Then, we can obtain an approximation of Eq. 5 with T Monte Carlo samples as (z.)), where 0g θ can be obtained by minimizing the expected risk.
~
In practice, sampling {0:) from the posterior (i.e., O, p(0; D)) is computationally expensive and sometimes even intractable since multiple training models need be built or extra approximation errors need be introduced. Inspired by recent Bayesian learning paradigm named SWAG [46), we propose to employ the information from the historical training trajectory to approximate the sampling process. More specifically, we train a model with ERM and save the prediction results fo, (æ:) of each sample each iteration epoch t. Then, to avoid the influence of inaccurate predictions the beginning of training, we estimate uncertainty with predictions after training T, I epochs with:
T.+T
U1 RV
<(y,, So, (2:)).
(6)
We empirically show that the proposed approximation could obtain reliable uncertainty in Sec. B.4 of the Appendix.
To obtain reasonable importance weights, assume that the samples with high uncertainty should be given a higher weight and vice versa. Therefore a reasonable importance weight could linearly positively related to the corresponding uncertainty.
Wi: = ryu; + C
where 7) E R4 is hyperparameter and e E R4 is a constant that keeps the weight to be positive. In practice, we set 1. The whole process for obtaining training importance weights shown in Algorithm 2
Algorithm 2: The process for obtaining training importance weights.
Input: Training dataset D, sampling start epoch T., the number of sampling T, and upweight
hyperparameter 7
Output: The training importance weights "wi."",,,,,,
for each iteration do
Train fa by minimizing the expected risk E{e(0, v:)}:
Save the prediction results (fo, ( of the current epoch
TT+t7 Obtain the uncertainty of each sample with u; + Dt=,, s(vi, fo, (a:)): Obtain the importance weight of each sample with w; 1u4 +
Remark, Total uncertainty can be divided into epistemic and aleatoric uncertainty (3) In the proposed method, the samples weighted only based epistemic uncertainty by sampling from
5

---

the model on the training trajectory, which can be seen as sampling from the model posterior in more efficient way. What's more, we consider that the training samples do not contain the inherent noise (aleatoric uncertainty) since it is usually intractable to distinguish between noisy samples and minority samples from data with subpopulation shifts.
Rethink why this estimation approach could work? Recent work has empirically shown that compared with the hard-to-classify samples, the easy-to-classify samples are learned earlier during training [8] Meanwhile, the hard-to-classify samples are also more likely to be forgotten by the neural networks 664]. The frequency with which samples are correctly classified during training can be used as supervision information in confidence calibration [3]]. Snapshot performs ensemble learning on several local minima models along the optimization path [26.. The proposed method is also inspired by these observations and algorithms. During training, samples from the minority sub- populations are classified correctly less frequently, which corresponds to higher training uncertainty. On the other hand, samples from the majority subpopulations will have lower training uncertainty due to being classified correctly more often. In Sec. Sec. 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. of the Appendix.
4 Experiments
In this section, we conduct experiments on multiple datasets with subpopulation shift to answer the following questions. QI Effectiveness (I). In the group-oblivious setting, does the proposed method outperform other algorithms? Q2 Effectiveness (II). How does UMIX perform without the group labels in the validation set? Q3 Effectiveness (III). Although our method does not use training group labels, does it perform better than the algorithms using training group labels? Q4 Effectiveness (IV). Can UMIx improve the model robustness against domain shift where the training and test distributions have different subpopulations. Q5 Qualitative analysis. Are the obtained uncertainties of the training samples trustworthy? Q6 Ablation study. What is the key factor of performance improvement our method?
4.1 Setup
We briefly present the experimental setup here, including the experimental datasets, evaluation metrics, model selection, and comparison methods. Please refer Sec. B in Appendix for more detailed setup.
We perform experiments on three datasets with multiple subpopulations, including Water- birds (58), CelebA 43 and CivilComments (9]. We also validate UMIX on domain shift scenario which is a more challenging distribution shift problem since there are different subpopulations be- tween training and test data. Hence, we conduct experiments on medical dataset called Camelyon17 17 33 that consists of pathological images from five different hospitals. The training data drawn from three hospitals, while the validation and test data are sampled from other hospitals.
Datasets.
Evaluation metries. To be consistent with existing works 33 56 70), we report the average accuracy of Camelyon17 over 10 different random seeds. On other datasets, we repeat experiments over 3 times and report the average and worst-case accuracy among all subpopulations. The trade-off between the average and worst-case accuracy is a well-known challenge [21]. In this paper, we lay emphasis on worst-case accuracy, which is more important than average accuracy in some application scenarios. For example, in fairness-related applications, we should pay more attention to the performance of the minority groups to reduce the gap between the majority groups and ensure the fairness of the machine learning decision system.
we assume the group labels of validation samples are available and select the best model based on worst-case accuracy among all subpopulations on the validation set. We also conduct model selection based on the average accuracy to show the impact of validation group label information in our method.
Model selection. Following prior works 41
72
Comparisons in the group-oblivious setting, Here we list the baselines used in the group-oblivious setting. (1) ERM trains the model using standard empirical risk minimization. (2) Focal loss 440 downweights the well-classified examples' loss according to the current classification confidences.

---

(3) DRO- bRO--ased methods including CVaR-DRO, DRO B8, CVaR-DORO and x² DDRO 72 minimize the loss over the worst-case distribution a neighborhood of the empirical training distribution. (4) JTT [4] constructs an error set and upweights the samples in the error set to improve the worst-case performance among all subpopulations.
Comparison in the group-aware setting. To better demonstrate the performance of the proposed method, we compare our method with multiple methods that use training group labels, including IRM [3], IB-IRM T. V-REx [4]. CORAL 63 Group DRO [38] DomainMix 69 Fish 60]. and LISA 77).
-ixup-based comparison methods. We compare our method with vanilla mixup and in in-group mixup, where vanilla mixup is performed any pair of samples and in in-group mixup is performed on the samples with the same labels and from the same subpopulations.
4.2 Experimental results
We present experimental results and discussions to answer the above-posede questions.
Q1 Effectiveness (I). Since our algorithm does not need training group labels, thus we conduct experiments to verify its superiority over current group-oblivious algorithms. The experimental results are shown in Table T and we have the following observations: (1) The proposed UMIX achieves the best worst-case accuracy worst-case accuracy of 85.3%, while the second-best is 81. 1%. (2) ERM consistently outperforms other methods in terms of average accuracy. However, it typically comes with the lowest worst-case accuracy. The underlying reason is that the dominance of the majority subpopulations during training leads to poor performance of the minority subpopulations. accuracy compared to other methods. For example, on CelebA, UMIX achieves the average accuracy of 90.10 1%, which outperforms all other IW/DRO methods.
all three datasets. For example, for the CelebA dataset, UMIX achieves
(3) UMIX shows competitive average
Q2 Effectiveness (II). We conduct the evaluation on the Waterbirds and CelebA datasets without using the validation set group label information. Specifically, after each training epoch, we evaluate the performance the current model on the validation set and save the model with the best average accuracy. Finally, we test the performance of the saved model on the test set. The experimental results are shown in Table 2 From the experimental results, we can observe that when the validation set group information is not used, the worst-case accuracy of our method drops a little while the average accuracy improves a little.
Q3 Effectiveness (III). We further conduct comparisons with algorithms that require training group labels. The comparison results are shown in Table According to the experimental results, is observed that the performance from our UMIX without using group label is quite competitive compared with these group-aware algorithms. Specifically, benefiting from the uncertainty- mixup, UMIX usually performs in the top three in terms of both average and worst-case accuracy. For example, on WaterBirds, UMIX achieves the best average accuracy 93.0% and the second-best worst-case accuracy of 90.0%.
Table 1:
Comparison results with other methods in the group-oblivious setting. The best results are in bold and blue. Full results with standard deviation are in the Table in Appendix.
Waterbirds
CelebA
CivilComments
Camelyon17 Avg.
Avg.
Worst
Avg.
Worst
Avg.
Worst
ERM
97.0%
63.7%
94.9%
47.8%
92.2%
56.0%
70.3%
Focal Loss 40
87.0%
73.1%
88.4%
72.1%
91.2%
60.1%
68.1%
CVaR-DRO
38
90.3%
77.2%
86.8%
76.9%
39.1%
62.3%
70.5%
CVaR-DORO 72
91.5%
77.0%
89.6%
75.6%
90.0%
64.1%
67.3%
2-DRR
38
88.8%
74.0%
87.7%
78.4%
89.4%
64.2%
68.0%
--DORO 72
89.5%
76.0%
87.0%
75.6%
90.1% %
63.8%
68.0%
JTT 41
93.6%
86.0%
88.0%
81.1%
90.7%
67.4%
69.1%
Ours
93.0%
90.0% %
90.1%
85.3 %
90.6%
70.1%
75.1 %

---

Table 2: Experimental results when the group labels in the validation set are available or not.
Group labels in validation set?
Waterbirds
CelebA
Average ACC
Worst-case ACC
Average ACC
Worst-case ACC
Yes No
93.00%
90.00%
90.10%
85.30%
93.60%
88.90%
90.40%
84.60%
Table 3: Comparison results with the algorithms using training group labels (Our method is not dependent on this type of information). Results of baseline models are from [70] The best three results are in bold brown or bold blue and the color indicates whether the training group labels are used. Full results with standard deviation are in the Table Appendix.
Group labels in train set?
Waterbirds
CelebA
CivilComments
Cam17
Avg.
Worst
Avg.
Worst
Avg.
Worst
Avg.
IRM 3
Yes
87.5%
75.6%
94.0%
77.8%
88.8%
66.3%
64.2%
IB-IRMM
Yes
88.5%
76.5%
93.6%
85.0%
89.1%
65.3%
68.9%
V-REx 34
Yes
88.0%
73.6%
92.2%
86.7%
90.2%
64.9%
71.5%
CORAL 63
Yes
90.3%
79.8%
93.8%
76.9%
88.7%
65.6%
59.5%
GroupDRO
588
Yes 91.8%
90.6%
92.1%
87.2%
89.9%
70.0%
68.4%
Domain.Mis
66
Yes 76.4%
53.0%
93.4%
65.6%
90.9%
63.6%
69.7%
Fish 660
Yes 85.6%
64.0%
93.1%
61.2%
89.8%
71.1%
74.7%
LISA 70
Yes
91.8%
89.2%
92.4%
89.3%
89.2%
72.6%
77.1% %
Ours
No 93.0%
90.0%
90.1% %
85.3%
90.6%
70.1 %
75.1 %
Q4 Effectiveness (IV). We conduct comparison experiments Camelyon17 to investigate the effectiveness of our algorithm under the domain shift scenario. The experimental results are shown in the last column of Table and Table respectively. In the group-obliviing setting, the proposed method achieves the best average accuracy on Camelyon17 as shown in Table For example, UMIX achieves the best average accuracy of 75. while the second is 70.3%. Meanwhile, in Table 3 benefiting from upweighting the mixed samples with poor performance, our method achieves quite competitive generalization ability on Camelyon17 compared with other algorithms using training group labels.
Group 0 size: 3498
Group size: 71629
Group 1 size: 184
40
Group size: 66874
40
Group 2 size:
Group size: 22880
Density
Density
30
Group 3 size: 1057
Group size: 1387
30
20
20
10
10
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Uncertainty
Uncertainty
(a) Waterbirds
(b) CelebA
Figure 1: Visualization of the obtained uncertainty with kernel density estimation Waterbirds and CelebA datasets, where group size refers to the sample number of the group.
Q5 Qualitative analysis. To intuitively investigate the rationality of the estimated uncertainty, we visualize the density the uncertainty for different groups with kernel density estimation. As shown in Fig- [ the statistics of estimated uncertainty is basically correlated to the training sample size of each group. For example, on Waterbirds and CelebA, the average uncertainties of minority groups much higher, while those of majority groups are much lower.

---

Q6 Ablation study. Finally, we conduct the ablation study in comparison with vanilla mixup and in-group mixup. The experimental results are shown in Table | Compared with ERM, vanilla mixup cannot significantly improve worst-case accuracy. After using the group label, the in-group mixup slightly improves the worst-case accuracy compared to ERM. The possible reason is that mixup-based methods do not increase the influence of minority subpopulations in the model objective function. Although our method does not use the group label the training samples, our method can still significantly improve the worst-case accuracy.
Table 4: Comparison with ERM and mixup based methods. Results of baseline models are from 70 The best results are in bold brown bold blue and the color indicates whether the training group labels are used. Full results with standard deviation are in the Table & in Appendix.
Group labels in train set?
Waterbirds
CelebA
CivilComments
Cam
Avg.
Worst
Avg.
Worst
Avg.
Worst
Avg.
ERM vanilla mixup in-group mixup
No
97.0% %
63.7%
94.9%
47.8%
92.2%
56.0%
70.3%
No
81.0%
56.2%
95.8%
46.4%
90.8%
67.2%
71.2%
Yes
88.7%
68.0%
95.2%
58.3%
90.8%
69.2%
75.5%
Ours
No
93.0%
90.0% %
90.1%
85.3%
90.6%
70.1
75.1%
Theory
In this section, we provide theoretical understanding of the generalization ability for UMIX. high level, we prove that our method can achieve a better generalization error bound than traditional IW methods without using mixup. For simplicity, our analysis focuses on generalized linear model (GLM). The roadmap of our analysis is to first approximate the mixup loss and then study the generalization bound from a Rademacher complexity perspective. To introduce the theoretical framework, we first present the basic settings.
Basic settings. Our analysis mainly focuses on GLM model classes whose loss function é follows E(0, v. y) = A x) yo x, where x E IRd is the input 0 E IRd is the parameter, E IR is the label and A--) 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 's proportion being ky and the g-th subpopulation follows the distribution Pg. Specifically, we have P = ΣG P=E.-1 kig Pg. Then we denote the covariance matrix for the g-th subpopulation Σx [ [x:x For simplicity, we consider the case where a shared weight wy 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 I Formally, we focus on analyzing the upper bound of the weighted generalization error defined as:
1 N
GError(O) (asooon)))(ra))))aee)
w(rt (:),
where the function w(z, y) is the weighted function return the weight of the subpopulation to which the sample (x, 3) belongs.
First of all, we present our main result in this section. The main theorem of our analysis provides seapulaatooooettoggeeeeee formally presented as:
dependent bound for the above generalization error. This theorem is
Theorem 5.1. Suppose A(-) is -Lipscchtt continuous, then there exists constants L, > 0 such that for any 0 satisfying 03 Exx S.1 7, the following holds with a probability of at least - 5,
GError ()) < 2L LA (max{( (5/2)
(5/2)
rank(E. Σx
log log(2/8) (2/8)
1/4
1123
+B
P
2n
where (6) is constant dependent on 8, Σx the data distribution, which will be formally introduced in Assumption 5.1
-
Ex and p is some constant related to

---

We will show later that the output of our Algorithm 1 can satisfy the constraint 0 Σx - Y and thus Theorem 5.T can provide a theoretical understanding of our algorithm. In contrast weighted ERM, the bound improvement of UMIX is the red term which can partially reflect the heterogeneity of the training subpopulations. Specifically, the red term would become vd/n n in the weighted ERM setting (see more detailed theoretical comparisons Appendix). Thus our bound can be tighter when the intrinsic dimension of data is small (i.e., rank(Ex) << d).
The proof of Theorem 5.1 follows this roadmap: (1) We first show that the model learned with UMIX can fall into specific hypothesis set Wy. (2) We analyze the Rademacher complexity of the hypothesis set and obtain its complexity upper bound (Lemmal A..3. (3) Finally, we can characterize
the generalization bound by using complexity-basee learning theory 7 (Theorem 8). More details of the proof can be found in Appendix.
As we discuss in Appendix, the weighted mixup can be seen as an approximation a regularization C 1E term ns
i=1 W; (x; 0)]0* Ex0 for some constant C compared with the non non-mixup algorithm,
which motivates us to study the following hypothesis space
W { 00 x. such that 0 satisfying E E (i)lu((z.v))"( (xT0))0' 0))oT Σx s 7)
for some constant 7.-
To further derive the generalization bound, we also need the following assumption, which is satisfied by general GL.Ss when has bounded & norm and adopted in, e.g., 76).
Assumption 5.1 (p-retentive). We say the distribution of is p-retentive for some p E (0, if for any non-zero vector E Rd and given the event that O E Wy where the 0 is output by our Algorithm 7 we have
Ex [T] P. (in x=)?].
Finally, we can derive the Rademacher complexity of the Wy and the proof of Theorem 5.1 is obtained by combining Lemma A.3 an3 the Theorem 8 of 77.
Lemma 5.1. Assume that the distribution of x; is p-retentive, iee, satisfies the assumption 5.1 Then the empirical Rademacher complexity of Wr satisfies
Rad(Wr,, S max{( (666) 1/2 (6) 1/23.
rank(Ex (Ex)
P
with probability at least 1 6.
6 Conclusion
In this paper, we propose a novel method called UMIX to improve the model robustness against subpopulation shift. We propose simple yet reliable approach to estimate the sample uncertainties and integrate them into the mixup strategy so that UMIX can mitigate the overfitting thus improving over prior IW methods. Our method consistently outperforms previous approaches used benchmarks. Furthermore, UMIx also shows the theoretical advantage that the learned model comes with leverage subpopulation information to improve UMIX can be a promising research direction.
commonly-
hetetttionn-eteergeeeety dependent generalization bound. In the future. how to
Acknowledgements
This work was supported in part by the National Key Research and Development Program of China under Grant 201911B2101900, the National Natural Science Foundation of China (61976151, 61925602, 61732011).
References
[1] Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Jean-Christophe Gagnon-Audet, Yoshua Bengio, loannis Mitliagkas, and Irina Rish. Invariance principle meets information bottleneck for out-of-distributic generalization. In NeurIPS, 2021.
10

---

[2] Javier Antorán, James Urquhart Allingham, and José Miguel Hernández-Lobato. Depth uncer- tainty in neural networks. In NeurIPS, 2020.
[3] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization. arXiv preprint arXiv: 197700290,, 2019.
[4] Raman Arora, Peter Bartlett, Poorya Mianjy, and Nathan Srebro. Dropout: Explicit forms and capacity control. In ICML, 2021.
[5] Peter Bandi, Oscar Geessink, Quirine Manson, Marcory Van Dijk, Maschenka Balkenhol, Meyke Hermsen, Bahak Ehteshami Bejnordi, Byungjae Lee, Kyunghyun Paeng, Aoxiao Zhong, et al. From detection of individual metastases to classification of lymph node status at the patient level: the camelyon17 challenge. IEEE transactions on medical imaging, 38(2):500-560, 2018.
[6] Solon Barocas and Andrew D Selbst. Big data's disparate impact. Calif. L. Rev. 104:671, 2016.
[7] Peter L Bartlett and Shahar Mendelson, Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463-482, 2002.
[8] Steffen Bickel, Michael Brückner, and Tobias Scheffer. Discriminative learning for differing training and test distributions. In ICML, 2007.
[9] Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. In WWW, 2019.
[10] Jonathon Byrd and Zachary Lipton. What the effect of importance weighting in deep learning? In ICML, 2019.
[11] Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, and Tengyu Ma. Learning imbalanced datasets with aabel-disribbutio--aaal margin loss. In NeurlPS, 2019.
[12] Yongqiang Chen, Yonggang Zhang, Yatao Bian, Han Yang, Kaili Ma, Binghui Xie, Tongliang Liu, Bo Han, and James Cheng. Invariance principle meets duist--dist generalization on graphs. arXiv preprint arXiv:2202.0544 2022.
[13] Yin Cui, Menglin Jia, Tsung- Yi Lin, Yang Song, and Serge Belongic. Class-balanced loss based on effective number of samples. In CVPR, 2019.
[14] Didan Deng, Liang Wu, and Bertram E Shi. Iterative distillation for better uncertainty estimates in multitask emotion recognition. In ICCV, 2021.
[15] John Denker and Yann LeCun. Transforming neural-net output levels to probability distributions. In NeurlPS, 1990.
[16] Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL, 2019.
[17] Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In ICML,, 2016.
[18] Yonatan Geifman, Guy Uziel, and Ran El- Yaniv. Bias-reduced uncertainty estimation for deep neural classifiers. In ICLR, 2019.
[19] Yu Geng, Zongbo Han, Changqing Zhang, and Qinghua Hu. Uncertainty-aware multi-view representation learning. In AAAI, 2021.
[20] Zongbo Han, Changqing Zhang, Huazhu Fu, and Joey Tianyi Zhou. Trusted multi-view classification with dynamic evidential fusion. IEEE TPAMI, 2022.
[21] Tatsunori Hashimoto, Megha Srivastava. Hongseok Namkoong, and Percy Liang. Fairness without demographics in repeated loss minimization. In ICML, 2018.
[22] Marton Havasi, Rodolphe Jenatton, Stanislay Fort, Jeremiah Zhe Liu, Jasper Snoek, Balaji Lakshminarayanan, Andrew Mingbo Dai, and Dustin Tran. Training independent subnetworks for robust prediction. In ICLR, 2021.

---

[23] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, 2016.
[24] Weihua Hu, Gang Niu, Issei Sato, and Masashi Sugiyama. Does distributionally robust super- vised learning give robust classifiers? In ICML, 2018.
[25] Zhaolin Hu and Jeff Hong. Kullback-leibler divergence constrained distributionally robust optimization. Available at Optimization Online, pages 1695-1724, 2013.
[26] Gao Huang, Yixuan Li, Geoff Pleiss, Zhuang Liu, John E Hoperoft, and Kilian Q Weinberger. Snapshot ensembles: Train 1, get m for free. In ICLR, 2017.
[27] Jiayuan Huang, Arthur Gretton, Karsten Borgwardt, Bernhard Schölkopf, and Alex Smola. Correcting sample selection bias by unlabeled data. In NeurlPS, 2006.
[28] Nathalie Japkowicz. The class imbalance problem: Significance and strategies. IJCAI, 2000.
[29] Gabriel Kalweit and Joschka Boedecker. Uncertainty-driven imagination for continuous deep reinforcement learning. In Conference on Robot Learning, pages 195-206. PMLR, 2017.
[30] Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for computer vision? In NeurIPS, 2017.
[31] Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In CVPR, 2018.
[32] Allison Koenecke, Andrew Nam, Emily Lake, Joe Nudell, Minnie Quartey, Zion Mengesha, Connor Toups. John Rickford, Dan Jurafsky, and Sharad Goel. Racial disparities in automated speech recognition. Proceedings of the National Academy of Sciences, 116414)77644-669 2020.
[33] Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: benchmark of in-the-wild distribution shifts. In ICML, 2021.
[34] David Krueger, Ethan Caballero, Joem-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang. Remi Le Priol, and Aaron Courville. -ut-o--dist generalization via risk extrapo- lation (rex). In ICML, 2021.
[35] Preethi Lahoti, Alex Beutel, Jilin Chen, Kang Lee, Flavien Prost, Nithum Thain, Xuezhi Wang. and Ed Chi. Fairness without demographics through adversarially reweighted learning. In NeurlPS, 2020.
[36] Balaji Lakshninarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In NeurIPS, 2017.
[37] Quoe V Le, Alex Smola, and Stéphane Canu. Heteroscedastic gaussian process regression. In ICML, 2005.
[38] Daniel Levy, Yair Carmon, John C Duchi, and Aaron Sidford. Large-scale methods for distributionally robust optimization. In NeurlPS, 2020.
[39] Kevin Li, Abhishek Gupta, Ashwin Reddy, Vitchyr H Pong, Aurick Zhou, Justin Yu, and Sergey Levine. Mural: Meta-learning uncertainty-wware rewards for outcome-die driven reinforcement learning. In JCML, 2021.
[40] Tsung- Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. In ICCV, 2017.
[41] Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In ICML, 2021.
12

---

[42] Wei Liu and Sanjay Chawla. Class confidence weighted knn algorithms for imbalanced data sets. In Pacific-Asia conference on knowledge discovery and data mining, pages 345-356. Springer, 2011.
[43] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In ICCV, 2015.
[44] Huan Ma, Zongbo Han, Changqing Zhang, Huazhu Fu, Joey Tianyi Zhou, and Qinghua Hu. Trustworthy multimodal regression with mixture of normal-inverse gamma distributions. NeurlPS, 2021.
[45] David JC MacKay. A practical bayesian framework for backpropagation networks. Neural computation, 4(3):448-472, 1992.
[46] Wesley J Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson, simple baseline for bayesian uncertainty in deep learning. In NeurlPS, 2019.
[47] Paul Michel, Tatsunori Hashimoto, and Graham Neubig. Modeling the second player in distributionally robust optimization. In ICLR, 2021.
[48] Paul Michel, Tatsunori Hashimoto, and Graham Neubig. Distributionally robust models with parametric likelihood ratios. In ICLR, 2022.
[49] Microsoft. Neural Network Intelligence, 2021.
[50] John Stuart Mill. Utilitarianism. Seven masterpieces of philosophy, pages 337-383. Routledge, 2016.
[51] Jooyoung Moon, Jihyo Kim, Younghak Shin, and Sangheum Hwang. Confidence-aware learning for deep neural networks. In ICML, 2020.
[52] Hongseok Namkoong and John Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. In NeurlPS, 2016.
[53] Radford M Neal. Bayesian learning for neural networks, volume 1 118. Springer Science & Business Media, 2012.
[54] David A Nix and Andreas Weigend. Estimating the mean and variance of the target probability distribution. In ICNN, 1994.
[55] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high performanoe deep learning library. In NeurIPS, 2019.
[56] Vihari Piratla, Praneeth Netrapalli, and Sunita Sarawagi. Focus on the common good: Group distributional robustness follows. ICLR, 2022.
[57] John Rawls. Justice as fairness: A restatement. Harvard University Press, 2001.
[58] Shiori Sagawa*, Pang Wei Koh*, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks. In ICLR, 2020.
[59] Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameteriizatoon exacerbates spurious correlations. In ICML, 2020.
[60] Yuge Shi, Jeffrey Seely, Philip HS Torr, N Siddharth, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. arXiv preprint arXiv:2104.099377 2021.
[61] Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log-likelihood function. Journal of statistical planning and inference, 90(2):227224, 2000.
[62] Jun Shu, Qi Xie, Lixuan Yi, Qian Zhao, Sanping Zhou, Zongben Xu, and Deyu Meng. Meta- weight-net: Learning an explicit mapping for sample weighting. In NeurIPS, 2019.
13

---

[63] Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pages 443-450. Springer, 2016.
[64] Mariya Toneva, Alessandro Sordoni, Remi Tachet des Combes, Adam Trischler, Yoshua Bengio, and Geoffrey J. Gordon. An empirical study of example forgetting during deep neural network learning. In ICLR, 2019.
[65] Vikas Verma, Alex Lamb, Christopher Beckham, Amir Najafi, Ioannis Mitliagkas, David Lopez- Paz, and Yoshua Bengio. Manifold mixup: Better representations by interpolating hidden states. In ICML, 2019.
[66] Junfeng Wen, Chun-Nam Yu, and Russell Greiner. Robust learning under uncertain test distributions: Relating covariate shift to model misspecification. ICML, 2014.
[67] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's 's transform- ers: State-of-the-art natural language processing. arXiv preprint 111.v:1910.0377,, 2019,
[68] Da Xu, Yuting Ye, and Chuanwei Ruan. Understanding the role of importance weighting for deep learning. In ICLR, 2021.
[69] Minghao Xu, Jian Zhang, Bingbing Ni, Teng Li, Chengjie Wang, Qi Tian, and Wenjun Zhang. Adversarial domain adaptation with domain mixup. In AAAI, 2020.
[70] Huaxiu Yao, Yu Wang, Sai Li, Linjun Zhang, Weixin Liang, James Zou, and Chelsea Finn. Improving out-of-distribution robustness via selective augmentation. In ICML, 2022.
[71] Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In CVPR, 2019.
[72] Runtian Zhai, Chen Dan, Zico Kolter, and Pradeep Ravikumar. Doro: Distributional and outlier robust optimization. In ICML, 2021.
[73] Runtian Zhai, Chen Dan, Zico Kolter, and Pradeep Ravikumar, Understanding why generalized reweighting does not improve over erm. arXiv preprint arXiv:2201.1229: 2022.
[74] Runtian Zhai, Chen Dan, Arun Suggala, Zico Kolter, and Pradeep Ravikumar. Boosted cvar classification. In NeurIPS, 2021.
[75] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. ICLR, 2018.
[76] Linjun Zhang, Zhun Deng, Kenji Kawaguchi, Amirata Ghorbani, and James Zou. How does mixup help with robustness and generalization? ICLR, 2021.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Sec. D in Appendix.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See Sec.
in Appendix.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full of assumptions of all theoretical results? [Yes] See Sec.
5
(b) Did you include complete proofs of theoretical results? [Yes] See Sec. in Appendix.
14

---

3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material as a URL)? [Yes] Code has been released.
(b) Did you specify all the training details (e.g.. data splits, hyperparameters, how they were chosen)? [Yes] See Sec. B in Appendix.
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [Yes] See Sec. [ in Appendix.
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)? [Yes] See Sec. B in Appendix.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any new assets either in the supplemental material or as URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [No] The datasets used are all publicly available datasets.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information offensive content? [No] The datasets used are all publicly available datasets.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A] We didn't 't conduct research with human subjects.
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A] We didn't t conduct research with human subjects.
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A] We didn't conduct research with human subjects.
15


Reviews:
[{"review_id": "sukZ66QsCY", "reviewer": "Reviewer_dN9V", "summary": "The paper proposes a new mixup method (UMIX) based on uncertainty estimation and importance weighting for subpopulation shift. The authors also prove theoretically that UMIX can achieve better generalization error bounds over IW methods without mixup. ", "questions": "\n1. While the idea proposed by the paper is straightforward, there are some missing details that are essential to justify the method’s effectiveness. For example, UMIX can be seen as a new data augmentation method/loss function, but how is the training procedure like? Is it similar as ERM, or is it adopted from some other algorithm. Adding this discussion would help to verify whether the comparisons are fair or not.\n\n2. The choices of parameters are missing and including them (such as T_{s} and T in Algorithm 2) would be nice.\n\n3. In Figure1, why is there hardly any samples with uncertainty from 0.3 to 0.8? Is it an effect of KDE, or is it due to the attributes of the datasets?\n\nMinor questions:\n\n1. Line59: contirbution $\\rightarrow$ contributions", "limitations": "The authors have thouroughly discussed the limitations and social impacts of their work in the appendix. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\nOriginality: Though the ideas of uncertainty estimation, importance weighting and mixup are not new, the combination of these ideas are novel and intriguing.\n\nQuality: The authors proved empirically and theoretically that their method is better than other baselines. The results are rather solid.\n\nClarity: The paper is well-written and easy to follow.\n\nSignificance: The paper provides a new method for augmenting mixup via uncertainty.\n\n**Weaknesses**\n\nI have not identified major weaknesses of this paper. However, I do have some concerns and listed them in the “Questions” part below.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "PQP9Ewfs2gk", "reviewer": "Reviewer_ptj2", "summary": "The paper extends the Mixup method by utilizing a dynamic uncertainty score as importance weight for Mixup samples. The uncertainty score (for a given x_i) is obtained by tracking the average prediction error of x_i through out the training trajectory, as an indication of how \"hard\" it is to learn this example.\n\nThe method is evaluated on 4 datasets with subpopulation shift. The proposed method outperforms others on \"worst case\" accuracy, but under-performs the baseline (ERM) on average accuracy.", "questions": "The uncertainty score u_i is updated every epoch, right? i.e. as the training epoch (T) progresses, you will get updated value of u_i(T). If so, please reflect that in Algorithm 1.\n\nIn Algorithm 2, you don't have to \"save\" all the predictions for every example at every epoch. Keeping a running average of $\\kappa(y_i, f_{\\theta_t}(x_i))$ over epoch t for each $x_i$ would be more memory efficient, right?", "limitations": "NA", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The overall writing is clear and easy to follow. However, I feel the originality and significance of the proposed method is thin. It is a straightforward extension of Mixup by combining importance weight and uncertainty score. The paper didn't explore or discuss other choices of this \"uncertainty score\" as importance weight. There's a whole literature on influence score (how influential is an example to the learning process), which can also play a similar role here.\n\nThe evaluation results are ok but not impressive. I'm not convinced that a better \"worst\" accuracy but degraded avg accuracy is a significant contribution, given all the extra compute required by UMix. Perhaps other reviewers could convince me how this is important and relevant to the research community.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "AyWDbSRAUBv", "reviewer": "Reviewer_Vm21", "summary": "The authors propose a importance reweighting scheme for mixup which promotes a fair training balance between subgroups in the data. The does not require group\nlabels during training.", "questions": "- Section 5 states \"Thus our bound can be tighter when\n315 the intrinsic dimension of data is small (i.e., rank(Σ) ≪ d).\" Is this really useful? Can you point to some intuition or concrete examples of when this would usefully lower the bound? When can we reasonably expect that to not be full rank? On the same note, is this referring to the dimension in the input space, the output space before the loss function, or something else entirely?\n\n- Are you actually MC sampling the saved SWAG models by building the Gaussian\nposterior throughout the training process in order to get the importance\nweights? OR are you just using the snapshots to get the model disagreement?\n\n- Section 3.1 says that the $N$ training examples are i.i.d. and sampled from\n$P$, but then immediately says that the subpopulation $g$ follows the\ndistribution $P_g$. Shouldn't the assumption then be that the dataset is not\ni.i.d. and only conditionally i.i.d. given $P_g$?", "limitations": "The authors have addressed both in the appendix. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n\n- The paper is well placed among recent works in both mixup, and importance reweighting methods.\n- The method performs well without group information during training which is impressive\n- The theoretical analysis provides a nice\n- Overall I think this is a solid work, but I some clarification on the issues raised in the following sections.\n\n# Weaknesses\n\n- L190 says that you assume a reasonable importance weight is linearly\npositively correlated to the corresponding uncertainty. What is a 'reasonable'\nimportance weight, and on what basis can this assumption be made? I realize\nthis is shown in Figure 1, but there should be some provided justification or\nreference to the fact that this will be empirically tested later.\n\n- How does UMIX perform without the group labels in the validation set? If the\nmethod is assumed to be oblivious to the group information, then it may not be\nreasonable to assume that they are always in the validation set.\n\n# Minor\n\n- L113: BNN's is proposed --> BNN's have been proposed.\n- L115: Performing ensembles --> ensembling them.\n- L122: upweight mixed samples thus can encourage --> upweight mixed samples, and thus can encourage.\n- L158: weights of UMIX are posed --> What does this mean?\n- L178: $p(theta; \\mathcal{D})$ --> shouldn't this be $p(\\theta | \\mathcal{D})$?\n- L180: with $T$ Monte Carlo Sampling --> with $T$ Monte Carlo samples\n- L193: in practice we could set $c = 1$ --> Does this mean that you do set $c=1$ in practice? It should be clearly stated if that is the case.\n- L295: proportion is --> proportion being\n- L296: in specific --> specifically\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "2wHsr1CHoxp", "reviewer": "Reviewer_cPPR", "summary": "The authors tackle the problem of subpopulation shift without subpopulation labels. They propose a method called uncertainty-aware mixup (UMIX), which augments regular mixup with importance weights computed based on predictive uncertainty. The authors benchmark their method against a large set of group aware and group unaware baselines on standard datasets, finding that their method shows competitive worst-case accuracy. Finally, the authors provide theoretical insight into why their method outperforms traditional importance weighting-based methods. ", "questions": "Please comment on points 2-6 from above. ", "limitations": "The authors adequately address the limitations in the appendix.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper is generally well-written and easy to follow.\n- The authors compare against a wide range of methods on the standard datasets, and report their results with standard deviations (in the appendix).\n- The method performs competitively relative to the baselines.\n\nWeaknesses:\n1. Though the method is novel, is it only an incremental change from what is proposed in JTT and LISA.\n2. The authors should provide more empirical justification for the uncertainty measure that they use. Though it is inspired by SWAG, it does seem quite different from SWAG. How does the measure of uncertainty compare with SWAG, and why can SWAG not be used directly in this method (by sampling from the posterior over weights)?\n3. It seems that the uncertainty measure takes both epistemic and aleatoric uncertainty into account. Would it be more ideal to weight based on epistemic uncertainty only, as some groups may inherently have more error than others?\n4. It would be interesting to explore how the uncertainty measure (Fig. 1) changes over time during the process of running Algorithm 2. Does this empirically support the intuition from Lines 195-198?  \n5. For Camelyon-17, I would suggest showing performance on the source domain test set along with the target domain results.\n6. It would be ideal if the authors could test on a couple more domain generalization datasets, as this seems relatively unexplored (only Camelyon-17), and many of the baseline methods are targeted towards the domain generalization setting. \n7. The authors should state the range of hyperparameters tuned, as well as the total tuning budget, in Appendix B3.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept