Mitigating Spurious Correlation via Distributionally Robust Learning with Hierarchical Ambiguity Sets
Authors: Sung Ho Jo
Seonghwi Kim
Minwoo Chae
Venue: ICLR 2026 Poster (2026)
Keywords: Spurious Correlation
Subpopulation Shift
Group Distributionally Robust Optimization
Wasserstein Distributionally Robust Optimization
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: Conventional supervised learning methods are often vulnerable to spurious correlations, particularly under distribution shifts in test data. To address this issue, several approaches, most notably Group DRO, have been developed. While these methods are highly robust to subpopulation or group shifts, they remain vulnerable to intra-group distributional shifts, which frequently occur in minority groups with limited samples. We propose a hierarchical extension of Group DRO that addresses both inter-group and intra-group uncertainties, providing robustness to distribution shifts at multiple levels. We also introduce new benchmark settings that simulate realistic minority group distribution shifts—an important yet previously underexplored challenge in spurious correlation research. Our method demonstrates strong robustness under these conditions—where existing robust learning methods consistently fail—while also achieving superior performance on standard benchmarks. These results highlight the importance of broadening the ambiguity set to better capture both inter-group and intra-group distributional uncertainties.

Published as a conference paper at ICLR 2026
MITIGATING SPURIOUS CORRELATION VIA DISTRIBU- TIONALLY ROBUST LEARNING WITH HIERARCHICAL AMBIGUITY SETS
Sung Ho Jo
Sconghwi Kim
Pohang University of Science and Technology 110ghhh81@@pstech.aae kr.
Pohang University of Science and Technology kahwi@postech.. ac kr
Minwoo Chae
Pohang University Science and Technology mchae@postech.al ac.kr .rr
ABSTRACT
Conventional supervised learning methods are often vuinerable to spurious cor- relations, particularly under distribution shifts in test data. To address this issue, several approaches, most notably Group DRO, have been developed. While these methods are highly robust to subpopulation or group shifts, they remain vulnera- ble to antra-group distributional shifts, which frequently occur in minority groups with limited samples. We propose a hierarchical extension of Group DRO that ad- dresses both inter-group and intra-group uncertainties, providing robustness to dis- tribution shifts at multiple levels. We also introduce new benchmark settings that simulate realistic minority group distribution shifts important yet previously underexplored challenge in spurious correlation research. Our method demon- strates strong robustness under these methods consistently fail-while also achieving superior performance on standard benchmarks. These results highlight the importance broadening the ambiguity set to better capture both inter-group and intra-group distributional uncertainties.
ceatttinns-herr existing robust learning
INTRODUCTION
In recent years, machine learning methods have achieved remarkable success across a wide range of applications. An important objective of many machine learning methods is learn model pa- rameters that minimize the population risk, which is the population expectation of the loss function. Given training data and model parameters, the population risk can be approximated by the empirical risk, defined as the sample-averaged loss. Therefore, model parameters can be learned by minimiz- ing the empirical risk, which is known as the empirical risk minimization (ERM) principle.
The underlying assumption of ERM-based methods is that the unseen future data, often referred to as test data, share the same distribution as the training data. However, many real-world problems, the test data may follow different distribution from the training data for various reasons. A notable example is subpopulation shift, where the training population consists of several groups (subpopula- tions), and the proportion of each group in the test data differs from that in the training data (Sagawa
et al. 2020 Cai et al. 2021 Yang et al. 2023).
In many instances of subpopulation shifts, the group indicator is spuriously correlated with the tar- get label or response variable. For example, in the widely studied Waterbirds dataset (Sagawa et al. 2020), the target label (e.g., "Waterbird" or "Landbirdd") is spuriously correlated with the background environment (e.g., water or land). As a result, FRM-hased models tend to associate "Waterbird" pri- marily with water backgrounds, leading significant performance degradation on minority groups, such as waterbirds appearing against land backgrounds. These vulnerabilities extend beyond con- trolled benchmarks, posing substantial risks in real-world, high-stakes domain such as healthcare
"Corresponding author.

---

