CODA: A Correlation-Oriented Disentanglement and Augmentation Modeling Scheme for Better Resisting Subpopulation Shifts
Authors: Ziquan OU
Zijun Zhang
Venue: NeurIPS 2024 poster (2024)
Keywords: deep neural networks
feature engineering
model generalizability
class imbalance
spurious correlations
Primary area: deep_learning_architectures

Abstract: Data-driven models learned often struggle to generalize due to widespread subpopulation shifts, especially the presence of both spurious correlations and group imbalance (SC-GI). To learn models more powerful for defending against SC-GI, we propose a {\bf Correlation-Oriented Disentanglement and Augmentation (CODA)} modeling scheme, which includes two unique developments: (1) correlation-oriented disentanglement and (2) strategic sample augmentation with reweighted consistency (RWC) loss. In (1), a bi-branch encoding process is developed to enable the disentangling of variant and invariant correlations by coordinating with a decoy classifier and the decoder reconstruction. In (2), a strategic sample augmentation based on disentangled latent features with RWC loss is designed to reinforce the training of a more generalizable model. The effectiveness of CODA is verified by benchmarking against a set of SOTA models in terms of worst-group accuracy and maximum group accuracy gap based on two famous datasets, ColoredMNIST and CelebA.

CODA: A Correlation--rientee Disentanglement and Augmentation Modeling Scheme for Better Resisting Subpopulation Shifts
Ziquan OU
Zijun ZHANG*
Department of Data Science
Department of Data Science City University of Hong Kong zi zijzhangocity.. edu. hk
City University of Hong Kong ziquanou2-cony cityu. edu.hk
Abstract
Data-driven models learned often struggle to generalize due to widespread subpopu- lation shifts, especially the presence of both spurious correlations and group imbal- ance (SC-GI). To learn models more powerful for defending against SC-GI, we pro- pose CoorellaionnOOiien
Disentanglement and Augmentation (CODA) mod- eling scheme, which includes two unique developments: (1) correlation-oriente disentanglement and (2) strategic sample augmentation with reweighted consis- tency (RWC) loss. In (1), a bi-branch encoding process is developed to enable the disentangling of variant and invariant correlations by coordinating with decoy classifier and the decoder reconstruction. based on disentangled latent features with RWC loss is designed to reinforce the training of a more generalizable model. The effectiveness of CODA is verified by benchmarking against set of SOTA models in terms of worst-group accuracy and maximum group accuracy gap based on two famous datasets, ColoredMNIST and CelebA.
(2), a strategic sample augmentation
1 Introduction
One grand challenge that impedes the generalization of machine learning models is 'Subpopulation shifts', the alterations the relative proportions of subpopulations between training and testing datasets [19.. In reality, one frequently encountered scenario is that specific groups in data are underrepresented, which we tag as the group imbalance (GI), contributed by class imbalance (CI),
attribute imbalance (AI), or their combination [29] [5). Due to GI, models trained to minimize average loss tend to favor the majority groups and thus fail to generalize across minority groups [25)
Another widely discussed issue in subpopulation shifts is the spurious correlations (SC), the incidental correlations between non-causal features and labels during training time that not hold in the real world. SC can significantly compromise the efficacy of model [ and has been observed in various applications, such as speech recognition I0, medical imaging (27), object recognition 25 g image captioning (7]. etc. For instance, object recognition models that classify animal species have been shown to generalize poorly to new environments since their classifications mistakenly favor image backgrounds over biological characteristics of animals 22 28). Learning such correlations can result in degraded performance on data breaking these spurious patterns, e.g., when animals are moved indoors to an uncommon context, such as cows on a beach instead of in pasture.
Prior research has extensively explored the developments of robust models performing uniformly well across subpopulations. Reported studies have varied from attempts of developing models without
"Correspondence Author
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

(a) Grass
(b) Beach
Desert
(d) Initialize
Trap
Disentangle
Figure 1: (a) ~ (c): Same bird embedded in different backgrounds. A robust model is expected to
predict consistently well on objects in all images. (d) ~ (f): The feature extraction and exchange
process of COD. Z. represents spurious correlations, and T represents causal correlations.
Swayles
Sundles
Ralust Classfer
You Ani
Yed Aud
wetk sarple
stge
Uoac mauu and
Figure 2: Overview of the CODA Framework on ColoredMNIST dataset. In ColoredMNIST, samples with digit less than are negative samples (Y = 0), and the rest are positive samples (Y 1). Each sample is either painted red (A = 1) or green (A = 0), with color spuriously correlated with the label in training set. In stage one, CODA first learns encode and disentangle causal correlations from spurious ones. The decoder reconstructs the input data from the two encodings, while the decoy classifier employed to lure the spurious information flowing towards the variance encoder. In stage two, CODA further enhances robustness by creating synthesized samples through the recombination of encoded features from different inputs, ensuring that the resultant samples maintain original class information but vary in spurious attributes. Finally, robust classifier is trained with novel reweighted consistency loss to deliver consistent performance across both original and synthetic samples, thereby reinforcing its resilience to spurious correlations.
methods leveraging knowledge of these spurious [5]. A critical observation in [21] unveiled that methods trained without utilizing spurious attributes still rely them in the model selection; otherwise, a significant performance drop would be observed. The potential of better utilizing spurious attributes further strengthen models in defending subpopulation shifts has yet to be fully explored, especially in terms of more innovations in utilizing spurious attributes model selection, addressing SC by mixing, or tackling GI via reweighting. Meanwhile, is also of great interest to explore the utilization of spurious attributes as a means to enhance model performance rather than impediments be overcome.
pre-identified spurious attributes [21 23 26 20
attributes to inform and guide training processes 25
B
30
14
This paper thus aims to develop a novel method of learning models for better defending the SC-GI type of subpopulation shifts from the aspect of innovating the mechanism of dealing with spurious attributes. Considering the object recognition task with backgrounds, addressing SC can be translated more straightforward question Can design model capable of reliably recognizing the same object across a range of spurious attributes, such as diverse backgrounds? To answer this question implies that the developed model needs to focus on the invariant and causal features pertinent to object identification, as conceptually illustrated in Figures Ta Ic Conducting such a study leads to two further pivotal technical questions:
Q1: How to synthesize samples that vary in spurious attributes yet retain accurate class information?

---

Q2: Based on the synthesized samples. how to design learning mechanism for developing models more powerful for addressing the SC-Gr GI subpopulation shifts?
To make a response, a novel framework, the Correlation Oriented Disentanglement and Aug- mentation (CODA), developed in this work to develop machine learning models better handling the SC-GI subpopulation shifts by disentangling variant and invariant features, coordinated with strategic sample augmentation process. The overview of the CODA framework is shown in Figure 2 Two unique developments in CODA are briefed as follows:
Correlation- -rrietted Disentanglement (COD) To facilitate the disentangling between class information and spurious information, a novel bi-branch encoding process called Correlation-Oriencte Disentanglement (COD) is developed in this paper. As shown in Figures [ddIi COD operates through a 'trap and disentangle' mechanism. The variance encoder, coordinated with a decoy classifier that predicts the spurious attribute and novel design of disentangling loss, extracts only the necessary information about the spurious attributes. The decoder and a reconstruction loss then simultancously regularize the invariance encoder to capture distinct and complementary information (including invariant features and label-irrelevant features). These bi-branch encoders, coupled with the learned decoder, facilitate the generation of rich and diverse synthetic dataset.
Strategic Sample Augmentation with Reweighted Consistency Loss Building on insights from recent studies 29 [3,, which highlight the effectiveness of upweighting sampling probabilities or losses for minority groups, we extend this concept by strategically generating additional synthetic samples for these groups. By integrating variant and invariant features extracted from different samples, the decoder can generate realistic synthetic data that exhibit different spurious information while preserving class fidelity, aligning with the insight shown in Figures [a we introduce novel weighted consistency loss designed to prompt the classifier to deliver uniform predictions across both the original and synthetic samples. Our experimental results affirm that CODA harmonizes effectively with pre-existing robust classification methodologies by integrating the proposed weighted consistency objective.
To complement this,
To the best of our knowledge, the work most akin to ours from the conceptual aspect is CAMEL 9 which applies CycleGAN for learning image transformations. CODA clearly differs from CAMEL in the following aspects: (I) CycleGAN is deterministic and yields a single fixed output for a given input sample, while CODA introduces substantially more flexibility in synthetic data by mixing variant features from multiple samples. (2) CycleGAN is limited to one-to-one transformations, necessitating (%) models to cover all potential mappings for n attribute values within one spurious attribute. In contrast, CODA can simultaneously learn multiple transformations and can be readily extended to accommodate multiple spurious attributes by incorporating additional decoy classifiers. (3) GANs are notoriously challenging to train with instability issues. In contrast, CODA offers more stable and expedient training process.
The primary contributions of this paper are as follows: (1) The potential of better utilizing spurious attributes, rather than treating them as impediments to be overcome, explored. A novel correlation- oriented disentanglement and augmentation framework for handling the SC-GI type of subpopulation (2) A bi-branch encoder network architecture coordinated with a decoy classifier and subsequent decoder is designed to realize disentangled representation learning for handling SC. (3) A strategic sample augmentation with a novel weighted consistency loss proposed for upweighting samples predicted with low confidence is developed to better handle the SC-GI. The proposed CODA framework can operate synergistically with existing robust classification methods. (4) Extensive experiments conducted confirm that CODA can drive the learning to focus on causal rather than spurious features and uplift the performance on defending extreme SC-GI subpopulation shifts.
shifts is thus developed.
2 Preliminaries
2.1 Formulation of learning models defending subpopulation shifts
Consider classification tasks in which we aim to predict label y E y based on an input x E X. Each instance in the dataset is also characterized by an attribute E A, which may influence the prediction. We define group by tuple g = (y, a) G. where each group has its own distribution Pg.
3

---

Then the training and testing data distributions can be expressed as mixtures of these subpopulation distributions:
Pirino - o;PP and Prest >. B P,, where a, 8 E Aig represent the mixture weights. Subpopulation shifts occur when the mixture weights for training and testing differ, i.e., + B. If the mixture weights of specific groups dominate, group imbalance occurs, potentially leading to a spurious correlation between the attribute and the label.
Given loss function f, our objective is to identify function f X + y that minimizes the worst-case expected loss over all possible mixtures [25), formally defined as:
argmin sup E gEA G
up (a.)~ un f()))
(I)
2.2 Benchmarks for defending subpopulation shifts
Empirical Risk Minimization (ERM) Consider a training dataset D {(x,, ys: 0z) 1a compris- ing N samples. Let Parain denote the empirical training distribution. Then common machine learning
practice applying ERM aims find a function f that minimizes the average empirical loss:
fiRM argmin (a.))~B... [C(y. fx))).
However, ERM typically fails in scenarios with subpopulation shifts due group imbalance and learning spurious correlations that do not hold in testing data.
Reweighting Groups (RWG) Given the highly imbalanced nature of subpopulation shifts, Idrissi et al. [13 studied the efficacy of data balancing techniques, especially RWG. RWG adopts group- balanced sampling strategy upweight minority groups, as shown in Eq. 43:
Pg
(3)
N
where g is the number of training samples in group g. and Ps represents the probability of sampling an instance from group g. Despite its simplicity, RWG achieves competitive performance and serves as strong baseline in studying subpopulation shifts.
Group Distributionally Robust Optimization (GDRO) The study 25 showed that optimizing Eq. ① is equivalent to minimizing the maximum expected loss over all groups The empirical optimization objective in GDRO is expressed as:
fODRO argmin max E
gEG P (L)
2.3 Disentangled representation learning
Disentangled representation learning targets extracting data representations that separate the distinct sources of variation and serves as the base of the COD development in this work. This section explores relevant models that have generated significant impacts such topic.
Variational Autoencoder (VAE) and variant Kingma and Welling [8] postulated that the data generative process is governed by an unobserved vector of latent codes with prior distribution p(z). The Vanilla Variational Autoencoder (VAE) is introduced to approximate the intractable true posterior with a variational posterior (2(1x) and to learn decoder ps 9 respectively. VAE's objective to maximize:
(s(x)2) (z]22, with parameters ó and
LA Φ) E (n)
(L((
(5)
V-E [1 incorporated a hyperparameter / to the KL term in Eq. (5) to control the trade-off between the reconstruction quality and the independence intensity between latent variables. When 3 > 1, the model prioritizes learning independent latent factors.
The study
12 showed that the KL term in Eq. can be decomposed into two components:
Epgn (w)
[(x; 2) + (L(go())Ilp(z))
(6)

---

in which (z) (=) and [x; 2) represents the mutual information between and under the distribution of Pdata (z)g(= x)) A higher value of 3 enforces a closer approximation of the latent codes the factorized prior p(z), leading to more independent latent codes. However, this also imposes a stronger penalty on the mutual information I (z; 2), which may reduce the amount of information about retained in and potentially worsen the reconstruction quality.
Flexibly Fair VAE (FFVAE) FEVAE [3] aimed to learn feature representations / that achieve subgroup fairness and targeted for sensitive attributes, with the following objective:
CO 1. = Nu (e (x]2,, (alle + a Po (a/=)]
(2, (/r))l( f))] (2 ()lg(() az)))
in which py (a|z) models the prediction of sensitive attributes. The 7-weighted KL divergence term encourages the latent variables 2j to be independent of both t and 2; for all i + j.
Due to the intractability of the second KL term in Eq. FFVAE employs adversarial density ratio estimation [16]. an estimator cannot differentiate whether a sample pair (2,t) comes from ((2 or the factorized distribution g(1) II ((z)), the second KI. term is minimized, suggesting the independence of and f. However, this approach inevitably introduces new density estimator component, along with GAN training objective, which can be challenging and unstable to train.
Correlation--rientte Disentanglement and Augmentation (CODA)
This section explains the novel mechanism of CODA utilizing spurious attributes for defending the SC-GI subpopulation shifts, which is composed of two parts, disentangling variant and invariant features well strategic sample augmentation with a reweighted consistency loss based on the disentangled features.
3.1 Correlation- Disentanglement (COD)
Given sample (x,,y, a, g), aim to learn encodings R and E R defined by stochastic parametric encoders qu (2(2) = (x))), (x)) and go = ((t)),())) (x)) respectively. Factorial priors p(z) for and p(t) for are assumed. Our model includes decoder, denoted by (0]22, t), and a decoy classifier pa (a[z), which predicts the spurious attribute a.
In this paper, we propose the following training objective for disentangled representation learning:
LCOD(O, w) = Epan (x) (z (f)x) (log (0)22, (z t) + a lo (a|2)]
(L
(8)
By applying Eq. 6 Eq. can be re-written as:
LCoD (0, 0, 1/, w) (u) E (o) log (o () to log pa
reconstruction
classification
11 (1;2) [7 KL (2) + KL (g (()))(())
mutual information
KL divergence losses
To clarify, Eq. (8) is the practical training objective for implementation, while the equivalent Eq. offers justification for motivation of the proposed COD. The objective function in Eq. (9) consists of four components: (1) reconstruction loss to ensure and t collectively encode the information in x, (2) a classification loss that encourages z to predict spurious attributes accurately, (3) mutual information terms that limit the extent to which encoders capture relevant information from and (4) KL divergence terms that penalize the encoders if they diverge from the factorial priors.
The working mechanism of the four components in Eq. (9) is described as follows:
The classification loss term rewards be maximally predictive of the spurious attribute. By assigning heavy weight
(7 1) minimizing [(x; z), we limit the expressive power of ensuring it encodes only the necessary information about the spurious attributes. contrast, we place a relatively low limit on minimizing (æ; t), thus ensuring that the reconstruction quality of the decoder is not degraded. While
5

---

the reconstruction loss term ensures that and t collectively capture all the information contained in I, is expected to carry distinct and complementary information. Furthermore, the two KL terms regularize the variational marginals to match the factorial priors, thus facilitating disentangling within each feature representation. preliminary theoretical exploration of the performance guarantee of the decoder training has been conducted with mathematical proof, which is offered in Appendix [
We emphasize that our goal is separate spurious and variant correlations from those that are causal and invariant, rather than to seek disentanglement within each individual representation, i.e., imposing each t; be independent t, (for all j # i). Consequently, a lower weight on the final KL divergence term is both justified and sufficient, as it supports our objective without imposing unnecessary constraints on the independence within the invariant feature space, Our proposed method circumvents the need for density estimation, thus avoiding the complexities and instabilities associated with training objectives based on the GAN-like density estimator.
3.2 Strategic sample augmentation with reweighted consistency loss
Strategic sample augmentation Through the disentanglement of variant and invariant features, as outlined in the previous section, we can now strategically generate synthetic samples for augmentation by leveraging the trained decoder. For a given batch of B samples {(T;, yr, 0z, 9:) we begin by extracting feature representation pairs {(2,, t,)) where 24 = qu(2/z;) and t; = (t)xr)) With hyperparameters L. which controls the number of synthetic samples per batch instance, randomly select indices from (1...., B} for each instance, yielding mixing set {( 1=)))) By combining features from different instances, we create the set of synthetic samples { = ) where æi.j = po ( , t;). These samples exhibit varied spurious information but retain their original class information. A model trained to perform well across augmented samples is likely to learn robust decision-making rules that do not depend on spurious correlations.
Reweighted consistency loss A natural intuition to prevent the final model fs from utilizing spurious correlations is to enforce the model to deliver consistent predictions on the original samples I; and the set of synthesized samples i; (fs.)/=1- To realize, we propose the following reweighted consistency loss:
(RC(IS; = (L
+ KL(fe(z:)lmm)) (10)
in which fe(Dig)y: is the predicted probability of the true class for رgg. and m. Normalize( Σラ= Je (e (((:))) the weighted mean prediction for the synthesized sam- ples, normalized to ensure that the probabilities sum to one.
In Eq. (OO.. we downweight the synthesized samples with low confidence in predicting the right class mitigate its effect jeopardizing th. Conversely, we upweight their KL losses penalize inconsistencies with other samples that have greater confidence in the ground-truth class.
At last, given loss function L. formulate final optimization objective as follows:
ECODA &i, v1) E(y,, fx.) + XfR f4, 91).
(11)
Pseudo codes of training and deployment of the proposed CODA are offered in Appendix B
4 Experiments
In this section, we evaluate the efficacy of the proposed CODA methodology. Specifically, we aim to answer questions: (1) Can CODA learn disentangled encodings that extract spurious and causal correlations, respectively? (Section 4.2); (2) Does CODA enhance robustness in addressing SC-GI subpopulation shifts, and what are the critical factors contributing to its effectiveness? (Section 4.3
4.1 Setup
Dataset description Experiments are conducted on the ColoredMNIST and CelebA datasets 22 Here we provide brief introduction, with comprehensive details available in Table and Appendix

---

Table 1: Dataset statistics. SC indicates spurious correlations, GI indicates group imbalance, and DGPS indicates the degree of group proportion shifts between training and testing sets.
Group
Group 2
Group 3
Group 4
SC
GT
DGPS
Training
21330 (42.26%)
3776 (7.55%%
3724 (7.45%)
21370 (42.74%)
ColoredMNIST
Validation
2550-125.15
2526 (25)3%)
2474 (24.7%)
2490 (24.9%)
High
Testing
761 (7.61
4378 (43.8%%
4122 (41.2%
739 (7.39%)
Training
71629 (44.01%)
22880 (14.06%)
66874 (41.08%)
1387 (0.85%
CelehA
Validation
8535 (42.969))
2874 (14.47%)
(41.66%)
182 (0.92%%
Low
Testing
9767 (48.93%)
2480 (12.42%)
75335 (37.75%)
180 (0.90%)
9
3
2 4 0
7
0
3
0
9
9 9
9
9
9
9
9
9 3 2 8 0 7077
9
9
33
3
3
3
3
3
3
3
3
3 2 8 0 7
2
2
2 8
2
2
2
2
2
2 8 0 7
8
8
8
8 0 7
8
8 0 7
8 0
0
0
0 7
0 7
0
7
7
7
7
0
0
0
3
3
3
3
3
3
3
3
3
3 0
3 0
0
0
0
0
0
0
0
0
(a) ColoredMNIST
(b) CelebA
Visualization of the synthesized samples. Images from the top-row and the leftmost column are real samples, while the remaining images are reconstructions. Each reconstructed image generated by combining t extracted from the corresponding leftmost sample and z from the corresponding top-row sample. The main diagonal images represent same sample reconstructions.
Figure 3:
ColoredMNIST: This variant of the MNIST dataset (6) incorporates color (red or green) into the original grayscale digits. Digits below five are labeled negative, while the rest are positive. Within the training set, 85% of the negative/positive samples are painted green/red, introducing color as spurious correlation. Conversely, in the testing set, this ratio is reduced 15%. Following []] labels are flipped with probability of 0.25.
CelebA: The CelebA dataset consists of large collection of celebrity face images coupled with 40 binary attributes. The task is to predict hair color (blond or non- non-blond), which exhibits strong spurious correlation with the gender attribute (male female).
Evaluation protocol In alignment with 25 we evaluate robustness of all methods using worst- group accuracy, the lowest test accuracy across all groups, as gold-standard. In addition, we also report the average accuracy and the maximum group accuracy gap. Model selection is based on the highest worst-group accuracy obtained during validation.
Hyperparameters tuning For COD, we employ Adam optimizer 17 and adopt pre-trained ResNet- 18 and ResNet-50 encoders for ColoredMNIST and CelebA, respectively, For learning the robust classifier, we train all approaches on both datasets using SGD optimizer [24] with momentum 0.9, and adopt 3-layer CNN and pre-trained ResNet-50 for ColoredMNIST and CelebA all approaches, respectively. Additional training details are available in Appendix D
4.2 Disentangling and synthesis abilities of CODA
CODA can generate samples varying in spurious attributes yet retaining accurate class informa- tion. Generated samples visualized in Figure 3 well justify the superior disentangling and synthesis capabilities of CODA. It is clear that the generated samples in Figure are visually appealing. It is
7

---

Table 2: Average test accuracy (%), worst-group test accuracy (%), and maximum test group accuracy gap (%) of the classifiers on the ColoredMNIST dataset. Variance and invariance classifiers take z and as inputs, respectively. Results averaged over three independent trials.
target= = Y (digit 2 or not) Worst.
target = A (color) Worst. Ace.
Classilier
Avg. Acc.
Max. Gap
Avg.
Gap 0.00 = 0.00 35.60 1:24
variance classifier
15.00 + 0.00
0.00 + 0.00
100.00 = 0.00
100.00 0.00
100.00 1.0.000000 0.00 35.73 + 0.35
invariance classifier
72.17 + 0.31
71.29 + 0.44
2.64 = 0.83
41.44 0.21
Table 3: Performance metrics (%) for ColoredMNISS and CelebA. Results are averaged over three independent trials. Red/blue/green represents the firs/second//hirrd highest performance. indicates that CODA+X shows improvements over base method denoted by X in terms of worst- worst-group accuracy and maximum group accuracy on both datasets, X E (ERM, RWG, GDRO).
ColSredMNIST Worst.
ColebA
method
Avg. Acc.
Gap
Avg.
Gap
LAF 23
67.64 64.64.2 5.19
50.91 1: 2.92
24.36 i 7.08
89.67 0.44
73.11 1: 1.43
20.72 1 2.40
JTT
21
72.11 0.36
71.01 +: 0.50
5.15 +: 1.83
92.10
0.26
76.45
1 0.75
20.90
0.92
ERM
17.02 + 0.76
2.08 = 0.70
92.70 0.71
96.05
+ 0.38
47.78 + 2.30
51.80 +:40+:20 2.45
RWG
13
72.91 0.45
0.15
2.70 = 1.06
0.23
83.41 1.55
10.78
GDRO
25
73.03 -. 0.25
71.08 +: 1.02
3.27 1: 1.13
93.01 + 0.08
88.22 1: 0.89
4,92 1.11
CODAFERM
72.20 0.57
71.74 1 0.24
2.45 1 1.02
91.89 1. 0.35
83.65 1 0.32
9.87
1.08
CODA-RWG
73.20 0.12
72.11 + 0.51
2.36 + 0.67
91.72 + 0.12
86.56 0.82
7.28 = 1.01
CODA+GDRO
73.02
71.98 0.57
2.37
0.91
90.91 + 0.21
89.26 0.26
4.01 0.24
Table 4: Performance metries (%) for ColoredMNIST v2, v3, and v4.
ColoredMNIST v2 Max. Acc. Gap
ColoreMNNSS
ColoredMNIST
method
Worst, Acc.
Woest. Ace.
Max. Acc. Gap
Worst, Acc.
Max. Acc. Gap
LFF
58.99 1 1.95
12.83
4.38
54.68 1. 7.32
15.32 :: 10.65
49.38
4:23
21.35 1: 9.29
JTT
71.09 + 0.73
2.59 + 1.83
69.81 + 0.79
6.9 + 0.63
62.07
0.52
12.82 + 1.81
ERM
2.94
+
0.02
99.98
100.00
0.00
RWG
71.33 1: 0.45
2.71 2.714.00.46
70.93 1: 1.22
2.8 d: 1.33
2.59 70.42
68.18
8.08 i 1.23
GDRO
71.84 0.65
0.27 3.32
69.65 1.50
3.96 + 2.05
0.78
4.17 0.75
CODA+ERM
71.1 + 0.84
3.3 0.58
71.12 + 0.92
3.21 + 1.91
71.03 0.53
3.59 + 0.53
CODA-RWG
71.91
2.42 0.32
2.42
CODA+GDRO
72.05 1 0.27
1.91 1. 0.44
71.86 1: 0.58
2.09 1: 0.22
71.44 3: 0.69
2.36 1 0.54
evident that the digit color/gender in the reconstructed images is consistent with the top-row samples, while the residual features such as digit, written style, and size/background, posture, hairstyle, and hair color-are are determined by the leftmost samples. This indicates that the variance encoder captures only the necessary information about the spurious attribute, while the invariance encoder encapsulates distinct and complementary information, including class-related and irrelevant details. These findings in harmony with our training objectives outlined in Section
CODA can separate spurious information from causal information. We further evaluate the ability of CODA to disentangle spurious and causal information through four classification tasks on ColoredMNIST. In each task, we fix the weights of the learned variance and invariance encoders and extract the corresponding latent codes (z and t) for each image. Here z and f are expected to carry color and digit information, respectively. A variance classifier and an invariance classifier, both structured as simple 3-layer MLPs, are then employed with: one using z as input and the other using t to predict the digit label (whether the digit 2 5) or the spurious attribute (color).
Results are presented in Table We observe that the variance classifier achieved perfect performance in predicting color. In contrast, it achieved 0% worst-group accuracy and 15% average group accuracy in predicting digit, which conformed to the color assignment ratio of 15% the testing set. Conversely, the invariance classifier exhibited robust label classification capabilities, even without applying any reweighting or sample augmentation technique. Its poor color classification performance, comparable to random guessing, justifies the effective disentanglement achieved by CODA.

---

6.8
6.7
0.6
0:13
0000
CODA-ERM
CODA+CIDOO
Figure 4: Sensitivity analysis on the weight of the reweighted consistency loss. When = 0, the methods degrade to vanilla ERM, RWG, and GDRO.
Table 5: Sensitivity analysis the number of synthesized samples per instance on MultipleCol-
oredMNIST. When L 0, the methods degrade to vanilla ERM, RWG, and GDRO.
Methods
L
Avg. Acc. (std)
Worst Acc. (std)
16.60 + 1.23
0.00 0.00
96.90 1 0.29
88.78 0.87
CODA-ERM
97.12 + 0.07
90.36 0.68
97.06 = 0.11
91.85 + 0.68
95.54 = 0.24
85,46 + 0.44
96.92 + 0.04
90.87 0.93
CODA-RWG
96.66 1: 0.28
91.35 0.70
96.94 = 0.08
91.10 + 1.16
91.99 = 0.33
82.16 + 4.60
CODA-GDRO
96. 68 + 0.03
89.88 0.50
96.70 1 0.16
90.10 1 1.10
96.85 + 0.05
90.42 + 1.39
4.3 Benchmarking studies and analysis
CODA enhances performance and robustness existing robust classification methods. Ta- BleBBillustratt the comprehensive performance metrics for CODA compared with benchmarking methods. On ColoredMNIST, ERM performed worst in all metrics due the intense subpopulation shift and the spurious correlation. While on CelebA, which has a less intense difference group proportions, ERM achieved the highest average accuracy but failed in worst-group accuracy due group imbalance and the spurious correlation. We observe that CODA consistently offered improved worst-group accuracy over vanilla ERM, RWG, and GDRO, and significantly curtailed the group accuracy discrepancy across both datasets. This enhancement in robust accuracy suggests that CODA be effectively integrated with prevailing robust classification methodologies. Scalability of CODA to the multi- coltticlasssfication scenario is further evaluated on dataset called MultipleColoredMNN which is detailed in Appendix [..T] Experimental results in Table present that CODA can be scaled up to scenarios with multiple spurious attribute values, demonstrating its great potential in solving complex subpopulation shifts.
Enhanced robustness of CODA extreme SC-GI subpopulation shifts. Additional experiments are conducted to assess the robustness of CODA against varying degrees of subpopulation shifts. Color bias introduced with probabilities p and p in training and testing sets, respectively. ColoredMNIST v2, v3, and v4 are then created with 0.8, 0.9, and 0.95, respectively. A higher value p indicates greater differences in group proportions, increased group imbalance, and more intense spurious correlation. Other dataset settings are consistent with ColoredMNISTT. Experimental results are shown in Ta results are presented Appendix F.2). As the bias intensity increased, CODA maintained more consistent performance metrics compared other methods,

---

which suggested that CODA is more robust to extreme subpopulation shifts compared with baseline methods and can effectively adapt its learning focus on causal rather than spurious features.
The critical role of the optimal coefficient selection for RWC to the success of CODA. Figure
presents a sensitivity analysis of the parameter À from Eq. (T1) on ColoredMNIST dataset. We observe that an elevated À value (A 2 100) compromises model performance due to excessive regularization. Conversely, diminutive À value fails to contribute significantly owing to insufficient regularization intensity. Thus, an optimal value is instrumental for the robust learning of the model.
Trade-off between robust classification performance and computational costs. Another critical hyperparameter in training CODA is the number synthetic samples per instance controlled by L. A larger L. introduces more variability to synthesized samples, driving the model to "forget" the spurious attributes decision-making dataset is provided in Table 5 It is observable that increasing L generally improves the worst-group accuracy. However, a larger also implies higher computational costs. Furthermore, excess L may not introduce more variability (depending 10 in MutiiIeCCllreaMMNIST can result in redundant synthesized samples. Thus, the value of needs to be selected carefully for best handling the trade-off between classification performance and computational efficiency.
sensitivity analysis of L on the MulipleCCoorrdMMNIST
the complexity of the dataset), e.g., L greater than
5 Conclusion
In this study, we introduced CODA, a novel approach designed enhance the robustness of ma- chine learning models against the SC-GI subpopulation shifts. Our extensive experiments based on the ColoredMNIST and CelebA datasets provided compelling evidence that CODA could success- fully disentangle variant and invariant feature representations and, more importantly, utilize these representations for sample augmentation to significantly improve model robustness.
However, it is important to acknowledge some limitations of CODA, notably its reliance on pre- identified spurious attributes. This dependency requires potentially costly labeling effort. Addition- ally, while CODA incurs no additional computational costs at deployment, the training phase is more resource-intensive due the necessity of data augmentation processes.
Despite these limitations, results of this study offered a promising direction for creating machine learning models that robustly generalize beyond biased training distributions and uphold robustness across diverse subpopulations.
Acknowledgments and Disclosure of Funding
This work was supported in part by the Shenzhen-Hong Kong-Macau Science Technology Category C Project with No. SGDX202205301 1 1205037, in part by the Hong Kong RGC General Research Fund Project with No. 11213124, in part by Hong Kong ITC Innovation and Technology Fund Project with No. ITS/034/22MS, in part by Guangdong Provincial Basic and Applied Basic Research Offshore Wind Power Joint Fund Project under Grant 2022A1515240066, in part by Guangdong Province Technological Project with No. 2023A05050300 and in part by InnoHK initiative, The Govermment of the HKSAR, and Laboratory for Al-Powered Financial Technologies.
References
[1] Martin Arjovsky, Léón Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization, 2020.
[2] Sara Beery, Grant van Horn, and Pictro Perona. Recognition in terra incognita, 2018.
[3] Elliot Creager, David Madras, Jörn-Herrik Jacobsen, Marissa Weis, Kevin Swersky, Toniann Pitassi, and Richard Zemel. Flexibly fair representation learning by disentanglement, In International conference machine learning, pages 1436-1445. PMLR, 2019.
10

---

[4] Tal Daniel and Aviv Tamar. Soft-introvae: Analyzing and improving the introspective variational autoencoder. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 4391-4400, 2021.
[5] Tal Daniel and Aviv Tamar. Soft-introvae: Analyzing and improving the introspective variational autoencoder, 2021.
[6] Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141-142, 2012.
[7] Robert Geirhos, Born-Herrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A. Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(1 1):6656673, November 2020. ISSN 2522-5839. doi: 10.1038 s42256-020-00255--2 URL http::/dx.... dwl...ddo..ooooo www.doi...ggiggoo 1038/ 442566 020-00257-
[8] Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A. Wichmann, and Wieland Brendel. imagenet-trained enns are biased towards texture; increasing shape bias improves accuracy and robustness, 2022.
[9] Karan Goel, Albert Gu, Yixuan Li, and Christopher Ré. Model patching: Closing the subgroup performance gap with data augmentation, 2020.
[10] Tatsunori Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang. Fairness without demographics in repeated loss minimization. In International Conference on Machine Learning, pages 1929-1938. PMLR, 2018.
[11] Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. In International conference on learning representations, 2016.
[12] Matthew D Hoffman and Matthew J Johnson. Elbo surgery: yet another way carve up the variational evidence lower bound. In Workshop in Advances in Approximate Bayesian Inference, NIPS, volume 2016.
[13] Badr Youbi Idrissi, Martin Arjovsky, Mohammad Pezeshki, and David Lopez-Paz. Simple data balancing achieves competitive worst-group-accuracyy 2022.
[14] Pavel Izmailov, Polina Kirichenko, Nate Gruver, and Andrew Gordon Wilson. On feature learning in the presence of spurious correlations, 2022.
[15] Bingyi Kang, Saining Xie, Mareus Rohrbach, Zhicheng Yan, Albert Gordo, Jiashi Feng, and Yannis Kalantidis. Decoupling representation and classifier for long-tailed recognition, 2020.
[16] Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In International Conference Machine Learning, pages 2649-2658. PMLR, 2018.
[17] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
[18] Diederik Kingma and Max Welling. Auto-encoding variational bayes. 2022.
[19] Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Sara Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang. WILDS; benchmark of in-the-wild distribution shifts. CoRR, abs/2012.07421,, 2020. URL abs 2tt.//222.....
[20] Daniel Levy, Yair Carmon, John C. Duchi, and Aaron Sidford. Large-scale methods for distributionally robust optimization, 2020.
[21] Evan Zheran 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, 2021.

---

[22] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang, Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
[23] Junhyun Nam, Hyuntak Cha, Sungsoo Ahn, Jaeho Lee, and Jinwoo Shin. Learning from failure: Training debiased classifier from biased classifier, 2020.
[24] Sehastian Ruder. An overview of gradient descent optimization algorithms. arXiv preprint arX14:1609.04477,, 2016.
[25] Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization, 2020.
[26] Amrith Setlur, Don Dennis, Benjamin Eysenbach, Aditi Raghunathan, Chelsea Finn, Virginia Smith, and Sergey Levine. Bitrate-constrained dro: Beyond worst case robustness unknown group shifts, 2023.
[27] Laleh Seyyed- Kalantari, Haoran Zhang. Matthew BA McDermott, Irene Y Chen, and Marzych Ghassemi. Underdiagnosis bias artificial intelligence algorithms applied to chest radiographs in under-served patient populations. Nature medicine, 27(12):117-2218) 2021.
[28] Kai Xiao, Logan Engstrom, Andrew Ilyas, and Aleksander Madry. Noise signal: The role of image backgrounds in object recognition, 2020.
[29] Yuzhe Yang, Haoran Zhang, Dina Katabi, and Marzyeh Ghassemi. Change is hard: A closer look at subpopulation shift, 2023.
[30] Huaxiu Yao, Yu Wang, Sai Li, Linjun Zhang, Weixin Liang, James Zou, and Chelsea Finn. Improving out-of disttiiiitiiitooonn robustness via selective augmentation, 2022.
12

---

A Theoretical exploration
In this section, provide theoretical exploration of the reconstruction performance property of CODA, which is stated in Lemma]) |
Lemmal shows that the decoder in CODA converges to an version of the underlying data distribution.
In the overall model development, the proposed CODA facilitates the robust training of the final classification model by conducting sample augmentation through VAE-style generative process. VAE-style generative models are known to have a stable training procedure, but they tend to generate blurry images that lack details. To facilitate realistic image generation based datasets containing samples more complicated contents (c.g., CelebA), we adopt the Soft-Intra VAE [4] training technique in which the encoder and decoder models are jointly trained in introspective way. No extra network component like a discriminator is needed in Soft-Intro VAE.
g(2, t|æ: = (z(xr)) (t)z). and a decoder = pp (zr|z, , t). factorized prior p(2, d) = p())p(t) is assumed for the latent codes. Let Pdata represent the underlying data distribution. We define pa(æ:) = Bp(z, (z]((z)22 ()] as the distribution corresponding to the samples generated by the decoder.
Recall that we aim to learn factorized variational posterior
In the proposed CODA, the ELBO is formulated as:
ELBO r; q, d) log po (1 (z) t)
(L - (L(oa(|)))lp(t)
(12)
By applying Soft- Intro VAE, the overall objectives of the variational posterior and the decoder formulated in Eq. 133 and Eq. (14.. respectively:
Lg d) (e))e<))))(((;(pppp
Epul exp(sELLBO(r::.)))
(13)
Ca(g, d) = Epsu(e) (EL d)] + nppp (ELBO(rqq.))
(14)
in which 1) 2 and 2 are hyperparameters for Soft-IntroVAE training.
A Nash equilibrium point (g",d") satisfies C, d²) 2 (g(g. and La(g" ,a(g",d')= 2 La(g", d) for all q and d. Define d' follows:
d' E ( (z) (IP) + H((((z)))
(15)
in which H(-) denote the Shannon entropy.
Lemma 1. Let d' be defined in Eq. 15). Denote 9" = pa- (2, Assume the encoders and the decoder have infinite capacities. Assume that there exists no region where Pdtsa(x) = O. and Pa- (z) S
Paoto (x))=x पोx for all x. Then (4" ,d") ) a Nash equilibrium point of the system formulated in Eq. 13 and Eq. 14
Proof. Considering the factorized property of the variational posterior and the prior, one can easily obtain the following:
<[) 1)) (I(()(())=))ll + (1(())()))))))))))
(16)
Based [4) by applying Eq. [6). Eq. 22) can be reformulated as:
ELBO(#; q, d) log pa pa(2)) (z) (g (pr)lll) 4x)) (7 (1k)
(17)
By plugging Eq. [17 into Eq. 13 we have:
Lg(9, d) Epan(()) (x) (x()) (L (|)
(r 1)KL((
Epul,
(r 1
(18)
Given fixed d, define 9" (d) such that Lg dd) > Ly(g,d) d) holds for all q. i.e,, (d) is
('(d)
the maximizer of Eq. 18). Now consider the function g(v) = ay f exp(Ry), where <bbxp(ka))
('()) a (xy))))) g (w)))
Pdata (x) 0. b p )() >
- -L(g(z, (p)]]la(zz, t|xr) - (7 ( & O. and hyperparamter x 2 1. We
have the following observations:
We ignore the classification loss for simplicity.
13

---

Algorithm Correlation--rienneed Disentanglement
Input:
training dataset D {(,, Ui, a 1i batch size B, encoders (0(2))) and (o(t))) decoder
po(x|2,t), decoy classifier (alz),, variant encoder KL weight y, classification weight Output: 9 (z|22), 4. (f)zz) and po (2) 2, t).
1: Randomly initialize vh, o, 0, and w.
2: while stopping criterion not met do
3: Randomly select batch {o Dr. a:)
4: for to do
5: (Nz, log o < gr (2(1;))
6:
(Mt, log of.) +- (t/ss)
7:
(eparameterize((ue log oo
8:
+ Reparameterizo((4e, log oi,)
9:
â (2(222, (.)
10:
0; + (a)
11:
end for
12:
Dream --
B - (ACE(S;,
1
Las ΣE (a (a;;aa)))) 0z)
13:
14:
KXL - 3 -mn+ j==
Idim(z) (1 + log o² 2d u2. exp(log 2
15:
LIKL 2i==
(1) (1 + log of pl,5 exp(log o?.
16:
Lwaal Lrocmm + aLas + + LIKE
e, 2. 8,0. backpropagate( (Loral)
18:
end while
return (2(2)x), 9o (t)x)) and po (x
If = 0 and a > 0, we have g' (v) 2 0, thus the maximum of g(v) is obtained at 0.
If' b 0 and 0, we have g'vv) = 0, thus 0 yields a maximizer of g(y).
If and a = 0, we have g' (1) < 0, thus there is no maximizer of g
If and a > we have saddle point y" = log f If E (0, 1), we have y" < 0, thus the maximum of g(y) is obtained at y² If Th 1, we have y 0, thus the maximum of g(y) is obtained at O.
We have
= 0 if and only if q(2, t|æ) = pa(z, f|x) and (=(2)) = p(2). Thus, given the assumption Lemma for any fixed d, 9" (d) satisfies q" (d)(2, pa(z, f|æz and g (2(x) p) = p(z). Thus we have ELBO (x: q² (d), d) log Pa(x).
Next, we consider
Lalg"(d), = (
ELBO ( + (ELBO(x; (:)
Epan (xx [logPa(z)) + T)lEp, (log
(L IPa) + Ep... [log Puaaa (x)] - 1/ (pa())))
(19)
Since Epan. [log Pdata Pdaa(x)) (x)] [ is a constant term, we conclude that d' E argmin, [KL (Paaa (2: (z) nH(pa(z))), and (q" d') forms Nash equilibrium point of the system formulated in Eq. [13 and
(x)))=))))
Eq. (4).
B Pseudo codes
This section presents the pseudo-codes for training the proposed CODA framework. The training process of the CODA framework is comprised of two stages:
The pseudo-code for training the correlation-orienteed disentanglement is detailed in Algo- rithm[]
14

---

Algorithm 2 Sample Augmentation
Input: batch size B, number of synthetic samples per instance L, original batch samples {z; 1.9 encoders qv (z|zz and qo (t)=z), decoder pe (z(22,,)) Output: synthesized samples {( )
1: for 1 to B do
2: (Mzy log o². + gu (2((a)
3: (pt, log o + 90 (t
4:
Reparameterize((4u, log 02 +
5: 4 < Reparameterizee(M:; log of
6:
end for
7: for = 1 to B do
8:
for j = to do
9: randomly sample/ h from
10: 3i.5 - Po (0(x)2),44
11: end for
12:
end for
13:
return {10,)
Algorithm 3 Robust Classifier with Reweighted Consistency
Input: training dataset D = {(x;, Ys, (a)) batch size B, number of synthetic samples per instance
L, consistency coefficient A, reweight coefficient, (3, encoders (z(xx) and (o(())) decoder pa (z|2, 1). final classifier fe.
Output: robust classifier fç-
Randomly initialize E.
2: while stopping criterion not met do
3: Randomly select batch {xx, yi, a:) h=)
4: Synthesize samples {(
1B from Algorithm 2
5: for = to B do
6: Oi fe(xz)
7: LCC, + (nc( ( 1/ vi)
based on Eq. 10
8: end for
9: Lets ΣE (y V1)
10: Coonsiscency
E/ LRC,
11: Liotal +- Loss *CConsistency
12: E { backpropagate
(Froal (roual)
13:
end while
14:
return robust classifier fe-
The pseudo-code for training the robust classifier with the proposed reweighted consistency loss in Eq. (TO) is outlined in Algorithm
C Dataset details
Experiments are conducted on the ColoredMNIST and CelebA datasets. Let DGPS denote the degree of group proportion shifts between training and testing sets. Denote pfr and p as the proportions of group g in training and testing sets, respectively. We then define the average group proportion shift (AGPS) Jे Ege|g 10 -p-e In this paper, for [G] we categorize DGPS as high AGPS exceeds 25%, low AGPS below 5%, and medium else. Key statistics of both datasets are summarized in Table
C.1 ColOTedMNIST
The ColoredMNIST dataset is a variant of the classic MNIST dataset handwritten digits. In ColoredMNIST, samples with digit less than 5 are labeled as negative samples (Y = 0), and those
15

---

with digit 5 or higher are labeled as positive samples (Y = 1). Additionally, each sample is colored either red (A 1) or green (A = 0). To introduce a spurious correlation, 85% the positive and negative samples are colored red and green, respectively, in the training set. This bias is adjusted in the validation set 50% for balanced distribution and reversed in the test set to 15% to challenge the learned correlation.
The training and validation sets are derived from the first 50,000 and last 10,000 samples of the original MNIST dataset, respectively. An independent test set of 10,000 samples is also provided. Consistent with the work of [T. label noise is added the dataset, with 25% flip rate.
Characteristically. as shown in Table]] ColoredMNIST presents challenges, such as group imbalance, significant variation in group proportions between training and testing sets, and strong spurious correlations between color and label. These features make it an ideal benchmark for testing the robustness and generalization of machine learning models against subpopulation shifts.
C.2 CelebA
The CelebA dataset is comprised of an extensive collection of celebrity face images, each tagged with 40 binary attributes. The task is to classify the hair color (blond or not-blond), which is highly correlated with gender (male female). We adhere to the standard splits for training (162,770 samples), validation (19,867 samples), and testing (19,962 samples). Like ColoredMNIST, CelebA also exhibits extreme group imbalance and spurious correlations, though with less variation in group proportions between the training and testing sets compared to ColoredMNIST. These characteristics render CelebA formidable dataset for assessing the robustness and generalization capabilities of machine learning algorithms.
D Training details
This section outlines the training parameters and computational details for our experiments. The experiments were conducted on a single NVIDIA L40 GPU.
D.1 Disentangled representation learning in CODA
For the COD training, utilized the Adam optimizer [17] and used a consistent batch size of 128 across 50 epochs on both datasets. Pre-trained ResNet-18 and ResNet-50 encoders were employed for the ColoredMNIST and CelebA datasets, respectively. The learning rates were set to 1e 3 with a corresponding weight decay of le - 3 for the decoy classifier on both datasets. The learning rates for other network components were le 4 for ColoredMNIST and 2e for CelebA, both with zero weight decay. To enhance visual quality in the generated images on CelebA, we applied the Soft-IntroVAI VAE training method (5] with Bmeg 1024 during the final 20 epochs (batch size is reduced to 48). Parameters a 0.1 and 7 = 5 were used in Eq. (9) for both datasets. The training times were approximately hour and 38 minutes with 14 GB of GPU memory for ColoredMNIST. and 45 hours and 29.4 minutes consuming 45.13 GB of GPU memory for CelebA, respectively.
D.2 Robust classifier training details
For the robust classifier, we trained all methods using the SGD optimizer 24 with a momentum of 0.9 and used consistent batch size of 64 across 50 epochs on both datasets. A 3-layer CNN and a pre-trained ResNet-50 were adopted for ColoredMNIST and CelebA, respectively. The learning rate was fixed 1e 5 with weight decay of 1e for CelebA, and both set to Le 3 for ColoredMNIST. The number of augmented samples, denoted by L, was set to 2 for ColoredMNIST and 4 for CelebA. Training times on the ColoredMNIST dataset were as follows: CODA+ERM completed in 26.36 minutes, CODA+RWG in 26.28 minutes, and CODA+GDRO in 27.32 minutes, all with GPU memory consumption of 488 MB. On the CelebA dataset, the training durations were 1 1..30 hours for CODM+ERM, 1 1.46 hours for CODA+RWG, and 19.35 hours for CODA+GDRO, each utilizing 36.35 GB of GPU memory. We conducted hyperparameter tuning for various robust classification approaches through a grid search method:
LIF: We grid searched over E {0.1, 0.3, 0.5, 0.7, 0.9).
16

---

Table Average accuracy (%) for Colored.MINST v2, v3, and v4.
Avg. Acc.
method
ColoredMNIST
ColoredMNIST
ColorsMNNSST
65.96 2.0
66.16
55.02 + 31.00
JTT
72.68
0.56
72.01 0.25
71.34 4: 1.32
ERM
45.69 1: 18.9
10.13 0.01
5.00 1: 0.00
RWG
72.96 0.14
72.62 0.00
72.03 0.76
GDRO
72.92 0.13
72.58 0.29
72.22 0.19
CODA+ERM
72.87
0.21
72.3 = 0.14
72.31 0.06
CODA+RWG
73.28 1: 0.12
73.34 0.16
72.39 0.33
CODA+GDRO
0.11
73.01
0.26
72.98 0.23
Table 7: Performance metries (%) for MultipleColorrdMMISS...
Avg. Acc. (sid)
Worst Acc. (std)
ERM
16.60 1.23
0.00 1: 0.00
RWG
95.54 = 0.24
85.46 0.44
GDRO
91.99 0.33
82.16 4.60
CODA+ERM
97.U6 U.11
91.85 0.68
CODA+RWG
96.94
0.08
91.10 1.16
CODA+GDRO
96.85 = 0.05
90.42 1.39
GDRO:
The group adjustment hyperparameter C' was selected from (1,2,3,4,5).
JTT: A grid search was performed for the identification epoch T over {4,2,5} and for the upsampling factor Aap over {5, 20, 100}.
CODA: The hyperparameter À searched within {10,50, 100, 200, 300, 400, 500} for ColoredMNIST and {10, 50, 100) for CelebA.
E Additional experiments and analysis
E.1 Scalability of CODA to the multi-classffication scenario
Scalability of CODA to the multi-classification scenario is evaluated on dataset called Multiple- ColoredMNIST. In this dataset, the task is to predict the digit y {0,.... 99. where [Y| 10. We define ten RGB colors a E {o..... 9} such that | = 10, resulting in a total of 100 groups. In the training set, each sample is painted with the color corresponding to its digit of 85%, and randomly assigned another for the remaining 15%. Labels are flipped with a probability of 25%, similar to the ColoredMNIST dataset. The training set is highly group imbalanced, with colors spuriously correlated with labels. The majority group (in a random run) constitutes 9.44% of the population, while the minority group only 0.13%.
...,,,
y with a probability
The worst group accuracy of the hypothetical optimal digit classifier f can as low as 62.07% (only 62.07% of the samples whose digits match the labels due to randomness in label flipping and color assignment). Thus, for validation and test sets, we used different setup, no label flipping and uniform color assignment. The worst group accuracy of f is 100% for both sets.
Experimental results are presented in Table 7 CODA demonstrates substantially higher performance compared to benchmark methods, indicating its scalability to scenarios with multiple spurious attribute values. These results underscore the effectiveness of CODA in addressing complex subpopulation shifts.
E.2 Robustness of CODA against different degrees of subpopulation shifts
To further assess the robustness of the CODA framework against different degrees of subpopulation shifts, we conducted series additional experiments. These experiments were designed to simulate extreme subpopulation shifts by introducing color bias into the training and testing datasets with probabilities p and p, respectively. Specifically, we generated three new versions of the ColOredMNIST Bataset-versin 2, 3, and with the probabilities set at 0.8, 0.9, and 0.95,
17

---

Table 8: Performance comparison (%) of different encoders to CODA.
Eicoder Type
Method
Avg. Acc. (sd)
Worst Acc. (sid)
CODA-ERM
08.89 0.59
67.50 + 1.08
3-CNN-layer Encoders
CODA+RWG
70.89 0.31
69.83 + 0.38
CODA+GDRO
71.36 : 0.48
70.51 + 0.52
CODA+ERM
72.20 0.37
71.74 0.21
Resnet 8 Encoders
CODA+RWG
73.20 1::00:00000 0.12
72.11 + 0.51
CODA+GDRO
73.02 0.23
71.98 + 0.57
1
9
2
6
2 7
4
2
6
0
8
1
1
1
1
7
1
1
1
8
7
4
4
4 4
4
4
4
4
4
4
4 2 5
9
9
9
9
9
9
9
9
9
9 9
Z
Z
Z
Z
Z
2
2
Z
2
2 5 7
6
6
6
6
6
6
6
6
S
S
$
2
5
2
$
$
5
2
2
2
2
2
2
2
2
2
2 7 6 4 2 5 7
2 7 6 1
7
7
>
7
7
7
>
7
7
/
/
7
/
/
7
1
1
/
6
6
6
6
6
6
6
6
6 1
2
2
9
9
9
9 6
9 6
9 8 8
1
1
1
1
1
1
8
6
8
6
6
6
6
6
2
222
2
2
2
2
2
2 5
0
0
6
3
3
9
5 7
8
8
8
7
7
7
7
7
7 7 7
7
1
(a) 3-CNN-layye Encoders
(b) Resnet18 Encoders
Figure 5: Visualization of the synthesized samples. Both encoder architectures can well disentangle the causal features from the spurious ones.
respectively. For each version, we maintained fixed validation ratio of 0.5 to ensure consistency in evaluation across different datasets.
The details of the dataset configurations are fully descrihed in Appendix The results from these experiments detailed in Table and Table The findings consistently show that CODA not only enhances the performance of baseline methods but also significantly improves the robustness of the model. Notably, CODA achieved the highest accuracy in the worst-performing groups and the smallest maximum accuracy gap across groups, regardless of the intensity of the subpopulation shifts. These results underline the effectiveness of CODA in handling challenging scenarios where traditional models often falter, further validating its utility in real-world applications with data bias.
E.3
The critical role of the reconstruction quality to robust classification of CODA
The efficacy of CODA relies on its ability synthesize samples with varied spurious attribute values alongside the original sample, thereby encouraging the final classifier to focus exclusively on causal features. This section evaluates the impact of the reconstruction quality robust classification performance.
To facilitate this comparison, simple 3- 3-CNN-layer architecture was employed for encoders, rep- resenting a less sophisticated alternative to the ResNet18 encoders. Evaluation on the Multiple- ColoredMNIST dataset revealed that the 3-CNN-layer architecture produced significantly higher pixel-wise reconstruction losses (0.06451 and 0.05424 for validation and test sets) compared to ResNet18 (0.00055 and 0.00057). Despite this gap in reconstruction quality, both architectures achieved similar disentanglement performance (see Fig 5) successfully separating causal and spuri- features.
Table presents the experimental results. The data indicate that the classification performance of CODA using ResNet18 encoders surpasses that of the 3-CNN-layer encoders. This observation suggests a positive correlation between reconstruction quality and final classification performance, with poorer reconstruction quality corresponding decreased classification accuracy.
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Claims in this paper reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Limitations are discussed in the conclusion part.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
19

---

Justification:
Full assumptions and proof' are provided in Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibiliyy
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Training details are available the Setup section in the main text and Training details section in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
20

---

Answer: [No]
Justification:
Data are publicly and easily available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips.com wttp:ccpppbccccceee guwdes/dde/eeee
swiiiiisisiyycyyyccc for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Training details are available the Setup section in the main text and Training details section in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Error bars are reported for all experimental results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
21

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπхо rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute resources are reported in the Training details section in appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://ss///acce....p//////////a
Answer: [Yes]
Justification: Research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper proposes method for addressing subpopulation shifts, which has no potential negative societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
22

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All related assets are cited in the paper.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
23

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
24


Reviews:
[{"review_id": "NvRAIINdbf", "reviewer": "Reviewer_p3rM", "summary": "This paper proposes CODA (Correlation-Oriented Disentanglement and Augmentation), a novel framework for addressing subpopulation shifts caused by spurious correlations and group imbalance (SC-GI). The key contributions are:\n\n1. A correlation-oriented disentanglement (COD) method that learns to separate variant (spurious) and invariant (causal) features using a bi-branch encoder architecture with a decoy classifier.\n2. A strategic sample augmentation technique that generates synthetic samples by mixing disentangled features, combined with a reweighted consistency loss.\n3. Extensive experiments on ColoredMNIST and CelebA datasets demonstrating CODA's effectiveness in improving worst-group accuracy and reducing maximum group accuracy gap compared to existing robust classification methods.", "questions": "1. While the paper provides some intuition behind the proposed method, could you elaborate on the theoretical understanding of the disentanglement process? How does CODA ensure that the learned representations are truly disentangled?\n2. The experiments focus on datasets with binary spurious attributes. How well does CODA scale to scenarios with multiple or continuous spurious attributes? Have you conducted any experiments in such settings?\n3. The paper shows a sensitivity analysis for the reweighted consistency loss parameter. Are there other critical hyperparameters in CODA? How sensitive is the method to their tuning?\n4. The paper mentions increased computational cost during training. Can you quantify this overhead compared to standard training procedures?", "limitations": "NA", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Novel approach that utilizes spurious attributes constructively rather than treating them as impediments\n2. Theoretically motivated disentanglement objective that avoids complex density estimation\n3. Comprehensive experiments across multiple datasets and varying degrees of subpopulation shifts\n4. Compatibility with existing robust classification methods like ERM, RWG, and GDRO", "weaknesses": "1. Reliance on pre-identified spurious attributes, which may require costly labeling\n2. Increased computational cost during training due to data augmentation processes\n3. Limited theoretical analysis of the proposed method", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful questions and positive evaluation.\n\n**Response to W1:** Thank you for this insightful comment. Labeling spurious attributes (SA) indeed require extra efforts. Yet, we would like to justify its necessity:\n\n1. Robust methods without using SA in training would suffer severe performance drops if SA were not provided in the validation set for model selection [1].\n\n2. There was a significant performance gap between methods with and without considering SA [2].\n\n[1] Liu, Evan Z., et al. \"Just train twice: Improving group robustness without training group information.\" International Conference on Machine Learning. PMLR, 2021.\n\n[2] Yang, Yuzhe, et al. \"Change is hard: A closer look at subpopulation shift.\" arXiv preprint arXiv:2302.12254 (2023).\n\n**Response to W2:** Thank you for your concern on extra computational cost due to Data Augmentation Process (DAP). Although DAP induces an increase of computational costs at training, the long-term benefits of a robust and fair model, such as reducing the frequency of model retraining and corrections, outweigh additional training costs caused by DAP and ultimately save resources. Furthermore, additional cost of CODA scales linearly and is controllable via hyperparameter L. A detailed analysis is offered in our response to your Q4.\n\n**Response to W3:** Thank you for pointing out this limitation. As a pioneering study of the CODA idea, this work aims at showing the empirical effectiveness of CODA. Although preliminary theoretical exploration has been offered in this work, we would like to conduct a deeper theoretical analysis in future studies with directions detailed in our response to Q1 to reinforce the credibility and applicability of CODA.\n\n**Response to Q1:** Thank you for your question. Eq. (9) justifies the disentanglement process (DP). The prediction loss and the mutual information I(x;z) weighted by large γ aim to learn z for predicting SA while minimizing the expressiveness of z on x. The reconstruction loss and the mutual term I(x;t) aim to emphasize more on the contribution of t in reconstructing x, pushing t to better describe x.\n\nThe effectiveness of CODA in DP has been numerically justified in Table 1. Yet, ensuring true disentanglement enquires more theoretical establishment, which is our next research pursuit. Our sketchy plan includes 1) developing a disentanglement measure (DM) tailored for CODA; 2) studying disentanglement performance bounds based on developed DMs; and 3) pursuing disentanglement performance bounds under a more generalized problem setting.\n\n**Response to Q2:** Thank you for this interesting question. For the continuous case, defining groups g= (a, y) can be challenging. One possible approach is to discretize the continuous values by bins, transforming it into a classification task.\n\nWe have evaluated the case of multiple SA values on MultipleColoredMNIST. The task is to predict digits 0-9, and we define ten RGB colors 0-9, resulting in 100 groups. In the training set, each sample is painted with the color a=y with prob 85%, and randomly assigned another for the remaining 15%. Labels are flipped with prob 25%. The training set is highly group imbalanced, with colors spuriously correlated with labels. The majority group constitutes 9.44% of the population, while the minority group only 0.13%. The worst group accuracy of the hypothetical optimal digit classifier f can be as low as 62.07% (only 62.07% of the samples whose digits match the labels due to randomness in label flipping and color assignment). Thus, for validation and test sets, we used a different setup, no label flipping and uniform color assignment. The worst group accuracy of f is 100% for both sets.\n\nExperimental results next indicate that CODA scales well to scenarios with multiple SA values, demonstrating effectiveness in solving SA.\n\n|Method|Avg. Acc. (std)|Worst Acc. (std)|\n|-|-|-|\n|ERM|16.60% (1.23%)|0% (0%)|\n|RWG|95.54% (0.24%)|85.46% (0.44%)|\n|GDRO|94.99% (0.33%)|82.16% (4.60%)|\n|CODA+ERM|97.06% (0.11%)|91.85% (0.68%)|\n|CODA+RWG|96.94% (0.08%)|91.10% (1.16%)|\n|CODA+GDRO|96.85% (0.05%)|90.42% (1.39%)|\n\n**Response to Q3:** Thanks for the insightful question. Another critical hyperparameter is L, which controls the number of synthesized samples for each data point. A larger L introduces more variability to synthesized samples, leading the model to \"forget\" the SA in decision-making. A sensitivity analysis of L on MultipleColoredMNIST is offered next:\n\n|Methods|Avg. Acc. (std)|Worst Acc. (std)|\n|-|-|-|\n|ERM|16.60% (1.23%)|0% (0%)|\n|RWG|95.54% (0.24%)|85.46% (0.44%)|\n|GDRO|94.99% (0.33%)|82.16% (4.60%)|\n|CODA+ERM|||\n|L=1|96.90% (0.29%)|88.78% (0.87%)|\n|L=2|97.12% (0.07%)|90.36% (0.68%)|\n|L=4|97.06% (0.11%)|91.85% (0.68%)|\n|CODA+RWG|||\n|L=1|96.92% (0.04%)|90.87% (0.93%)|\n|L=2|96.66% (0.28%)|91.35% (0.70%)|\n|L=4|96.94% (0.08%)|91.10% (1.16%)|\n|CODA+GDRO|||\n|L=1|96.68% (0.03%)|89.88% (0.50%)|\n|L=2|96.70% (0.16%)|90.10% (1.10%)|\n|L=4|96.85% (0.05%)|90.42% (1.39%)|\n\nIt is observable that increasing L generally improves the worst-group accuracy. However, a larger L also implies higher computational costs. Furthermore, excess L may not introduce more variability (depends on the complexity of the dataset), e.g., L greater than 10 in this case can result in redundant synthesized samples. Therefore, the trade-off between classification performance and computational efficiency should be considered carefully.\n\n**Response to Q4:** The increased computational cost in training the robust classifier is caused by augmented samples. Assume that C is the computational cost per sample. The standard training cost is C * N where N is the number of original samples. The training cost of CODA is then C * (L+1) * N. Here the cost of the RWC loss in Eq. (10) is ignored, as it is much smaller than that of classifying and computing the cross-entropy loss. Thus, the cost increases linearly with L, resulting in an overhead of L times standard training cost.", "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": "p8wSLP9T9s", "reviewer": "Reviewer_gJ1P", "summary": "This paper proposes a novel framework for addressing subpopulation shifts. The framework contains two key components. The first one is Correlation-Oriented Disentanglement, which separates the spurious information from class information. Then, with disentangled spurious and class embeddings, the framework can augment the original dataset with synthesized images and train a robust classifier for improved worst-group performance.", "questions": "(1) Did the authors investigate the relationship between the disentanglement performance and the worst-group classification accuracy? For example, does higher disentanglement performance always lead to more robust worst-group accuracy? If there is a trade-off concerning the quality of the synthesized images due to high KL loss, what would be the best strategy for selecting the hyper-parameter γ (the coefficient for the KL term) for different datasets, especially more complex ones like CelebA?\n\n(2) Did the authors study the diversity and fairness of the learned spurious embeddings? Could it also be biased? For example, could the VAE reconstruct images of the major group better? How could this affect the final classification performance?", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) The idea of translating the research problem from overcoming the effects of spurious attributes to utilizing them, so that the model is encouraged to perform equally well on them, is clever, inspiring, and promising.\n\n(2) The design of the Correlation-Oriented Disentanglement is a good contribution and well-suited for the purpose of the proposed framework.\n\n(3) I enjoyed the writing of this paper very much; it is clear and easy to follow.", "weaknesses": "The major weakness of this paper is its evaluation:\n\n(1) The qualitative results do not seem to support the disentanglement very well. In Figure 3, while the results for MNIST are pretty good, the gender consistency on CelebA is not so obvious, which leads to my question (1).\n\n(2) The paper could be strengthened by evaluating datasets that have more complex spurious information, like the popular Waterbird dataset, which is also used as an example in the introduction. This would better help readers evaluate how well the proposed method addresses complex real-world problems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful feedback and encouraging evaluation.\n\n**Response to W1:**\n\nThank you for your observation regarding the concern about the gender consistency on CelebA presented in Figure 3. The reconstruction results suggest that causal and label-irrelevant features such as background, posture, hairstyle, and hair color are determined by the leftmost samples, while the variance embeddings primarily capture the facial characteristics, which are crucial for determining biological gender in this face dataset. This implies that the variance embeddings carry significant information related to the gender.\n\nOne factor contributing to the less consistent gender representation could be the heavily regularized mutual information term I(x;z), controlled by the hyperparameter γ. The variance embeddings are expected to carry the most essential information about gender. However, if I(x;z) is overly minimized during training time due to a high γ value, it may lead to reduced expressive power of the variance embeddings, thereby affecting the consistency and quality of the reconstructions. Furthermore, CelebA is a significantly more complex dataset compared to ColoredMNIST. The variety of attributes such as age, hairstyle, facial expressions, and lighting conditions adds to the complexity, making disentanglement more challenging.\n\nCareful tuning of the hyperparameter γ is essential to prevent overly minimizing I(x;z) and ensure the variance embeddings retain sufficient expressive power for accurate reconstruction. Conducting experiments to find an optimal balance for γ that maintains the expressive power of the variance embeddings while achieving effective disentanglement could lead to better results.\n\n**Response to W2:**\n\nWe highly appreciate your constructive suggestion of evaluating CODA via additional datasets with more complex spurious information, such as the popular Waterbirds dataset. Yet, given the limited rebuttal time, we decided to first conduct computational experiments to clarify concerns on CODA performances based on datasets considered and we faced insufficient time of conducting a thorough evaluation on additional datasets. However, we acknowledge the importance of your suggestion and we plan to include this evaluation in our future research.\n\n**Response to Q1:**\n\nThank you for the input. Higher disentanglement performance, primarily assessed through visual inspection of reconstructed images, generally leads to improved worst-group classification accuracy. The key is to drive synthesized samples to retain as much information as possible from the original image with spurious attributes excluded.\n\nWhile disentanglement performance is sensitive to the hyperparameter y (the coefficient for the KL term), reconstruction quality is less so. The y-weighted KL loss limits the expressive power of the variance embeddings, but the other KL loss term, being less regularized, helps maintain reconstruction quality, which is crucial for robust classification (see response to Q2).\n\nTo better control reconstruction quality, assigning extra weight to the reconstruction loss term can be considered. Selecting y can start with a higher value if spurious attributes carry minimal information (e.g., colors in the ColoredMNIST dataset) and a lower value if they include more complex patterns. Next y can be adjusted with decreasing it if the variance embeddings capture insufficient variation of spurious attributes and increasing it if they capture more than just the spurious attributes.\n\n**Response to Q2:**\n\nThank you for highlighting the potential bias in the learned spurious embeddings. To address this concern, we conducted an evaluation of the pixel-wise reconstruction loss for each group on both validation and test sets of both datasets. Results are reported in the following table:\n\n||Group 1|Group 2|Group 3|Group 4|\n|-|-|-|-|-|\n|**CelebA**|||||\n|Validation Set|0.0128|0.0119|0.0134|0.0128|\n|Test Set|0.0129|0.0121|0.0128|0.0129|\n|**ColoredMNIST**|||||\n|Validation Set|0.00057|0.00054|0.00057|0.00054|\n|Test Set|0.00056|0.00058|0.00056|0.00056|\n\nThe pixel-wise reconstruction loss is relatively consistent across all groups on both the validation and test sets. While there are minor variations, they do not suggest a significant bias toward any particular group.\n\nHowever, poor reconstruction quality can indeed decrease the final classification performance. We applied a simple 3-CNN-layer architecture for encoders (less powerful than the Resnet18 encoders used in the paper) on the ColoredMNIST dataset, the average pixel-wise reconstruction losses were 0.06451 for the validation set and 0.05424 for the test set, both significantly higher than those of Resnet18 encoders. Notably, the reconstruction images indicate lower reconstruction quality but good disentanglement performance, comparable to that of Resnet18 Encoders. Classification results are presented in the below table:\n\n||Avg. Acc. (std)|Worst. Acc. (std)|\n|-|-|-|\n|**3-CNN-layer Encoders**|||\n|CODA+ERM|68.89% (0.59%)|67.50% (1.08%)|\n|CODA+RWG|70.89 (0.31%)|69.83% (0.38%)|\n|CODA+GDRO|71.36% (0.48%)|70.51% (0.52%)|\n|**Resnet18 Encoders**|||\n|CODA+ERM|72.20% (0.57%)|71.74 (0.24%)|\n|CODA+RWG|73.20% (0.12%)|72.11 (0.51%)|\n|CODA+GDRO |73.02% (0.23%)|71.98 (0.57%)|\n\nWe observed that the classification performance of CODA using ResNet18 encoders is higher than that of the 3-CNN-layer encoders. This implies that poor reconstruction quality does indeed hinder the final classification performance. Building on this insight, we argue that bias towards embeddings of certain groups (i.e., poor reconstruction quality for specific groups) will negatively impact the final classification performance.", "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": "9QVYDJeOW3", "reviewer": "Reviewer_CR4y", "summary": "To accommodate the presence of spurious correlations and group imbalance, the paper introduces a method CODA to reliably recognize the same object across a range of spurious attributes. The paper disentangled invariant and causal features concerning the object identification from spurious features. The performance of the proposed scheme is tested on extensive experiments.", "questions": "Please see Weakness", "limitations": "The authors adequately addressed the limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The presentation of the paper is clear and easy to follow. Problem is well-motivated and each component of the proposed method is clearly justified. The paper introduces an innovated scheme to disentangle the spurious features and causal features related to the predefined spurious attributes. The authors conducted a variety of experiments showing the effectiveness and robustness of the method.", "weaknesses": "1, Compared to disentanglement of the latent variables, the binary classification of the latent space has practically no transfer learning ability and limited for downstream tasks. Please discuss this aspect of the limitation further.\n\n2, The average accuracies in the experimental results sometimes are not as good as the baseline models. In comparison, the worst group accuracy is higher for the proposed method. Please explain this phenomena in details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful comments and positive assessment.\n\n**Response to W1:**\n\nThank you for this interesting comment, which inspires us to think of CODA under transfer learning. To achieve disentanglement, the variance encoder is regularized by an auxiliary label prediction task, which is problem-specific. Therefore, compared with classical machine learning models, in transfer learning, CODA needs to consider the variance encoder-auxiliary label predictor together with the remaining structure for the downstream classifier. Hypothetically, given a source domain and multiple target domains, the CODA can be first trained based on source domain data with source domain spurious attributes identified to produce a CODA backbone model. To transfer such a backbone model in to target domains, the specific target domain spurious attributes need to be first identified to fine-tune the auxiliary label predictor. Next, the target domain data can be supplied to customize the CODA backbone towards the target domain task. This is a quick thought of helping the CODA gain transfer learning ability. Meanwhile, we highly appreciate this great comment because both extending CODA to gain transfer learning ability and innovating transfer learning to accommodate CODA are interesting and valuable problems which will be considered in our future research. Meanwhile, we will update our limitation discussions on this item in the revised paper.\n\n**Response to W2:**\n\nThank you for this insightful comment. We appreciate the opportunity of clarifying the observed phenomenon regarding the average accuracies and the worst group accuracies in our experimental results.\n\nWe acknowledge that, on CelebA, the CODA method in terms of accuracy might not be always higher that of the baseline models (see Table 3). However, on ColoredMNIST and its variants, our method achieves higher average accuracy compared to the baseline models (see Tables 3 and 5 in Appendix). More importantly, our method consistently achieves higher worst group accuracy across all datasets than baseline models (see Table 3 and Table 4). This trade-off between average accuracy and worst group accuracy is a well-documented phenomenon in the literature on subpopulation shifts, which commonly considered the worst-group accuracy as the golden standard while average accuracy as a reference in evaluation.\n\nCODA performances can be interpreted from following two aspects:\n\n1. Modeling Focus and Learning Mechanism:  \n   The primary focus of our work is to develop a robust classification model that performs well across all subpopulations, particularly those that are underrepresented. Traditional models often optimize for overall average accuracy, which can lead to suboptimal performance on minority groups. In contrast, baseline methods such as RWG and GDRO in the paper apply techniques that upweight the sampling probabilities of minority groups (while inevitably compromising the sampling probabilities of majority groups). This drives the model to focus more on minority groups, resulting in higher worst-group accuracy, but potentially lower average accuracy. Our proposed learning mechanism may further emphasize minority groups to ensure no subpopulation is disproportionately disadvantaged.\n\n2. Dataset Characteristics and Group Proportion Shifts:  \n   According to the dataset statistics in Table 1, the degree of group proportion shifts (DGPS) between training and testing sets is low for CelebA and high for ColoredMNIST and its variants. In ColoredMNIST and its variants, the majority/minority groups in the training set become the minority/majority groups in the test set. While in CelebA, the majority/minority groups in the training set remain the majority/minority groups in the testing set. DGPS and the insight from point 1 explain the lower average accuracy of our method on CelebA. On the other hand, high DGPS has less impact on worst group accuracy since group-wise accuracy is not sensitive to changes in sample size.\n\nIn summary, the higher worst group accuracy of CODA, despite the occasional lower average accuracy, aligns with our goal of developing a robust and fair classification model. This trade-off is justified and beneficial for ensuring equitable performance across all subpopulations, especially in critical applications like healthcare, where fairness and robustness are paramount, and where the cost of misclassification is high for certain groups.\n\nThank you once again for this valuable feedback.", "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:

Meta-review: This paper proposes a novel method for addressing subpopulation shifts, which contains a few major components: Correlation-Oriented Disentanglement which separates the spurious information from class information; augmenting the original dataset with new synthesized data from disentangled spurious and class embeddings; training a robust classifier for improved worst-group performance. Extensive experiments on ColoredMNIST and CelebA demonstrate the effectiveness of proposed method, in improving worst-group accuracy and reducing maximum group accuracy gap.
Decision: Accept (poster)