Published as a conference paper at ICLR 2026
Q
N8 Ca
Q2
Pa
(a) Group DRO Ambiguity Set
(b) Proposed Ambiguity Set
Figure 1: Comparison of the Group DRO ambiguity set (a) and our hierarchical extension (b). While Group DRO restricts uncertainty to mixtures of group distributions, our approach introduces additional within-group uncertainty (indicated by red dashed arrows), offering robustness to both inter-group and intra-group distributional shifts. (For visualization, we assume the 3-dimensional space in the figure represents a probability space, where each point corresponds to probability distribution.)
Zech et al. 2018 Badgeley al. 2019). fairness (Hashimoto et al. 2018
Buolamwini & Gebru
2018 Obermeyer et al. 2019). and autonomous driving (Zhang et al. 2017
Over the past few years, a growing body of research has focused on mitigating these spurious corre- lations to ensure more reliable and stable model performance. Among the proposed methodologies, group distributionally robust optimization (Group DRO) (Sagawa et al. 2020) has emerged as foundational approach. By partitioning the data into predefined groups and optimizing for the worst group loss, Group DRO effectively minimizes the model's s reliance on spurious correlations tied to specific subsets of data. This framework has inspired wide range of subsequent methods, such as JTT (Liu et al.) (02I). SSA (Nam et al., 2022), PG-DRO (Ghosal Li, (023), DISC (Wu et al.] 2023) and GIC (Han & Zou 2024), which commonly employ two-step strategy: first identify- ing latent groups and then utilizing established robust training itself-to enhance model robustness.
faaechess-fruuulllly Group DRO
While these methods have significantly advanced the field, they primarily focus on minimizing the worst-group loss under the assumption that each group's training distribution reliably represents its true underlying distribution. However, this assumption often fails in practice-especiaaee esspccilll for minority groups with limited samples-where consequence of underrepresennaation in the training data (Ben-Tal & Nemirovski] 2002, Devroye al. 2013, Duchi Namkoong, 2021). This limitation highlights the need for more flexible and robust approach that accounts for uncertainty not only across groups, but also within them.
group distributional shifts naturally arise as
In this work, we address these limitations by introducing a hierarchical ambiguity set within the Group DRO framework, capturing both inter-group and intra-group uncertainties. As illustrated in Figure [ while conventional Group DRO focuses on robustness only to shifts in group proportions by minimizing the worst-group risk, approach extends this perspective by additionally modeling uncertainty within individual groups
Technically, employ nificant theoretical and empirical support for its efficacy in designing distributionally robust learning methods (Volpi al. 2018) Kuhn al. 2019 Blanchet et al.) 2022 Bai et al. (024). By defining a semantically meaningful cost function in a latent space, this formulation flexibly accommodates variations in the underlying data-generating mechanisms within each group. Consequently, our hi- erarchical ambiguity set enables the model to maintain robustness across a broader spectrum of distributional deviations, particularly for minority groups that are underrepresented in the training data.
bnsserstein-ddisance--boe
formulation, which has recently garnered sig-
Our main contributions are threefold:
We propose novel hierarchical ambiguity set that simultancously captures inter- group and intra- group uncertainties, constituting a fundamental advancement over previous methods built on Group DRO.
2

---

Published as a conference paper at ICLR 2026
We develop tractable minimax optimization algorithm that is computationally efficient, en- abling the practical realization of our DRO framework with a hierarchical ambiguity set.
We introduce evaluation settings standard benchmarks that reveal previously over- looked yet important failure mode: minority-grome distribution tuufts-eveed when simply alter- ing the train-test split without synthetic noise external shifts- where state-of-the-art methods collapse but our approach achieves consistently strong performance.
RELATED WORK
Using explicit group labels is well-established approach to achieving robust performance on un- derrepresented subpopulations. (Sagawa al. (2020) pioneered partitioning data by known group annotations and minimizing worst-group loss. Subsequent works extend this paradigm in multiple directions: Along similar lines, K et (2021) expands Group DRO's ambiguity set from convex to affine combinations of group distributions, although it still assumes fixed conditional distributions within each group. Meanwhile, Kirichenko et al. (2023) rebalances distributions via subsampling. Deng et al. (2024) employs a staged expansion of group-balanced set, and Yao et al.) 2022) uses Mixup-banc based augmentations (c.g.. CutMix, Manifold Mix) to learn more invariant fea- tures. Izmailov et al. (2022) shows that even standard ERM can yield robust feature representations when combined with selective reweighting, and Piratla et al. (2022) further leverages inter-group interactions identify shared features that enhance distributional robustmess.
In parallel, a growing body of work addresses scenarios where group annotations are unavailable or prohibitively expensive, motivating the need infer or approximate group structure. Building on Kirichenko et al. (2023). LaBonte et al. (024) extends last-layer retraining to settings with minimal or no group annotations. Some methods employ limited validation sets to discover latent groups and then apply Group DRO (Sagawa et al. 2020) for robust optimization (Nam et al. [2022] Ghosal & Li, 2023 Indeed, the model's loss (or its alternatives) often helps identify underrepresented subpop- ulations; for example, Nam et al. 2020); ; Liu et al. (202T) : Qiu et al. (2023) use high-loss samples to recognize minority groups. Other approaches draw on diverse cues: AAhed et al.) (202T): [Ceaa
ger et al. (2021) infer groups by maximizing violations of the invariant risk minimization principle (Arjovsky al. 2019), while Asgari et al. (2022) employs masking to reduce reliance on spurious features. Sohoni et al. (2020); (Sco et al. (2022) instead cluster feature embeddings to discover la-
tent groups and identify pscudo-attribute for debiasing. Likewise, [Zhang et al.) (2022) proposes a two-stage contrastive learning framework by aligning samples with the same class but different spu- rious attributes, and Wu et al. (2023) constructs a "concept bank" of candidate spurious attributes for robust partitioning. Another line of work identifies spurious features by comparing the train- ing set with a carefully selected reference dataset (Han & Zou 2024) or removes examples that disproportionately degrade worst-group accuracy (Jain et 2024)..
Beyond these established techniques, emerging efforts explore new failure modes within the spu- rious correlation framework. For instance, recent works have examined imperfect group partitions (hou et al. 2021). multiple spurious features (Paranjape et al. 2023 Kim et al.) 2024). low spu- rious signal strength (Mulchandani Kim, 2025), and spurious aware out-of-distribution detection (Zohrabi et al. 2025). Relatedly, recent work has shown that the presence of spurious correlations can itself constitute a failure mode for domain adaptation methods (Kim et al., (026). In this con- text, our work addresses critical yet underexplored issue: minority-group data, by nature extremely limited in size, may not faithfully reflect their underlying distributions-a natural but overlooked phenomenon that must be explicitly accounted for.
3 PRELIMINARIES
Problem Setup. We consider a supervised learning problem where each observation consists of input features X E X' and a label Y E y. Let {(x,, y be the training data and θ be the parameter space. Assuming that the test data follows the same distribution as the training data, an important goal of various machine learning methods is to minimize the risk (also referred to as the test or generalization error)
Ep [p[c( [x(((((X)))))
3

---

Published as a conference paper at ICLR 2026
over O where f is function parametrized by 8, L is standard loss function, such the cross- entropy, and Ep denotes the expectation under the population distribution P. To achieve this, one may solve the following optimization problem:
inf AEB
E Eple(0: (X,
E
e(0: (2t,
often referred to as the ERM problem, where P is the empirical measure and E(0; (X, Y)) is short- hand for (X),,
In addition to the above basic setting, we assume that the data are partitioned into multiple groups, with an indicator variable G & G; thus, the training data can expressed as {[r,, Mhi gi (:)))))) )== In particular, assume that this indicator variable is spuriously correlated with the label Y. More specifically, in all our examples, we assume the existence of spurious attribute A, and that the group variable G (Y, A)) is pair involving the response variable Y. Hence, g y A. With adjusted notation, we write G [I,..,,
In the Waterbirds dataset, for example, the task is classify birds Waterbird Landbird, with a spurious attribute being the background type (Water Background or Land Background), which results in four distinct groups. This dataset provides classic example of spurious correlation: most waterbirds (Y Waterbird) are found against water backgrounds (A = Water Background), leading models to rely excessively background features. This overreliance substantially diminishes the performance of ERM on underrepresented groups, such as waterbirds on land backgrounds.
Group DRO. Group DRO (Sagawa et al. 2020) was devised to address the aforementioned issue caused by the spurious attribute. In the Group DRO, the training data are modeled instances from a mixture distribution P = Σø= agPa. where P. denotes the conditional distribution of (X, Y) given = 9. and a = (a),.... am represents the mixing proportion. Thus, each group forms a subpopulation within the training data. Instead of minimizing the population risk (I). Group DRO aims to minimize
inf max Ep, [E(0; ()), GE@ gEğ
which corresponds the risk of the worst-performing group. Consequently, this procedure is highly robust to the subpopulation shift described in the introduction.
The Group DRO formulation involves optimization the discrete group variable g. posing a computational challenge for practical use. Sagawa et al. 2020) showed that the problem can be reformulated into an equivalent form
inf AEO
sup BE BEAm-1
Σ By [((0: (X, (X,))
g=1
that involves continuous variables only, where A = (3 : 3, 2 0, Em simplex.
g= Bg = =1] is the (m 1)-
Group DRO can be understood as instance of standard DRO (Ben- Tal al. 2013 Duchi al.
2021 with a specific ambiguity set. Note that the standard DRO formulation is given as
inf sup EQ [C(B: (X,Y))..
6Ee QEo
where Q is a class of distributions, commonly referred to as the ambiguity (or uncertainty) set. The Group DRO formulation 2 specific case of DRO with
Y12
BgP BE BEA
Note that in frequently used DRO frameworks, the ambiguity set Q is often defined as small neighborhood with respect to standard (pseudo- )eetiic, such as the Wasserstein distance (Moha- jerin Esfahani & Kuhn 2018 Gao et al. 2024) and f-divergence (Namkoong & Duchil 2016) Mehta al. (0044).

---

Published as a conference paper at ICLR 2026
4 PROPOSED METHOD
4.1 HIERARCHICAL AMBIGUITY SETS
In this section, we propose a hierarchical extension of Group DRO to be robust to distribution shifts multiple levels. The proposed method is devised to capture both inter-group and intra-group uncertainties in modeling the distributional shifts.
High-level Formulation. As in Group DRO, we model the training distribution a mixture of the form P = --nn1 ag Pg. To model the distributional uncertainty, we consider the DRO formulation 3) with a hierarchical ambiguity set 2, defined as
d1 a) sp. d2(Qgr Pg) S Cy vg
Q
8gOg
(5)
where di and d2 are suitable metrics on A and the class of distributions for (X, Y), respectively, and p. Cg > 0 are radii that determine the size of the ambiguity set.
The ambiguity set Q has two-level hierarchy. The first level is controlled by the mixing proportion 3. It accounts for uncertainty in the proportion of each subpopulation or group. Such uncertainty can arise, for example, if certain minority groups appear more frequently in evaluation settings than in the training set, thereby increasing their probability of occurrence and potentially amplifying spurious correlations if not properly addressed (Sugiyama & Storkey (006). At the second level, the distributional shift in each group is considered to capture within-group variability.
By jointly accounting for changes in the group proportion and the conditional distributions {Po (Pg)g= the proposed framework provides two levels of robustness: inter-group generalization and resilience intra- intra-group group variability. This dual modeling of real-world uncertainties enables the proposed method to address a broader range distributional shifts compared to Group DRO ( alone standard DRO B), which uses standard (pseudo-) metrr neighborhood as its ambiguity set.
Relationship to Group DRO and Standard DRO.
The ambiguity set 4 used in Group DRO is special case of the proposed ambiguity set (5). particular, (4) can be obtained by setting p = co and Cg 0.
While standard DRO, which uses standard metric neighborhood as an ambiguity set, can also be understood as special case of the proposed method, the philosophy of the proposed ambiguity set differs from that of the standard ones. In standard DRO, the ambiguity set is taken as small neighborhood with respect to a standard (pseudo-) metric.. In contrast, we allow large value for p. the radius that determines robustness to group proportions. Hence, distributions that are far from P with respect to standard metrics can also belong to the ambiguity set (
The choice of dj is not critical because most reasonable metrics Am- lead to similar behavior in our formulation. The constraint d, (3, o)) S p controls how much the test- time group proportions 8 are allowed to deviate from the training mixture a. When prior knowledge suggests that the test-time group proportions will not depart substantially from those observed in training (for example, they are expected to remain within a small neighborhood of a), one may choose a finite p to restrict attention to such mixtures. In the spurious-correlation correlation setting we study, however, it is standard assume no reliable prior on test-time group proportions and evaluate robustness via worst-group accuracy, which can be viewed as examining performance under the extreme mixture that concentrates all probability mass on the worst-performing group. For this reason, in the remainder of this paper we instantiate our framework with p
Detailed Formulation.
For d2, we use a Wasserstein distance. Other options such as f-divergences (e.g.. KL divergence or total variation distance) are in principle also possible, but they induce ambiguity sets with structural limitation
our setting. Ambiguity sets defined via an f-divergence require candidate distributions to be absolutely continuous with respect the empirical distribution; since the empir- ical distribution has finite support, such sets can only reweight the observed training samples. In contrast, Wasserstein distances do not impose absolute continuity and therefore allow support shifts,
5

---

Published as a conference paper at ICLR 2026
enabling the ambiguity set include meaningful variations that may appear in the training data. This choice is thus better aligned with the atra-group distributional shifts we aim to capture.
Among Wasserstein distances, consider the onfinite-order Wasserstein distance for computational convenience. Recall the definition the Wasserstein distance of order p E [1, 00):
P) =inf inf { { (((x)))) (el (d.1)) ((
W P)
where the infimum is taken over every coupling 7 of Q and P. and c(.,)) :) is a cost function. The infinite-order Wasserstein distance defined as W 00 (Q, P) supp21 W (Q, (Q,P), P), with variational representation
Wo (Q, = inf O Q(A) P(A)
(6)
for every Borel set A
where A denotes the e-enlargement A: see Givens Shortt (1984).
The cost function is defined in a latent semantic space, which is more effective than defining it in the space of raw data (Zeiler & Fergus 2014, Krizhevsky et al. 2017). Specifically, we employ
deep neural network f of depth L, defindd as
fz))) = f2 (fh-1 (x).
and take the output of the (L 1)-th layer (before the final fully connected layer) the semantic representation:
z(x) == fh. (f2 (...
We then define the cost function c(:,,) as
=(x')ll,
yf
c((x, (2,,
(00
otherwise.
Note that under definition, QP, oo if the marginals P and Q of Y differ. In all applications, the group indicator G is defined as a pair A); hence, this definition does not cause any issues.
Proposed Hierarchical DRO Formulation. To sum up, the proposed hierarchical DRO can be written the standard form 3 with the ambiguity set
BE 2 B,Q
(8)
(o, 6g
The flexibility in 3 E A allows the group proportion to differ from a, which enables adaptation to new changing subpopulation frequencies without introducing entirely new groups. With the
constraint (Q, Pg) I Cg. we accommodate plausible instance-level shifts within each group. Leveraging the semantic cost c(-,-) allows the model to capture meaningful perturbations in high- dimensional feature spaces without conflating different class labels.
4.2 ALGORITHM
In this subsection, we provide an algorithm to solve the proposed DRO with the ambiguity set (8). We set cg = c/ /īg. where ng is the size of group in the training data, and e is tunable hy- perparameter that controls the degree robustness to within group distributional shifts. Intuitively, the fewer samples a group has, the more cautiously its potential distributional variations must be accounted for. The choice of is a common challenge across robust learning methods, particularly DRO-based approaches. While our framework remains effective across a broad range of values of we propose a simple yet effective heuristic for its selection to further enhance the practicality of framework. Further details the selection of c provided in Appendix G
Due to the hierarchical structure of the ambiguity set in g). solving the resulting DRO problem is not straightforward. We therefore begin by reformulating the hierarchical DRO into tractable optimization problem. We formally state the resulting formulation in the following theorem. The proof is provided in Appendix

---

Published as a conference paper at ICLR 2026
Theorem 4.1. Let Q be the ambiguity defined in (8. robust optimization problem
Then, the corresponding distributionally
inf sup Eg
[e(; (X, X))
BE QEQ
upper-bounded by the following surrogate objective:
=' [(B(E))))
inf AEO
sup
Br
sup
BEAm-1
Intuitively, Theorem 4.1 shows that the worst-case risk over our hierarchical ambiguity set can be conservatively wer-approximmted by an adversarial perturbation problem in the latent space, where the inner maximization is weighted by the worst-case group proportions B.
Remark (Tightness of the relaxation). mization from
The upper bound in (9) relics on relaxing the inner maxi-
Ep,
((f(((())) I Ep, (x))
sup
sup
((2(6)))))
This relaxation can be striet in general when the feature map z(-) is not surjective onto the eg-ball. However, expect this effect be negligible in practice for embedding distributions arising in modern deep networks. It is commonly assumed that the pushforward distribution of X under z(-) possesses positive Lebesgue density on, or in neighborhood of, its effective support.
Although raw images may concentrate near a low-dimensional manifold in input space and thus typically do not admit a Lebesgue density, the situation is different after they pass through several nonlinear layers. Empirically and conceptually, internal representations produced by deep networks behave "thickened" manifolds and are typically modeled as distributions with nonzero volume in the ambient feature space.
Formally proving the existence such a density for arbitrary deep embeddings is challenging. Nevertheless, under this standard modeling assumption, the relaxation above becomes nearly tight. The overlap between the baall and the image of =(-) is substantial, making the resulting upper bound in (9) close approximation of the original hierarchical objective.
We therefore minimize the surrogate objective
9 via coordinate-wise procedure, as detailed next.
Proposed Iterative Training Procedure.
For given 6, let 2/ denote the maximizer of the map 4 ((f vs)
over the set {2' || - 2(>;)]| I ). To solve the optimization problem P. we iteratively update 3, A and semantic variables z! coordinate-wis. provided in Algorithm
below. A pseudo-code for minibatch size of is
Update of For given 0, 2 can be approximated by one-step projected gradient ascent, ensuring that || - =xx:)]l I <s. (Lines 6-8)
2. Update of 3. For given 0 and z/. 8 can be computed using exponentiated gradient ascent, variant of mirror descent with negative Shannon entropy (Sagawa et al. (220)). (Lines 10-12)
3. Update of 0. For given and z/G, we update 0 using stochastic gradient descent. (Line 13)
convergence guarantee under convexity assumptions established in Appendix B showing that the algorithm achieves an ((//VT) convergence rate.
EXPERIMENTS
5.1 DATASET
We conduct experiments on three widely used benchmark datasets, CMNIST, Waterbirds, and CelebA, each exhibiting known spurious correlations between the label and irrelevant attribute. All datasets include minority group that is underrepresentee,, rendering them susceptible to distri- butional shifts.

---

Published as a conference paper at ICLR 2026
Algorithm DRO with Hierarchical Ambiguity Set
Input: Step sizes 1/8, Tjo, Vls initial parameters 0 310; number of iterations
2: for t= to T do
3: Sample g Uniform(1.....))
4: Sample (x) ~ Pg
Initialize 2' ((x)
6:
x' - z' + 77 V (JE (x').))
7:
if || [ Eg then
Proj ( (z)
9:
end if
10:
Update 8' + B(--1)
11: Update 3y - 8y exp (1 [1= (1 (z))
12:
Normalize ((0) + 81 8.
13: Update ((0) 0(--1) 99 (I (=)..))
end for
Original Datasets.
CMNIST (Arjovsky et al., 2019): A colored variant of MNIST, split into four groups based on digit label (digits 4-4 as label 0, and digits 5-9 label 1) and color (red vs. green). The color is spuriously correlated with the digit label in the training set.
Waterbirds (Sagawa et al. 2020): Created by combining bird images from CUB (Wah et al., 2011) with backgrounds from Places (Zhou et al.]| (2017). yielding four groups based on (bird rype, buckground). The minority group (waterbird, land background) typically has few samples.
CelebA (Liu al., 2015): A facial attribute dataset used here for classifying blond blond hair, where gender acts as a spurious attribute. The minority group (blond hair, male) is significantly underrepresented.
Modified Datasets with Minority Group Shifts. To rigorously test our approach under more realistic distribution shifts, we construct modified versions of the above datasets by inducing intra- group shifts specifically in each minority group:
Shifted CMNIST: Rotate all images in the minority group (label red) by 909 at test time, while
keeping them unrotated training time.
Shifted Waterbirds: Restrict the training set's minority group (waterbird, land background) to only waterfowls, and the test set's 's minority group to only seabirds.
Shifted CelebA: For minority group (blond hair. male), include only no- no-glasses images in training and only with glasses images at test time.
These modifications reflect real-world scenarios where underrepresented groups not only appear more frequently but also exhibit subtle changes. Further details and illustrative examples are pro- vided in Appendix D
5.2
BASELINES
We compare our method to several representative baselines: ERM, Group DRO (Sagawa et al. 2020 JTT (Liu al., 202T), CnC (Zhang al.) a (022), SSA (Nam et 2022). LISA Yao al.
2022 DFR (Kirichenko et al. 2023), PDE (Deng et al. (2044, and GIC (Han & Zou 2024). These methods range from direct robust learning (e.g.. Group DRO) to two-step pipelines that first
fer group membership and then apply robust training (e.g.. SSA, GIC). Detailed descriptions provided in Appendix
For our newly constructed datasets incorporating minority-group distribution shifts, conducted experiments focusing on Group DRO, LISA, DFR, and PDE. Unlike methods that infer group la- bels and then rely separate robust training step. these four baselines like proposed ap- proach- proach-directly utilize known group information. This distinction provides a more consistent and fair comparison in scenarios where explicit group labels available.
8

---

Published as a conference paper at ICLR 2026
Table 1: gorst-group and average accuracy on CMNIST, Waterbirds, and CelebA under shifted distributions. All results are averaged over three runs with different random seeds. Boldface indicates the best performance, while underlined numbers denote the second-best.
Shifted CMNIST
Shifted Waterbirds
Shifted CelebA
Method
Group label
Worst Acc
Average Acc
Worst Acc
Average Acc
Worst Acc
Average Acc
DRO
65.9
74.0107
91.7.103
94.9
59.8132
92,4:103
LISA
V
42.9=10.0
59,81445
79.1
94.2.003
60.6111
92.1
DFR
28.0a44
47.8
89.2
96.310.4
50.3
90.5
PDE
65.3
71.3262
84.4
92.0
56,3a112
91.6404
Ours
71.8+22
75,0+04
93.7.402
94.6
72.1+20
91.340.1
Table 2: Worst-group and average accuracy CMNIST, Waterbirds, and CelehA under their original (un- shifted) distributions.
CMNIST
Waterbirds
CelebA
Method
Group label
Worst Ace
Average Acc
Worst Acc
Average Acc
Worst Ace
Average Acc
ERM
X
3,4
12.9108
62.6103
97.3
47.7
94,9103
JTT
X
67.3
76.4233
83.8+12
89.3
81.5
88.1 18.
CnC
X
88.5 88.5
90.9
88. 899
89. 89.9205
GIC
72.2
73.22:02
86.3
89.621.
89.440a
91.9201
SSA
71.1 0..4
75,0+-03
89.0#06
92.2#0.9
89.8+11
92,8+0.1
GroupDRO
73. +.03
74.8++02
90.6+02
92.7+0.1
89.3+13
92,6403
LISA
73.3
74.0+=a.
89.2+ao
91.8=0.
89,3+1.
92,4+04
DFR"
59.8
62.1
90.2+08
97.0=0.
80.7++4
90.6407
PDE
72.6+0.7
73.0+04
90.3+03
92.4
91.0+04
92.0#06
Ours
V
73.6
75.1
90.8
92.6=0.2
90.4+au
92.7400
5.3 EVALUATION
Metrics. We consider two metrics: worst-group accuracy and average accuracy. The worst-group accuracy is obtained by evaluating accuracy on each group and taking the minimum across all groups, providing insight into how a method performs if the test distribution is heavily skewed to- ward the most challenging subgroup. Meanwhile, the average accuracy is computed as the weighted average of group accuracy, where the weights are proportional to the group sizes in the training data, reflecting overall performance but offering less visibility into group-specific disparities.
Model Selection. Following Sagawa et al. (2020) and related methods, we select hyperparameters and stopping criteria based on the highest worst-group validation accuracy. In particular, for scenar- ios involving minority group shifts, we adopt the data-driven tuning procedure from Appendix G determine the perturbation parameter
5.4 RESULTS
Performance on Shifted Distributions. Under shifted distributions (Table I.. our method demon- strates clear superiority worst-group accuracy across all three benchmarks (CMNIST. Waterbirds, and CelebA). On CMNIST, LISA and DFR degrade substantially, highlighting their vulnerability to intra-group shifts. By contrast, our framework maintains high worst-group accuracy.
For Waterbirds, which involves moderate shift in species composition within the minority group, most baselines experience notable drops in worst-group accuracy. In contrast, our approach main- tains robust worst-group accuracy, indicating its capacity to adapt to intra-group variability. Interest- ingly, both Group DRO and our method report higher worst-group accuracy in the shifted case than in the original Waterbirds dataset (Table 2); however, this discrepancy arises from known misla- beling issue (Asgari et al., 2022). where three bird species labeled as "waterbirr" should actually be "landbird." " To verify this, correct the mislabeled samples in the original dataset and report results

---

Published as a conference paper at ICLR 2026
in Table under the corrected labels, both Group DRO and our method exhibit the expected pattern, performing better in the unshifted setting than under the minority- group shift. Notably, our approach outperforms Group DRO in both scenarios, confirming its robustmess even after label corrections.
On the more challenging CelebA bench- mark, advantage grows more pro- nounced. While PDE shows slightly higher worst-group accuracy on the orig- inal dataset (Table 2). its performance drops sharply (by about 34.7%) when the minority-grus group distribution is shifted. These observations underscore the im- portance of modeling both inter-group and intra-group eepeetaiintessesseiilllll given that minority groups in Waterbirds and CelebA constitute only about 1% of the data and thus are more susceptible to distributional changes. Furthermore, results highlight that relying on pre-defined test splits with uniformly distributed attributes may offer an overly optimistic view of real-world robustness.
Table 3:
Worst-group and average accuracy Waterbirds under minority group shifted distributions and Corrected Wa- terbirds on the original dataset with corrected labels.
Shifted Waterbirds
Corrected Waterbirds
Method
Worst
Avg Acc
Worst Ace
Avg Acc
Group DRO Ours
91.710.3
94,920.1
1006
94.7+000 96.3+0000
93.7+00.2
94.6+0.1
95.1=004
Performance on Original Distributions. On the original (unshifted) versions of CMNIST, Wa- terbirds, and CelebA (Table 2) our method consistently achieves top-tier worst-group accuracy. It secures the highest or near-highest scores across all three benchmarks, confirming that the proposed framework not only excels under distributional shifts but also remains effective when gropu distributions are stable. Notably, even in these unshifted settings, methods such as Group DRO rely on the strong assumption that each group's 's training distribution remains valid at test time. As results show, explicitly modeling distributional uncertainty within minority groups can yield more reliable robustness, highlighting the limitations of approaches that treat group distributions as fixed. By addressing potential discrepancies at both the inter-groun group and intra-group levels, our framework provides a stronger foundation for real-world applications.
CONCLUSION
Summary. We introduce novel distributionally robust optimization framework with a hierar- chical ambiguity set that explicitly models both inter-group and intra-group distribution shifts-a previously overlooked but practically crucial scemario for underrepresented subpopulations. We re- formulate the resulting DRO into tractable and computationally efficient optimization problem, making it feasible for practical use. Our analysis shows that even modest changes in how minority partitioned between training and testing can severely degrade the performance
group samples of existing robust methods. In contrast, our approach consistently maintains strong performance on both standard and modified benchmarks.
Outlook. While experiments foeus image datasets with explicit feature labels that make intra-group shifts observable, extending our framework to domains such as text-where such labels difficult to obtain- represeenss interesting avenue for future work. More broadly, we hope this work motivates future efforts to systematically uncover and address overlooked failure modes in robust learning.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. RS-2023-00240861), a Korea Institute for Advancement of Technology (KIAT) grant funded by the Korea Government (MOTIE) (RS- (RS-2024004409992, 2024 HRD Program for Industrial Innovation), and Institute of Information & Communications Technol- ogy Planning & Exaluation((ITP) Data-X Leader HRD program grant funded by the Korea government (MSIT)
REFERENCES
Faruk Ahmed, Yoshua Bengio, Harm Van Seijen, and Aaron Courville. Systematic generalisation with group invariant predictions. Proc. International Conference on Learning Representations, 2021.
M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Pazz Invariant risk minimization. arXiv preprint arXiv: 1907.02893, 2019.
Saeid Asgari, Aliasghar Khani, Fereshte Khani, Ali Gholami, Linh Tran, Ali Mahdavi Amiri, and Ghassan Hamarneh. Masktune: Mitigating spurious correlations by forcing to explore. In Proc. Advances in Neural Information Processing Systems, 2022.
Mareus Badgeley, John R Zech, Luke Rayner, Benjamin Glicksberg, Manway Liu, William Gale, Michael V McConnell, Bethany Percha, Thomas M Snyder, and Joel T Dudley. Deep learning predicts hip fracture using confounding patient and healthcare variables. NPJ Digital Medicine, 2(1):31, 2019.
Xingjian Bai, Guangyi He, Yifan Jiang, and Jan Obloj. Wasserstein distributional robustness of neural networks. Proc. Advances in Neural Information Processing Systems, 2024.
A. Ben- Ben-Tal, D. Den Hertog, A. De Waegenaere, B. Melenberg, and G. Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 2013.
Aharon Ben-Tal and Arkadi Nemirovski. Robust optimizatiion-mettooollgy and applications. Math-
ematical Programming, 92:453 480, 2002.
Jose Blanchet, Lin Chen, and Xun Yu Zhou. Distributionally robust mean-variane portfolio selce- tion with wasserstein distances. Management Science, 68(9):6382-641 2022.
Joy Buolamwini and Timnit Gebru. Gender shades: Intersectional accuracy disparities in commer- cial gender classification. Proc. Conference on Fairness, Accountability and Transparency, pp- 77-91, 2018.
Tianle Cai, Ruiqi Gao, Jason Lee, and Qi Lei. A theory of label propagation for subpopulation shift. In Proc. International Conference on Machine Learning, pp. 1170-1182, 2021.
Elliot Creager, Jörn-Henrik Jacobsen, and Richard Zemel. Environment inference for invariant learning. In Proc. International Conference Machine Learning, pp. 2189-2200, 2021.
Yihe Deng, Yu Yang, Baharan Mirzasoleiman, and Quanquan Gu. Robust learning with progressive data expansion against spurious correlation. In Proc. Advances in Neural Information Processing Systems, 2024.
Lue Devroye, László Györfi, and Gábor Lugosi. A Probabilistic Theory of Pattern Recognition. Springer Science & Business Media, 2013.
C. Duchi, P. W. Glynn, and H. Namkoong. Statistics of robust optimization: A generalized empirical likelihood approach. Mathematics of Operations Research, 2021.
John C Duchi and Hongseok Namkoong. Learning models with uniform performance via distribu- tionally robust optimization. Annals of Statistics, 4978).137811466, 2021.
Rui Gao, Xi Chen, and Anton Kleywegt. Wasserstein distributionally robust optimization and variation regularization. Operations Research, 72(3):1177-1191 2024.

---

Published as a conference paper at ICLR 2026
Soumya Suvra Ghosal and Yixuan Li. Distributionally robust optimization with probabilistic group. In Proc. AAAI Conference on Artificial Intelligence, pp. 11809-11817, 2023.
Clark R Givens and Rae Michael Shortt. A class of wasserstein metries for probability distributions. The Michigan Mathematical Journal, 31(2):231-240, 1984.
Yujin Han and Difan Zou. Improving group robustness on spurious correlation requires preciser group inference. In Proc. International Conference Machine Learning, pp. 17480-17504, 2024.
Tatsunori Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang. Fairness with- out demographics in repeated loss minimization. In Proc. International Conference on Machine Learning, pp. 1929-1938, 2018.
Pavel Izmailov, Polina Kirichenko, Nate Gruver, and Andrew G Wilson. On feature learning in the presence of spurious correlations. In Proc. Advances in Neural Information Processing Systems, 2022.
Saachi Jain, Kimia Hamidieh, Kristian Georgiev, Andrew Ilyas, Marzyeh Ghassemi, and Aleksander Madry. Improving subgroup robustness via data selection. In Proc. Advances in Neural Informa- fion Processing Systems, 2024.
Nayeong Kim, Juwon Kang, Sungsoo Ahn, Jungseul Ok, and Suha Kwak. Improving robustness to multiple spurious correlations by multi-objective optimization. In Proc. International Conference on Machine Learning, pp. 24040-24058, 2024.
Seonghwi Kim, Sung Ho Jo, Wooseok Ha, and Minwoo Chae. Distributionally robust classification for multi-source unsupervised domain adaptation. In Proc. International Conference on Learning Representations, 2026.
Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations. In Proc. International Conference on Learning Represen- tations, 2023.
W. Koh, S. Sagawa, H. Marklund, S. M. Xie, M. Zhang, A. Balsubramani, W. Hu, M. Yasunaga, R. L. Phillips, Gao, al. Wilds: A benchmark of wn-the-wild distribution shifts. In Proc. International Conference on Machine Learning, pp. 5637-5664, 2021.
A. Krizhevsky, I. Sutskever, and G. Hinton. Imagenet classification with deep convolutional neural networks. Communications of the ACM, 60(6):8499, 2017.
David Krueger, Ethan Caballero, Zhang, Remi Le Priol, and Aaron Courville. Out-of-distributie generalization via risk extrapo- lation (rex). In Proc. International Conference on Machine Learning, pp. 5815-5826, 2021.
Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai
Daniel Kuhn, Peyman Mohajerin Esfahani, Viet Anh Nguyen, and Soroosh Selfieezadeh-Abadeee Wasserstein distributionally robust optimization: Theory and applications in machine learning. In Operations Research & Management Science in the Age of Analytics, pp. 130-166. Informs, 2019.
Tyler LaBonte, Vidya Muthukumar, and Abhishek Kumar. Towards last- retraining for group robustness with fewer annotations. In Proc. Advances in Neural Information Processing Systems, 2024.
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 Proc. International Conference on Machine Learning, pp. 6781-6792, 2021.
Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In Proc. Interna- fional Conference on Computer Vision, pp. 3730-3738, 2015.
12

---

Published as a conference paper at ICLR 2026
Ronak Mehta, Vincent Roulet, Krishna Pillutla, and Zaid Harchaoui. Distributionally robust opti- mization with bias and variance reduction. In Proc. International Conference on Learning Rep- resentations, 2024.
Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the wasserstein metric: Performance guarantees and tractable reformulations. Mathematical Progrumming, (71(1):11- 2018.
Varun Mulchandani and Jung-Eun Kim. Severing spurious correlations with data pruning. In Proc. International Conference Learning Representations, 2025.
Junhyun Nam, Hyuntak Cha, Sungsoo Ahn, Jaeho Lee, and Jinwoo Shin. Learning from failure: De-biasing classifier from biased classifier. In Proc. Advances in Neural Information Processing Systems, 2020.
Junhyun Nam, Jachyung Kim, Jacho Lee, and Jinwoo Shin. Spread spurious attribute: Improving worst-group accuracy with spurious attribute estimation. In Proc. International Conference on Learning Representations, 2022.
Hongseok Namkoong and John C Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. In Proc. Advances in Neural Information Processing Systems, 2016.
Arkadi Nemirovski, Anatoli Juditsky, Guanghui Lan, and Alexander Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal optimization, 19744:5774 1609, 2009.
Ziad Obermeyer, Brian Powers, Christine Vogeli, and Sendhil Mullainathan. Dissecting racial bias in an algorithm used to manage the health of populations. Science, 366(6464):447 453, 2019.
Bhargavi Paranjape, Pradeep Dasigi, Vivek Srikumar, Luke Zettlemoyer, and Hannanch Hajishirzi. AGRO: Adversarial discovery of error-prone groups for robust optimization. In Proc. Interna- tional Conference on Learning Representations, 2023.
Vihari Piratla, Praneeth Netrapalli, and Sunita Sarawagi. Focus on the common good: Group dis- tributional robustness follows. In Proc. International Conference on Learning Representations, 2022.
Shikai Qiu, Andres Potapczynski, Pavel Izmailov, and Andrew Gordon Wilson. Simple and fast group robustness by automatic feature reweighting. In Proc. International Conference on Machine Learning, pp. 28448-28467 2023.
S. Sagawa, W. Koh, T. B. Hashimoto, and P. Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In Proc. International Conference on Learning Representations, 2020.
Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based local- ization. In Proc. International Conference on Computer Vision, pp. (18-626,77 2017.
Seonguk Seo, Joon- Lee, and Bohyung Han. Unsupervised learning of debiased representa- tions with pssudo-attributes. In Proc. Conference on Computer Vision and Pattern Recognition, pp. 16742-16751, 2022.
Nimit Sohoni, Jared Dunnmon, Geoffrey Angus, Albert Gu, and Christopher Ré. No subclass left behind: Fine-grained robustness in coarse-grained classification problems. In Proc. Advances in Neural Information Processing Systems, 2020.
M. Staib and S. Jegelka. Distributionally robust deep learning as a generalization of adversarial training. In NeurlPS Workshop on Machine Learning and Computer Security, 2017.
M. Sugiyama and A. J. Storkey. Mixture regression for covariate shift. In Proc. Advances Neural Information Processing Systems, 2006.
13

---

Published as a conference paper at ICLR 2026
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. Journal of Machine Learning Research, (100):2579-2005. 2008.
Riccardo Volpi, Hongseok Namkoong, Ozan Sener, John C Duchi, Vittorio Murino, and Silvio Savarese. Generalizing unseen domains via adversarial data augmentation. In Proc. Advances in Neural Information Processing Systems, 2018.
C. Wah, S. Branson, P. Welinder, P. Perona, and S. Belongie. The caltech-ucsd birds-200-2011 dataset. Technical Report CNS-TR-2011- 1-001, California Institute of Technology. 2011.
Shirley Wu, Mert Yuksekgonul, Linjun Zhang, and James Zou. Discover and cure: Concept-aware mitigation of spurious correlation. In Proc. International Conference Machine Learning, pp- 37765-37786, 2023.
Yuzhe Yang, Haoran Zhang, Dina Katabi, and Marzych Ghassemi. Change is hard: A closer look at subpopulation shift. In Proc. International Conference on Machine Learning, pp. 39584-39622, 2023.
Huaxiu Yao, Yu Wang, Sai Li, Linjun Zhang, Weixin Liang, James Zou, and Chelsea Finn. Improv- ing out-of-disttibution robustness via selective augmentation. Proc. International Conference on Machine Learning, pp. 25407-25437, 2022.
John R Zech, Marcus A Badgeley, Manway Liu, Anthony B Costa, Joseph J Titano, and Eric Karl Oermann. Variable generalization performance of a deep learning model to detect pneumonia in chest radiographs: A cross-sectional study. PLoS Medicine, 15(1]):e 1002683, 2018
M. Zeiler and Fergus. Visualizing and understanding convolutional networks. In Proc. Euro- pean Conference on Computer Vision, pp. 818-833 833, 2014.
Michael Zhang, Nimit S Sohoni, Hongyang R Zhang, Chelsea Finn, and Christopher Re. Correct- N-Contrast: contrastive approach for improving robustness to spurious correlations. In Proc. International Conference Machine Learning, pp.
2022.
Yang Zhang, Philip David, and Boqing Gong. Curriculum domain adaptation for semantic segmen- tation of urban scenes. In Proc. International Conference on Computer Vision, pp. 2020-2030, 2017.
B. Zhou, A. Lapedriza, A. Khosla, Oliva, and Torralba. Places: 10 million image database for scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, pp- 1452-1464, 2017.
Chunting Zhou, Xuezhe Ma, Paul Michel, and Graham Neubig. Examining and combating spurious features under distribution shift. In Proc. International Conference Machine Learning, pp- 12857-12867, 2021.
Reihanch Zohrabi, Hosein Hasani, Mahdich Soleymani Baghshah, Anna Rohrbach, Marcus Rohrbach, and Mohammad Hossein Rohban. Spurious-aware prototype refinement for reliable out-of-distribution detection. arXiv preprint arXiv:2506.2388// 2025.
14

---

Published as a conference paper at ICLR 2026
LLM USAGE
We used large language models solely aid polishing the writing of this paper.
A PROOF OF THEOREM 4.
Proof. We begin with lemma adapted from Staib & Jegelka (017)) with minor adjustments to match our framework. This lemma provides an equivalent form for the inner supremum problem of DRO with W -ee-hiighborhood, which is closely related to the representation (6) of w.
Lemma A.0.1. (Staib Jegelka 2017 Proposition 3.1) Let be fixed model parameters, and let
e) be a metric the input space X. For any distribution P on X x y and for any 0,
Ep
sup (X,)
2(0; (I, 2))
sup (Q.()
(X,
where (, y) = ((x',))) c((( v), (x',8') E
With the ambiguity set (8) the DRO 3 is equivalent to
inf
sup BE BmAm-1
sup Wx (Q PS)S g I....
(10)
AEB
where Q = Eg= PgQg.
For fixed 0, the double supremum in (10) can equivalently be written as
172
sup AEAm
sup
Bg Eg, (e(00 (x(,
Po) M
g=1
sup 3
B
sup W (g),
Eq, 1e(0; YJ]]
By applying Lemma A. one can upper-bound the inner supremum in the previous display
E. x:1|((a)=
sup
(10 (z)
EEp,
sup x:11=(a)- ((X])S,
-(f((2())))
cUR(t).n)
S Ep,
sup =1 Ise,
where & =(x) denotes the feature map defined in upper-bounded by
Thus, the original optimization problem is
inf sup 1EG BEAm 9=1
Σ ByE,,
m
sup -
((P(()))))
and completes the proof.
B CONVERGENCE ANALYSIS OF ALGORITHM
We analyze convergence via ěy of the average iterate 0 (1:7)
(( (1:1)
ST = max
.
min max L(0,
BEAm.
8E SEA..
where L(0,3) = Ex= Ep sup - C (J2(=). In the convex setting, our method achieves O(1/VT).
15

---

Published as a conference paper at ICLR 2026
Proposition B.1 (Convergence of Algorithm 1). Suppose (ff,(2), y) is non-negative, convex in 8, Bx-Lipschirz in 0, and bounded by Be for all (x, v) in X x V. In addition, let ||01|2 Be for all 0 in some convex set & R2, and assume the feature map 2(x) is fixed w.r.f. Then, the average iterate of Algorithm I achieves expected error at the rate
10 (Bg B + B2 logm) m)
CT S 2m
Proof. Each iteration samples G Unif m) and (X,Y) Pc. The resulting joint sample Ç = (X,Y,G) is drawn i.i.d. from the mixture distribution == - mm 2g= P..
For each group ".....,,, define the stochastic loss function
E,) = m 1[G = g)
sup | [
((8((=))))
and let
-UB(E))))
fo(0) = -Ep,
sup (X")
The total objective is then L) = (0).
We now verify the conditions required to apply the standard online mirror descent (OMD) regret bound (Nemirovski al. 2009):
(A) Convexity. For each g. the inner function (2)), is convex and non non-negative in 8, and the supremum preserves convexity via Danskin's theorem. Thus, f0) convex.
(B) Expectation form. We have
m
E E (0:EE)
Σ =1=1
EXXYPP
m 1. = g]
sup (1
-(F((2)) = (o()).
m
(C) Unbiased subgradients. By Danskin's theorem, the mapping 0 + supy (=)((='), Y) is subdifferentiable. Hence, Vo {) is an unbiased subgradient:
E [ (o:E(()))) (e),
With the conditions (A)-(C) established, and using the boundedness assumptions:
||0||2 I Be,
|VoC| < Bv,
L B
the standard OMD regret bound (Nemirovski al. 2009
Sagawa et al.
2020 yields
10 BB BG + B& log m)
E E(er] <
2m
T
completing the proof.
The convergence guarantee in Proposition] B.1 is derived under the assumption that the feature map =(-) fixed with respect 0. Under this assumption, the resulting optimization problem becomes convex which allows us to establish the stated (1/VTT) convergence rate. Accordingly, Propo- sition B. should be interpreted characterizing the behavior of Algorithm i the idealized regime where the representation is fixed, rather than as a guarantee for the full nonconvex joint optimization of all network parameters.
16

---

Published as a conference paper at ICLR 2026
C INTERPRETING LATENT PERTURBATION REGULARIZATION
To clarify the intuition behind our latent perturbation framework, we employ a first-order Taylor expansion of the loss function. This approximation shows that the innermost supremum in opti- mization problem can be interpreted as the original loss (f (x). y) plus an additional regularization term involving the dual norm of the gradient with respect to the latent representation. Specifically,
sup ((L(()))))= >=
(()(((()) + + =xx)))
sup
|| (10
(f(((()))) + V.L v) ( (s' (s'- || (1))=(a))
sup
(f v) + sup
v= =x))
C(f((z) v) + + (
(z(()
where || ||- denotes the dual norm corresponding to This added regularization term penalizes large gradients in the latent space, promoting robustness by ensuring that small perturbations x(x) do not lead to significant changes in the loss. By minimizing this term alongside the original loss, the model gains stability and improved performance under real-world distributional shifts.
D DATASET DETAILS
D.1 ORIGINAL DATASET
Colored MNIST (CMNIST) (Arjovsky al. 2019). The CMNIST dataset is designed for noisy digit recognition task, incorporating color as spurious attribute. The dataset is divided into four distinct groups based class and color: 91 f0, green]. 92 = {1, green),
), 93 {0, red), and I4 = {1, red). involves two classes: class 0 includes the digits (0, 1, 2, 3.4) 4) and class 1 includes consists of 30,000 samples, where for class O, the ratio of red to green samples is 8 2, while for class I, this ratio is 2 : 8. The validation set, which comprises 10,000 samples, maintains red to green samples for each class. The test set includes 20,000 samples and introduces a more pronounced group distribution shift: class 0 has red to green sample ratio of 9, and class has ratio of 9 1. Following the approach proposed by Arjovsky et al. (2019). labels in the dataset are flipped with probability 0.25.
the digits (5, 6, 7, 8, 9). The training
equal distribution of color across both classes, with 1 1 ratio of
3
3
9
9
(a) 0, green
(b) 1, green
(c) 0, red
(d) 1, red
Figure 2: Example images from the CMNIST dataset. The groups are gi greenn 92 =
green) gs = (0, red), and 94 red..
Waterbirds Sagawaa al. 2020). The Waterbirds dataset is designed to classify images of birds into two categories: "waterbirds' and "landbirds", with a deliberate introduction of spurious correlations between the bird type and the background. The dataset is divided into four distinct groups based on bird type and background: 9i = [landbird. land), 92 = (landbird, water), 93 = (waterbird, land), and 94 (waterbird, water). This synthetic dataset is created by combining bird images from the
17

---

Published as a conference paper at ICLR 2026
Caltech-UCSD Birds 200-2011 (CUB) dataset (Wah et al. 2011) with backgrounds from the Places dataset (Zhou et al. 2017). Waterbird species, such as albatross, auklet, cormorant, frigatebird, classified as landbirds. The dataset and others, are grouped together, while all other species comprises 4,795 training samples distributed as follows: 3,498 landbirds on land backgrounds, 1 1,057 waterbirds water backgrounds, 184 landbirds backgrounds. This setup highlights the minority groups and the inherent spurious correlations. In contrast to the training set, the validation and test are constructed to have an equal number of samples for each group within each class. The minority group, waterbirds on land, emphasizes the skewed distribution of the dataset, making it suitable for studying the impact of spurious correlations on model performance. The Waterbirds dataset is accessible through the Wilds library in Py PyTorch (Koh et al. ).21).
water backgrounds, and 56 waterbirds land
(a) landbird, land
(b) landbird, water
(c) waterbird, land
(d) waterbird, water
Figure 3: Example images from the Waterbirds dataset. The groups are 91 = (landbird, land),
g2 (landbird, water], 93 {waterbird, land], and 91 { (waterbird, water].
CelebA (Liu et al. 2015). The CelebA dataset is used for hair-color prediction task with fa- cial images of celebrities, where the target labels are "blond" and "non-blond" hair colors. For experimental purposes, the dataset is divided into four distinct groups based hair color and gen- der: 91 = [non-blond hair, female]. 92 (non-blond hair, male), 93 = (blond hair, female), and 04 (blond hair, male]. Gender serves spurious feature, introducing correlations between the hair color and gender of individuals. The training set consists of 162,770 images distributed as fol- lows: 71,629 females with non-blond hair, 66,874 males with non-blond hair, 22,880 females with blond hair, and 1,387 males with blond hair. The validation set includes 19,867 images, with 8,535 females with non-blond hair, 8,276 males with non-blond hair, 2,874 females with blond hair, and 182 males with blond hair. The test set comprises 19,962 images, with 9,767 females with non-blond hair, 7.535 males with non-blond hair, 2,480 females with blond hair, and 180 males with blond hair. The minority group in this dataset is males with blond hair, which constitutes a small fraction of the data, highlighting the skewed distribution and the presence of spurious correlations.
(a) non-blond, female
(b) non-blond, male
(c) blond, female
(d) blond, male
Figure 4: Example images from the CelebA dataset. The groups are 91 = [non-blond hair, female), f2 (non-blee blond hair, male), 9s {blond hair, female), and 94 {blond hair, male).
18

---

Published as a conference paper at ICLR 2026
D.2 MODIFIED DATASETS
Building on the previously introduced datasets CMNIST, Waterbirds, and CelebA we con- structed modified versions of these datasets by applying conditional distribution shifts the mi- nority groups, simulating real-world scenarios. Below, we detail the modifications for each dataset and illustrate these shifts with corresponding figures.
Modified CMNIST. In the CMNIST dataset, we created modified version where the minority group's s images (label red) were rotated by 90 degrees in the test set, while they remained unrotated in the training set. This manipulation simulates conditional distribution shifts often encountered in real-world applications. Figure [ggureeppros provides an illustration of this shift, showing example images from the train and test sets.
9
7
U
(a) Train set
(b) Test set
Figure 5: Example of conditional distribution shift in the CMNIST dataset, where the minority group (label I, red) images are rotated by 90 degrees in the test set, while they are unrotated in the training set.
Modified Waterbirds. For the Waterbirds dataset, we constructed modified version where the minority group (waterbird, land background) was designed to have a shift in species composition between the train and test sets. Specifically, the training as Gadwall, Grebe, Mallard, Merganser, and Pacific Loon, while the test set contained exclusively seabird species, including Albatross, Auklet, Cormorant, Frigatebird, Fulmar, Gull, Jaeger, Kitti- wake, Pelican, Puffin, Tern, and Guillemot. During the dataset construction process, we identified and corrected mislabeling issue involving three species-Westt Western Wood-Pewee, Eastern Towhee, and Western Meadowlark-
included only waterfowl species, such
had been incorrectly labeled waterbirds instead of landbirds (Asgari et al., (0022). Figure iillsstrates this shift, highlighting the separation of species between the train and test sets.
(a) Train set
(b) Test set
Figure 6: Example of conditional distribution shift in the Waterbirds dataset, where the minority group (waterbird on land background) consists of waterfowl in the training set and seabirds in the test set.
To further highlight the impact of this modification, Figure compares the original distribution and the modified distribution shift scenarios. In the original dataset (Figure ((a). bird species in the minority group relatively evenly distributed across train, validation, and test sets. However, in
19

---

Published as a conference paper at ICLR 2026
the modified version (Figure 7(b). the training set contains only waterfowl, while the test set is composed entirely of seabirds, creating distinct distribution shift.
Speckes
Speckes
(a) Original distribution minority group.
(b) Distribution shift of the minority group (only wa- terfowl training, only seabirds in testing).
Figure 7: Comparing the original and shifted distributions the minority group (waterbird, land background) in the Waterbirds dataset (left: seabirds, right: waterfowl, split by dashed line).
Modified CelebA. the CelebA dataset, we modified the minority group (blond hair, male) to have different attributes between the train and test sets. Specifically, the training set contained only images without glasses, while the test contained only images with glasses. This modification flects real-world distribution shifts where rare attributes in small minority groups may change across different distributions, impacting model performance. Figure shows example images demonstrat- ing this shift.
(a) Train set
Test set
Figure 8: Example of conditional distribution shift in the CelebA dataset, where the minority group (blond hair, male) included only images without glasses in the training set and images with glasses in the test set.
Figure provides detailed comparison of the original and modified distributions for the CelebA dataset. In the original distribution (Figure Пa)). the minority group is predominantly represented by the "Without Eyeglasses" category across train, validation, and test sets, with relatively few examples in the "With Eyeglasses" category. In the modified version (Figure 9b)), the training set consists exclusively of "Without Eyeglasses" images, while the test set contains only "With Eyeglasses", creating clear disjoint in key attributes between training and testing phases.
By introducing these conditional distribution shifts, our modified datasets simulate real-world chal- lenges, particularly in scenarios where small minority groups are highly susceptible to such changes. These constructions not only reflect practical settings but also provide realistic benchmarks for eval- uating the robustness and generalization capabilities of machine learning models under diverse and challenging conditions.
E BASELINE DETAILS
We compare method against range of representative baselines:
20

---

Published as a conference paper at ICLR 2026
100
Treeende of Cpeclaaes
Preserce Cyedlaeee
(a) Original distribution of the minority group.
(b) Distribution shift of the minority group (only "Without Eyeglasses" in training, only "With Eye- glasses" in testing).
Figure 9: Comparing the original and shifted distributions of the minority group (blond hair, male) the CelebA dataset (left: "Without Eyeglasses", right: "With Eyeglasses").
ERM: ERM optimizes average accuracy on the training set without any robust objective or group- specific considerations.
Group DRO (Sagawa et al. 2020): A canonical approach for mitigating spurious correlations using known group labels. By partitioning data into predefined groups and minimizing the worst- case group loss, Group DRO aims to improve the worst-group accuracy relative to standard ERM.
JTT
(Liu et al. 2021: A two-step method that first trains ERM model to identify misclassified samples (viewed as proxies for minority groups), then upsamples these samples and retrains a classifier.
CnC (Zhang et al. 2022): Identifies samples that share the same true class but differ in spuri- ous attributes by analyzing ERM outputs, then trains a robust model with a contrastive learning objective. This does not require explicit group labels.
SSA (Nam et al. 2022): Infers latent groups via loss-based criterion, then applies Group DRO to improve robustness. This method partially automates the discovery of group boundaries without needing full group labels.
LISA Yao al. 2022): Mitigates spurious correlations by using Mixup strategies. Depending on the dataset, LISA employs different Mixup variants (e.g., classic Mixup, CutMix, Manifold Mix) to interpolate images within the same label same spurious attribute, thereby reducing reliance on superficial cues.
DFR (Kirichenko et al. 2023 Balances the dataset by subsampling match the minority group size (the "Subsample" strategy), then retrains an ERM model on this balanced data. This simple yet effective approach substantially improve worst-group performance. For fair comparison, following Deng et al. (2024). we evaluate DFR using only the training dataset for both training and fine-tuning, ensuring consistency across methods, which is denoted DFR
PDE (Deng et al.) 2024); Progressively expands the training dataset during the training process, starting with balanced subset prevent the model from learning spurious correlations. This approach aims to enhance robustness across all groups, including underrepresented ones.
GIC (Han Zou 2024): Uses two-step pipeline where group membership is partially inferred, then robust optimization (e.g., Group DRO) is applied. Similar to LISA, it can incorporate tailored Mixup strategies depending on the dataset's characteristics.
F IMPLEMENTATION DETAILS
For experiments involving newly constructed datasets, reimplemented both proposed method and the relevant baselines. When certain baselines lacked reported results for given dataset, used the performance from Zhang et al. (2022) and Han & Zou (2024) if available; otherwise, performed our own reimplementations under consistent settings. In particular, for the original CM- NIST dataset, reimplemented experiments for DFR and PDE, since their original papers did not
21

---

Published as a conference paper at ICLR 2026
include CMNIST results. In all other cases, referenced performance metries from each baseline's s primary source. All experiments were conducted on an NVIDIA GeForce RTX 3090 GPU, and we make our code available at httttts://cciiii com Sung- Ho Ho-Jo/hierarchiccce... H
Across all datasets, we employed the torchvision implementation of ResNct-50 pretrained on ImageNet, training with SGD at momentum of 0.9 and batch size of 128, following Sagawa et al.] (2020). Our approach also introduces a perturbation parameter с to control t/yīgg where 7g represents the
within-group uncertainty. Specifically, we define Eg size of group 9 in the training data. To determine 6, we performed grid search over the set (12/255, 24/255, 36/255, 48/255, 60/255, 72/255, 84/255, 96/255), scaling each value by //mmin Here, Amin = ming n2 denotes the smallest group size in the training set. Additionally, we tuned the generalization adjustment parameter C over {0. 1, 2, 3), as described in Section 3.3 of Sagawa et al. (2220). This setup was applied consistently across every dataset.
For CMNIST, we conducted grid search over learning rates {10-4, 10-3, 10-3) and &g penalties {10-1, 10 10= for 50 epochs. Due to instability in training with the selected parameter com- binations the original Group DRO implementation, we applied ReduceLLROnPlatean scheduler starting at learning rate of 0.01, using it consistently for both our method and Group DRO to ensure fairness. For Waterbirds, the learning rate was tuned over {10-3 3 10 10 100 and the le penalty over {10- 10 1). with training conducted for 300 epochs. For CelebA, the learning rate was tuned over (10-1, 10-5) and the f2 penalty over (10-1, 10-2, 1} for 30 epochs. We referred to prior works including Yao al. (2022) aiosal & Li 2023) guide these hyperparameter search ranges.
G TUNING AND ANALYSIS OF
G.1 SELECTION OF
As is common in DRO problems, selecting the size of the ambiguity set is challenging. Since the extent of distributional shifts is generally unknown, many robust learning methods commonly rely on heuristic or even arbitrary choices in practice. We propose a simple data-driven procedure that partitions the training data using a one-dimensional t-SNE (Van der Maaten & Hinton 2008) ordering. This avoids reliance shifted or test-time data while still creating validation splits that mimic intra-group distribution shifts. Our design is inspired by prior work (Duchi & Namkoong. 2021). which used training-date partitioning to simulate distributional shifts.
Specifically, we project each 2(x;) onto a one-dimensional space using t-SNE, rank samples within each group, and split them into five quantiles. The extreme quantiles (top 20% and bottom 20%) are alternately held out as validation sets, with the remaining 80% used for training. This procedure creates realistic validation shifts that disproportionately affect minority groups
Finally, we select the value of that maximizes minority- group accuracy across these validation setups, ensuring that the chosen perturbation radius provides meaningful robustness for underrepre- sented subpopulations.
G.2 VISUALIZING T-SNE ORDERING FOR MINORITY GROUPS
To highlight the utility of t-SN ordering in simulating realistic distribution shifts, we present visual examples of waterbirds from the top 20% and bottom 20% quantiles of the t-SNE projections. This approach effectively partitions samples based on semantically meaningful intra-group differences, enabling validation splits that closely mimic real-world distribution shifts.
As shown in Figure 10 the t-SNE ordering captures nuanced intra-group differences within this minority group. In the top 20% quantile (Figure [O(a)). waterbirds with longer beaks dominate; in the bottom 20% quantile (Figure I0(b). shorter-beaker waterbirds are more prevalent. This contrast illustrates how LSNE-Eriiee partition can create validation splits that mimic real-world distribution shifts. This method not only emphasizes variations within groups but also systematically evaluates the model's robustness under challenging real-world conditions.
22

---

Published as a conference paper at ICLR 2026
(a) Example images from the top 20% of t-SNE ordering
(b) Example images from the bottom 20% of t-SNE ordering
Figure 10: L-SNE-based ordering reveals subtle distinctions within the minority group. (a) The top quantile features waterbirds with longer beaks, while (b) the bottom quantile features those with shorter beaks.
G.3 IMPACT OF ON ROBUSTNESS
The perturbation parameter plays critical role in improving robustness under minority group shifts. Figures [I(a) and [1bb show how increasing affects worst-group accuracy for the Water- birds and CelebA datasets, respectively. Notably, both datasets achieve significant gains in worst- group accuracy when is set above zero, indicating enhanced resilience to distributional shifts.
450
300
ACE
(a) Waterbirds dataset under minority group shift.
CelebA dataset under minority group shift.
Figure 11: Impact of robustness. The x-axis represents values scaled by 255, and the y-axis indicates accuracy. Each point is the mean of 3 runs (solid lines), and the shaded regions show the standard deviation. For this analysis, the learning rate and &g penalty were fixed to isolate the effect of e.
As illustrated in Figure [aa larger values consistently improve worst-group accuracy Water- birds, enabling the model to better manage intra-group variations and subpopulation shifts. similar trend appears in Figure [1(b) for CelebA, further validating the robustness gained by appropriately increasing c.
These findings underscore the importance of incorporating conditional distribution uncertainty into the training framework. By effectively capturing within-group variability, our approach significantly enhances worst-group performance, making it well-suited for handling realistic distributional shifts.
23

---

Published as a conference paper at ICLR 2026
H GRAD-CAM RESULTS AND ANALYSIS
To gain further insight into where each model focuses its attention under minority-group shifts, we visualize Grad-CAM (Selvaraju al., 2017) heatmaps misclassified examples (by Group DRO) that our method classifies correctly. Figure 12 shows examples on the Waterbirds dataset, while
Figure 13 presents examples from CelebA.
Input
Group
Figure 12:
Grad-CAM visualizations for Waterbirds test images from a minority-group shift sce- nario. Each column shows input image (top row), Grad-CAM for Group DRO (middle row), and Grad-CAM for our method (bottom row).
Waterbirds. In Figure 12 the minority-grae group shift involves species changes not observed in the training set. While Group DRO often localizes on narrow region of the sometimes near the torso or background our method exhibits more distributed attention, covering details like the wings, beak, or feet. This broader localization helps the model rely on features invariant to previously unseen waterbird species, enabling robust classification despite changes in the specific types waterbirds encountered.
Input
Group DRO
Ours
Grad-CAM visualizations for CelebA test images from the minority- group shift scenario. Each column shows the input image (top row), Grad-CAM for Group DRO (middle row), and Grad- CAM for our method (bottom row).
Figure 13:
CelebA. Figure 13 shows examples from the minority group (blond-hair, male) in which the test images include glasses an attribute absent from the training set. In these cases, Group DRO erro- neously directs attention toward the facial eyewear regions rather than focusing hair color. By contrast, our method more reliably highlights the hair region, aligning with the intended classifica- tion objective and enabling correct predictions even under previously unseen attributes.
Overall, these visualizations confirm that, under challenging distribution shifts, our hierarchical DRO framework is less prone confounding features and more successful in focusing the task-
24

---

Published as a conference paper at ICLR 2026
relevant regions. This broader and more contextually aligned attention helps maintain strong perfor- mance even when encountering unseen or spurious attributes.
I COMPARISON WITH STANDARD WASSERSTEIN DRO
To contextualize the performance gains achieved by our hicrarchical formulation, we compare against standard Wasserstein DRO, using the same latent-space cost function and identical training hyperparameters. The comparison is carried out on the Waterbirds and CelebA benchmarks under the minority group shift setting considered in experiments, which reflects the primary distribu- tional scenario targeted by our method. Figures [4(a) [4bb) [ summarize the behavior of standard Wasserstein DRO across wide range Wasserstein radii.
Two structural limitations underlie the weak performance observed in these plots. First, standard Wasserstein DRO models uncertainty only through perturbations within single Wasserstein ball around the empirical distribution, without incorporating any uncertainty over group proportions. In spurious-correlation settings, however, variation group proportions constitutes fundamental mode of distribution shift: a minority group that is scarcely observed during training may appear far more frequently at test time, and ignoring this possibility prevents standard Wasserstein DRO from addressing such mixture-level changes. Second, in spurious-correlation benchmarks, different groups have substantially different sample sizes. Our method designed to reflect this by assigning a distinct within-gree group radius each group (c, /n))).g) so that smaller groups whose empirical distributions are statistically less reliable are assigned a larger perturbation radius. In contrast, standard Wasserstein DRO uses single global radius and therefore cannot adapt to group-specific statistical uncertainty, which is crucial in the minority-shift setting.
For both Waterbirds and CelebA, varied the Wasserstein radius report the corresponding worst-group accuracy curves below, Across all radii, standard Wasserstein DRO achieves significantly lower worst-group accuracy than our method, which attains 93.7% and 72.1% on Waterbirds and CelebA, respectively. These results highlight that perturbing the uncondi- tional distribution alone insufficient for capturing the types of shifts present spurious-correlation settings.
broad range of values and
or Al
Word AL
(a) Waterbirds dataset.
(b) CelebA dataset.
Figure 14: Performance of standard Wasserstein DRO across a wide range of radii.
These findings reinforce the main conclusion of our study: achieving robustness under the complex and realistic distribution shifts that arise in spurious-comrelation benchmarks-partiicuaar those in- volving minority-gre -rropp shifts-requires modeling uncertainty over group proportions as well as group-specific conditional distributions.
25


Reviews:
[{"review_id": "42oezqzAAC", "reviewer": "Reviewer_wr5K", "summary": "This paper proposed a robust optimization scheme with hierarchical distribution shifts. It optimizes the worst-case distribution from a mixture of shifted group distributions Eq.(5). As a concrete instance, the mixture coefficients can be arbitrary and the group drift is constrained within a Wasserstein ball as in Eq.(8). A practical algorithm is introduced with convergence guarantee. Experiments on common robust optimization benchmark datasets show that the algorithm can outperform several baselines in the robust optimization literature.", "questions": "What other options would be possible and practical for the general framework specified in (5)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Writing and presentation are very clear and easy to follow\n\n2. The algorithm is theoretically motivated, and a practical algorithm is introduced with a convergence guarantee. \n\n3. The algorithm performs better than several baselines in terms of worst-case accuracy.", "weaknesses": "1. A hierarchical framework is introduced but with only one instantiation. It would be better to showcase the flexibility of the framework by using the $\\rho$ in (5). More importantly, it is unclear why such a $\\rho$ constraint is meaningful in real-world scenarios.\n\n2. Similarly, only the Wasserstein distance is considered as the measure of distribution shift. It would be worthwhile to discuss other options like f-divergence, and whether they will be feasible in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "Z1ahJtMpRs", "reviewer": "Reviewer_GbgA", "summary": "This paper addresses a limitation of existing methods for mitigating spurious correlations, such as GroupDRO, which are robust to shifts in group proportions (inter-group shifts) but fail when the data distribution within a minority group changes between training and testing (intra-group shifts). The authors propose a hierarchical distributionally robust optimization (DRO) framework that models both levels of uncertainty through a hierarchical ambiguity set. This set forces the model to be robust not only to the worst-case mixture of groups but also to the worst-case data variations within each group, formulated using a Wasserstein distance in a latent space. To validate their approach, they introduce new challenging benchmark datasets designed to exhibit these minority-group shifts, demonstrating that their method achieves stronger robustness where state-of-the-art methods fail, while also maintaining top performance on standard benchmarks.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper presents a well-motivated extension to distributionally robust optimization (DRO) for mitigating spurious correlations. It identifies and addresses a failure mode in existing robust learning methods: vulnerability to intra-group distribution shifts. \n\nWhile prior work like Group DRO focuses on robustness to changing group proportions (inter-group shifts), this paper argues that the distribution within small minority groups can itself be non-stationary. \n\nThe proposed hierarchical ambiguity set is a theoretically sound solution that unifies robustness at both the inter- and intra-group levels. The authors introduce new, modified datasets specifically designed to induce the targeted minority-group shifts, providing a direct test of their hypothesis. \n\nThe motivation is made intuitive through clear diagrams (Fig. 1), the mathematical formulation is precise, and the results are analyzed thoroughly, including qualitative insights from Grad-CAM.", "weaknesses": "I've identified several gaps and issues:\n\nThe proof applies Lemma A.0.1 (from Staib & Jegelka 2017) but then immediately relaxes it:\n\n$$E_{P_g}\\left[\\sup_{x: d(z(x), z(X)) \\leq \\epsilon_g} L(f_\\theta^L(z(x)), Y)\\right] \\leq E_{P_g}\\left[\\sup_{z': d(z', z(X)) \\leq \\epsilon_g} L(f_\\theta^L(z'), Y)\\right]$$\n\n- This inequality can be strict (not tight) when $z(\\cdot)$ is not surjective onto the $\\epsilon$-ball\n\nAlso proposition B.1 assumes:\n- \"the feature map $z(x)$ is **fixed** w.r.t. $\\theta$\"\n\nBut $z(x)$ is defined (equation 7) as:\n\n$z(x) := f_\\theta^{L-1}(f_\\theta^{L-2}(\\ldots f_\\theta^1(x)))$\n\nThis is a direct contradiction - $z(x)$ clearly depends on $\\theta$ through layers 1 to $L-1$. The convergence analysis is therefore invalid for the actual algorithm being run.\n\nThe cost function (page 5) sets:\n\n$c((x,y), (x',y')) = \\begin{cases} \\|z(x) - z(x')\\|, & \\text{if } y = y' \\\\\\\\ \\infty, & \\text{otherwise} \\end{cases}$\n\nThis means $W_p(P, Q) = \\infty$ whenever marginal distributions of $Y$ differ. \n\n- The paper claims \"this definition does not cause any issues\" because $G = (Y, A)$, but this severely restricts the ambiguity set - distributions can only shift within same-label groups\n\nThe paper doesn't compare against:\n- Standard Wasserstein DRO methods (Kuhn et al. 2019, cited but not compared)\n\nAlso as a notation note, $P$ vs $\\hat{P}$ usage is sometimes unclear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "DoCF2BkdTa", "reviewer": "Reviewer_tVrP", "summary": "The authors propose a hierarchical uncertainty set approach that takes care of inter-group and intra-group shifts based on robust optimization. Their method is effective in improving the worst group accuracies under certain data shifts scenarios, especially those with shifts in distribution beyond just changing group proportions. They have evaluated their approach on multiple datasets and obtain small improvements against competing approaches.", "questions": "- How do the authors determine the size of the within-group ambiguity set? This could be a difficult parameter to tune.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed method is effective in improving the worst group accuracy in handling distribution shifts, compared to GroupDRO and competing methods on a variety of datasets including CMNIST, Waterbirds, and CelebA. \n\n- The main idea of the method is clear and well-presented: in addition to uncertainty in shifts in group proportions, add a second layer of uncertainty in the within-group samples.", "weaknesses": "- Compared to vanilla GroupDRO, the improvement is relatively small (Table 2) for the unshifted case. \n\n- As with other forms of robust optimization that handles uncertainty in the samples, there is a tradeoff between optimizing for average and worst-case accuracies. Although we could obtain better worst case accuracies with the authors' method, it suffers worse average accuracies compared to GroupDRO in Table 1. It could be difficult to determine this tradeoff in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)