Chroma-VAE: Mitigating Shortcut Learning with Generative Classifiers
Authors: Wanqian Yang
Polina Kirichenko
Micah Goldblum
Andrew Gordon Wilson
Venue: NeurIPS 2022 Accept (2022)
Keywords: deep generative model
variational autoencoder
generative classifier
shortcut learning
spurious correlation
simplicity bias

Abstract: Deep neural networks are susceptible to shortcut learning, using simple features to achieve low training loss without discovering essential semantic structure. Contrary to prior belief, we show that generative models alone are not sufficient to prevent shortcut learning, despite an incentive to recover a more comprehensive representation of the data than discriminative approaches. However, we observe that shortcuts are preferentially encoded with minimal information, a fact that generative models can exploit to mitigate shortcut learning. In particular, we propose Chroma-VAE, a two-pronged approach where a VAE classifier is initially trained to isolate the shortcut in a small latent subspace, allowing a secondary classifier to be trained on the complementary, shortcut-free latent subspace. In addition to demonstrating the efficacy of Chroma-VAE on benchmark and real-world shortcut learning tasks, our work highlights the potential for manipulating the latent space of generative classifiers to isolate or interpret specific correlations.

Chroma-VAE: Mitigating Shorteut Learning with Generative Classifiers
Wanqian Yang
Polina Kirichenko
Micah Goldblum
Andrew Gordon Wilson
New York University
{wanqian, pk1822, goldblum) edu, andreugv@cims nyu edu
Abstract
Deep neural networks are susceptible shorteut learning, using simple features achieve low training loss without discovering essential semantic structure. Contrary to prior belief, we show that generative models alone are not sufficient to prevent shortcut learning, despite an incentive to recover more comprehensive representa- tion of the data than discriminative approaches. However, we observe that shortcuts are preferentially encoded with minimal information, a fact that generative models can exploit to mitigate shortcut learning. In particular, we propose Chroma- VAE, two-pronged approach where VAE classifier is initially trained isolate the shortcut in a small latent subspace, allowing a secondary classifier to be trained the complementary, shortcut-fres free latent subspace. In addition demonstrating the efficacy of Chroma-VAL on benchmark and real-world shortcut learning tasks, our work highlights the potential for manipulating the latent space of generative classifiers to isolate or interpret specific correlations.
Introduction
As science fiction writer Robert Heinlein quipped in his 1973 novel Time Enough for Love, progress is made not by "early risers", but instead by "lazy men trying to find casier ways to do something" Indeed, we can accuse modern machine learning models of emulating the same behaviour. There are notable examples of models that wind up learning the wrong things. For example, images of cows standing on anything other than grass fields are commonly misclassified because the combination of cows and grass fields is so prevalent training data that the model simply Iearns to rely the background as a predictive signal [4]. more concerning example involves predicting pncumonia from chest X-ray scans. Hospitals from which training data is collected have differing rates of diagnosis, a fact that the model easily exploits by learning to detect nospital-specine metal tokens in the scans rather than signals relevant pneumonia itself [52].
Deep neural networks can learn brittle, unintended signals under empirical risk minimization (ERM) [11, 16, 13]. well-known phenomenon observed and studied by various communities. These signals often possess two key attributes: (i) they are spuriously correlated with the label and are therefore strongly predictive [6, 50, 31], despite having no meaningful semantic relationship with the label, and (ii) they are learnt by the neural network as a result of its inductive biases [43, 10]. A recent unifying effort by Geirhos et al. [10] coins the term shortcut to describe such signal. Networks that learn shortcuts fail to generalize to relevant or challenging distribution shifts.
Prior work has sought to alleviate this problem under various formalizations, most commonly in the settings of group robustness [e.g. 39, 28, 7. 25, 20] or adversarial robustness je.g. 45]. In this paper, we motivate different approach to mitigate undesirable shortcuts, where we seek instead to learn shortcut-invariant representation of the data - that is, "everything but the shortcut".
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

i
dassifirr
partal
Lttnn sutepacc 2j Ioans sharica Ioean
inpulcc
shartuut osler (inalyslind
bas lawek
recmubudionn
q
Letent space
Latat sutsppce mokr-avrrant only cartans digit reprssstrttinionnn
Figure I: A visual representation of Chroma- VAE. The latent space is partitioned into two subspaces, one of which is incentivized encode the shortcut under ERM. The complementary subspace learns a shortcut-invarian representation.
To this end, we first observe empirically that deep neural networks preferentially encode shortcuts as they are the most efficient compression of the data that strongly predictive of training labels. This observation helps to explain why shorteut learning is so prevalent amongst discriminative classifiers.
However, this same preference he exploited in generative classifier learn shortcut-invariant compression. Specifically, we consider the Variational latent representations of the data. Our key insight to back-propagate the classifier's loss through latent subspace, while reconstructing with the entire latent space. The model minimizes classification error by encoding the shortcut in this subspace. Since the shortcut representation is isolated, the complementary subspace is free to learn a shortcut-invariant representation through reconstruction. The final classifier then trained on this invariant representation. We dub our approach Chroma- VAE, inspired by the technique of chromatography, which separates the components of a chemical mixture travelling through the mobile phase. Figure summarizes our approach. We note that our pipeline is not in principle restricted to VAEs, and could be used in conjunction with other families of generative models to compartmentalize a latent representation into a shortcut and semantic structure.
(VAE) [23], capable of learning
Our contributions are as follows:
We show empirically that: (i) deep neural networks preferentially encode shortcuts as the most efficient compression (Section 3.1), and (ii) the shortcut representation can be sequestered in a latent subspace of a VAE classifier (Section 3.2).
These observations allow us to propose Chroma- (Section 3.3), an approach designed learn classifier shorteut-ines innvaratninv representation of the data. We demonstrate the effectiveness of Chroma- VAE on several benchmark datasets (Section 5).
Our code is publicly available at https: htttpp://ggttub..co////////// http.////aaaaaaaa/cccraoaaaaaeea public.
2 Background and Notation
Even though shortcuts present in many domains within deep learning 50, 2, 34, 30, 12), restrict our discussion to image classification tasks, as (i) VAEs are widely applied to modeling image distributions, and (ii) shortcut learning is ubiquitous in the vision domain.
Let D,, = {x,, y denote i.idd training data. A variational auto-encoder (VAE) [23] (Eg, Dp) models with latent variable: p(x, z) p(z)p(x(z). The prior p(z) is typically A/ and the likelihood p(x|z) is implicitly modeled by the decoder network Dz(z) (zo diag (o)(2)
as (x) diag (z)(z)) Maximizing p(x) directly is intractable due the integral Instead, we use an encoder network Ex(x) (xx()), diag (o²(x)) as an (amortized) variational approximation go (z) (x), diag(o? and maximize the Evidence Lower BOund (ELBO):
log px x)
ELBO d,x) x) Ezz-ga
g0
by computing unbiased estimates of its gradients.

---

(b)
(c)
(a)
Figure 2: (Left) (a) Examples of reconstructions by decoder trained independently on the dis- criminative model's hidden layer. The shortcut (blue square) is clearly visible when it is present. Reconstruction particularly poor on images with the shorteut the reconstructed faces seem to collapse to the same archetype that does not resemble the original images. (Right) Grad-CAM heatmaps on a CelebA test image. (b) The original image with the shortcut patch. (c)-(e) Averaged activations across all training epochs on three separate models, with a bottlenecked hidden layer of sizes 32, and 128 respectively. The smaller the bottleneck, the more likely is that the model only focuses on the shortcut patch. With more capacity, the model focuses on other regions of the image.
In hybrid VAE classifier (Es, Dg, Co), we model Dir as p(z, y) p(z)p(x(zz)p(//s p(y(x) is modeled by classifier network Cpp that takes the latent mean (4g (x) as input and outputs class probabilities, i.e. Py,o(v(x) = CE y. Cip (Ha(x))), where CE is the cross-entropy loss. Again, maximizing p(x, y) directly is intractable and the training objective becomes:
0.0.).))= =
Σ ELBO( o x) + X log Po.
(2)
(x.
where À is scalar multiplier widely used in practice to account for the fact that p(x) is typically magnitudes smaller than p(yl x) on high-dimensional image datasets.
Group Terminology Let E S denote the shortcut (spurious) label. The group of an input x is g (s,y). v). the combination of its shortcut and true labels. Majority groups refer examples with the dominant correlation between and on the training data, while minority groups refer to the small number of examples with the opposite correlation that ERM models typically misclassify. Group robustness refers to the ability of models to generalize to distribution shifts where shortcuts are no longer predictive. Methods improve group robustness can make use of group annotations; however, our method does not assume access to group labels at train or test time.
3 Chroma- VAE: Separating Shortcuts Generatively
The key intuition behind our approach is to exploit shortcut learning for representation learning, by using a VAE to learn invtriatinvariantt representation of the data. Two central ideas, supported by empirical observations, motivate our method: (i) deep neural networks preferentially encode shorteuts under finite and limited representation capacity, and (ii) the shortcut representation can be sequestered latent subspace of the VAE when jointly trained with a classifier.
3.1 Shortcuts are preferentially encoded under limited representation capacity
Why do shortcuts exist? As Geirhos et al. [10] note, shortcuts arise because the model's inductive biases (the total of interactions between training objective, optimization algorithm, architecture, and dataset) favour learning certain patterns over others; e.g., convolutional neural networks (CNNs) prefer texture over global shape [3, 9].
We consider shortcuts from an information- iioomatiin-thhoreticc perspective, Deep neural networks are commonly thought of as representation learners that optimize the information bottleneck (IB) trade-off, i.e. they aim to learn a maximally compressed representation (minimally sufficient statistic) that fits the labels well [47, 42]. Our key empirical observation is that in datasets where shorteuts exist, they often efficiently compressed, and the compressed information is predictive of labels. As such, deep models preferentially encode shortcuts, especially under limited representation capacity.
We consider an experiment on the CelebA dataset [29], where the task is predicting hair color ("blonde" or not). We inject synthetic shortcut in the form of a 10 10 blue "patch", superimposed

---

(d)
(a)
(b)
(c)
Figure 3: (a) Original image containing shortcut. (b) Grad-CAM with pa (x) gradients. (c) Grad- CAM with (x(x)2 gradients. (d) Partial reconstructions X1- (e) Partial reconstructions X2- While no patches are present in samples of X2 here, we have observed that they can occur with small probability, Conversely, across our many repeated samplings, the shortcut patch is always present samples of x1. Note also that x samples contain a variety of hair colors, whereas X2 samples faithfully reconstruct the celebrity's blonde hair.
onto the image with 0.9 probability for the positive class (see Figure 2b). We first train a discriminative classifier on the labels, and then independently learn decoder on the (fixed) hidden layer of the trained discriminative model.
Figure 2a shows the reconstructions of the decoder. The shortcut is accurately replicated in sharp detail where it exists, but other regions of the reconstructions In particular, note that reconstructions are even less accurate for examples of the positive class (with the shortcut), where mode collapse has happened. Since the decoder was learned separately, this observation implies that the classifier's latent compression primarily encoded for the shortcut's presence.
not as faithful to the originals.
This result is further supported by visualizing the activated regions of the input images during training. We implement Grad-CAld [41], which computes linear combination of activation maps weighted by the gradients of output logits (or any upstream parameters), to produce a heatmap superimposed on the original image, showing regions of the image with the greatest positive contributions to the network's 's prediction.
In Figures 2b-e, visualized these heatmaps on models with bottleneck hidden layer of different sizes. As we can see, the smaller the bottleneck size, the stronger the activations on the shortcut patch. With additional model capacity, the model is more likely to encode other relevant (or spurious) features in addition to the shorteut itself.
From these experiments, we see that deep neural networks tend to compress input data in a shortcut- preserving manner when such shortcuts are both (i) compressible with little information and (ii) predictive of training labels. This tendency is exacerbated under limited latent capacity and is the reason why discriminative classifiers are susceptible shortcut learning under ERM.
3.2
VAE classifiers can sequester shortcuts into specific latent dimensions
What presents problem for diseriminative modeling be exploited into advantage using a generative model. The key insight here is that we can train a VAE classifier where the latent space partitioned into two disjoint subspaces where only one subspace feeds into classifier and is backpropagated through using labels. This subspace encodes the shortcut precisely because this information maximizes the classifier's performance (i.e. simultaneously exhibits high mutual information with the labels and small mutual information with the data). The remaining subspace therefore encodes a partial representation of the image without the shortcut.
Formally, we modify the standard VAE classifier described in Section as follows: we partition the latent space
(21, 22). The classifier Cip uses only 2₁ as input, i.e. for given data point x, the predicted output C.- (x)) . As such, 21 is used for both reconstruction and classification, whereas Z₂ is used for reconstruction only.
In Figure 3, we visualize the results of this experiment in two ways: (i) Grad-C.a -rad,CCMM, where we use gradient weights both of Mo (x), and of μa (x)2. and (ii) sampling partial reconstructions from the VAE: given encoding Hs (x), we replace each half px(x)) with standard Gaussian noise and

---

Algorithm Chroma-VAL
Input: data Der, VAE classifier (Es. Do, Cip), zz-classifier ((2 optimization hyperparameters I/ Training initial VAE classifier initialize (0, db sp)
using Adam, minimize Equation (2), noting that Cpp only takes He (x)1 as input Training Z2 -2-clsssif
initialize xp
compute the (fixed) imput Hg (x)2 using Es
using Adam, minimize L(y') CE CE(3. C (1)(x)2))
return C (2)
sample from the decoder:
81 Do
Do (Ius(x)) I)(O.))))
(3)
82 Do Do (1V(O) [=(x)2]
We use the same predictive task and synthetic shortcut described in Section 3.1. As can be seen from Figures 3b-c, Grad-CAM shows that (ta [xa(x) is most sensitive to pixels around the blue patch, whereas Mg (x)2 sensitive to pixels spread out other regions of the image (including certain regions of the celebrity's blonde hair). The representation of the patch is largely isolated in Z1.
Partial reconstructions in Figures 3d-e support this observation. In Figure 3d, where μu (x)) is fixed and we sample only 22 N NO, I), the decoded samples X represent a variety of faces that do not resemble the original image. However, all samples contain the shortcut patch in the bottom-right right corner. Conversely, in Figure 3e where (x(x)2 is fixed, the reconstructed samples X, greatly resemble the original image, but may or may not contain the shortcut. In other words, H (x)x strongly correlated with the shortcut's 's presence, whereas pa (x) is incorrelated with the shortcut.
Together, Grad-CAM and the partial reconstructions suggest that the latent representation of the shortcut patch has largely been sequestered into 21. This result is intuitive - as the classification loss on C₂ only back-propagates through Zj, it will learn representations that most useful for prediction. This representation is dominated by the shortcut patch, as we showed in Section 3.1. Since the model is also simultaneously learning reconstruct the image, most other information describing the image is partitioned into Z2, resulting in &₂ samples being very similar to the original.
3.3 Chroma-
These results suggest simple approach: we only need to train separate, secondary classifier ((22
(
on Z2 after the initial VAE classifier (Eg, Do. C.p) has been trained, i.c. Z₃ is fixed input into that is not back- pock-propagated through x. Section 3.2 shows that the initial VAE classifier learns Z2 in way that minimizes shortcut representation but contains other salient features (by virtue of Z2 learning to reconstruct). As such, c( will be shorteut-invariant while remaining predictive.
We name this approach Chroma-VALe We provide full description of the training procedure in Algorithm and corresponding diagram in Appendix A. From here on, we will refer to Cpp as the -classifier (zi- not used for prediction) and c(( as the -ccassiif (Z2-clf; shortcut-invariant classifier used for prediction).
Chroma-VAF enables high flexibility with respect the zæ-classifier. Since Z2 is latent vector, the 2y-classifier need not be a deep neural net. Indeed, we will see in our experiments that simpler models perform better smaller datasets. Furthermore, instead of training
like k-nearest neighbors
c(22) on Z2, we can generate partial reconstructions X₂ and train a classifier directly in X'sspace. This procedure allows us to exploit deep models since the inputs are now images. We term this latter approach the X classiiiff (xx-
5

---

Hyperparameters and Group Labels Our method has two hyperparameters: the dimensionality dim(zı)
(0, 1), which controls the relative sizes of z1 and zg. Compared to existing work, Chroma- VAE does not require training group labels. While validation group labels are technically necessary for hyperparameter tuning, we empirically found that it is possible to tune Chroma-VAL without needing them, for two key reasons: (a) In most domains, where the shortcut is known priori (e.g. image background), we can generate and visually inspect partial reconstructions for hyperparameter tuning, by selecting dim(z) to ensure good reconstruction, and 2p to confirm that the shortcut has been isolated in 21. (b) Even where the shortcut is unknown, worst-group accuracy is far less sensitive to Chroma- VAE hyperparameters than for existing methods such as [28]. This insensitivity is likely because most small values of Zp typically suffice to isolate the shortcut z Table summarize the group label requirements.
of z, dim(z), and the partition fraction, 2p dim(z:)+ddimzz)
Example
Group-DRO [38]
JTT [28]
Chroma-VAB
Group Labels on Train Data Group Labels on Validation Data
x
X Vix
Table Comparison of group label requirements. While Chroma- VAE can be tuned using validation worst-group accuracy, we found that partial reconstructions and average accuracy typically suffice.
4 Related Work
Generative Classifiers Depending on how decompose pe px y). generative classifiers can be divided into two categories: (i) class-conditionl models po(y) or (ii) hybrid models po(x)po (y)p(y))...rr)))))) x). The majority of previous work on deep generative classifiers [e.g. 40, 54] focus on the former category, e.g. for semi-supervised learning [24, 19] or adversarial robustness [40]. The most notable work in the latter category is Nalisnick al. [32], which applies hybrid models OOD detection and semi-supervised learning. Our work uses VAEs as the deep generative component, since we require explicit latent representation. We only focus on hybrid models we desire single model that learns the shortcut representation across all classes.
Group Robustness Out-of-distribude (OOD) generalization is a broad area of study [e.g. 8, 14, 37], depending on what assumptions are made on the relationship between train and test distributions as well as the information known at train time. In this work, we primarily consider distribution shifts arising from the presence of shortcuts or spurious correlations. where signals predictive in the train distribution are no longer correlated to the label in the test distribution. Prior work has generally approached this from the group robustness perspective, where the objective is to maximize worst-group accuracy (typically the minority groups) while retaining strong average accuracy. To the best of our knowledge, Chroma- VAAE is the first method to tackle shorteut learning from a supervised representation learning approach using generative classifiers.
Methods in this space can be distinguished by the assumptions they make. Some work rely on having group labels for training data [1, 38, 53, 39], for example, Sagawa et al. [38] optimize worst-group accuracy directly. However, as group annotations can be expensive to acquire, other approaches relax this requirement [28, 7, 44, 51, 33J. For example, Liu et al. [28], which we compare to, treat misclassified examples by an initial model as a proxy for minority groups; these samples are up- weighted when training the final model. However, their approach is brittle to hyperparameter choices, requiring group labels on validation data (from the test distribution) for hyperparameter tuning.
5 Experiments
In Section 5.1, we present results on the ColoreaMNIST benchmark, oroof-of-concept which we use to highlight some key observations and comparisons. Section 5.2, we apply Chroma- VAE to two large-scale benchmark datasets (CelebA and MMIST-FashionMIISS Dominoes), as well as real-world problem involving pneumonia prediction using chest X-ray scans. Appendix C.I contains further results the CelebA synthetic patch -rof--f--ooneept that was presented earlier in Section 3.

---

Method
Dnn.
D Dout
Theoretical UB
75
75
Invariant
60.8
65.6
6
3
3
Naive-Class
87.8
17.8
Naive- Naive-VAE-Clan VAE-Class
89.0
13.2
Naive-Independent
89.7
11.4
JTT
63.2
63.8
(a)
Chroma- VAE (2) cll)
89.0
14.5
Chroma-VAE (z2-clf)
72.5
72.4
3.5
Absolute Differance of (.b) for Dy and D/
2)
Table 2: Predictive accuracies on Din and Dout- Chroma-VAL (Zg-clf) is the best-performing classifier. Our strongest zy-classifier uses k-nearest neighbor classifier (k = VN) rather than MLP. Note that the best that classifier rely- ing only digit do is 75% since Pa = 0.75 by construction.
3.0
2.5
1.-
1.0
0.5
0.0
10
15
20
30
(c)
Figure 4: (a) Original image. (b) Partial reconstructions &1 (top row) and &2 (bottom row). All samples of 1 preserve the color of the original image but not the digit. Conversely, samples of xo preserve digit but can be variably green or red. (c) Plot of absolute differences between p(x) of Dtr and Di,, (the latent means) against dimension. Dimensions of Z₂ have smaller differences than Z
Detailed experimental setups can be found in Appendix B. In considering baselines, we avoid comparing to methods that rely on group labels [38, 36], multiple training environments [1]. or counterfactual examples [46]. We select the following baselines:
1. Naive discriminative classifier Nlass--Class:: standard ERM classifier
2. Naive VAE classifier (Naive-VAE-Clan VAE-Class: standard VAE classifier
3. Naive VAE + classifier "Waive-Independentl standard VAE is first trained on unlabelled data {x, J.N the classifier is then separately trained on the latent projection {xs(x), y
4. Just Train Twice (JTT) [28]: classifier is trained with limited number of epochs, mis- classified training points are upweighted and a second classifier trained again
5.1 ColOredMNIST
Setup. Following Arjovsky et al. [1], (i) first we binarize MNIST [26] labels as (with digits 0-4 and 5-9 as the two classes), (ii) we then obtain actual labels (used to train and evaluate) by flipping ý with probability Pa 0.25, and finally (iii) we obtain color labels e by flipping with variable probability Pe- is used to color each image green or red. In the training distribution Din. Pe = 0.1, hence is more strongly correlated y than ý is, and color becomes the shortcut. the adversarial OOD test distribution Dentt Pc. 0.9, i.e. every digit is more likely to be shaded the other color. Table 2 shows that the "invariant" classifier (trained on black-and-white images) performs similarly on both D... and Dout. hence degraded performance on D is solely due to shortcut learning.
The main takeaway from our results in Table 2 is that Chroma- VAE vastly outperforms all other methods under the adversarial distribution D Dout- We want to demonstrate that the zy-classifier learning from largely shortcut-invarian representation of the data in attaining this performance. First, we observe that the shortcut is heavily present in Z1, as the Z claassifier performs just as poorly as the naive ERM approaches. Furthermore, show that the shorteut representation isolated in Z1, we can sample and inspect partial reconstructions. We plot one input example in Figure 4b; we provide further examples in Appendix C.2. We observe that samples of x; are color-invariant whereas samples of X2 are digit-invariant, suggesting that the bulk of color representation has been isolated in 21. As yet further evidence, we create D',,, a dataset that identical to the training set Drr except that every image has its color flipped. We then measure (p(x) (1x)(1)) for every dimension of
7

---

3p 1/4
3p
1/2
Ap
3/4
dim(z) =
dim(z)
2
2
2
9
6
dim(z) 16
C
6
7
4
O
6
7
G
dim(z) z)) = 32
6
&
2
Table 3: Partial reconstructions for varying choices of dim(z) and Zp. In each cell, the top row contains X₁ samples and the bottom row contains X2 samples. The original image is a green "6". Increasing 2p relaxes the bottleneck such that X₂ samples less resemble the original image, and increasing dim(z) allows more information for higher fidelity reconstructions. The optimal hyperparameters that produced our results in Table are dim(z) 32 and 2p 4'
Din
D Dout
3 100
49.4
3 = 10
63.1
62.8
69.8
70.3
3 1
72.5
72,4
3 = 0.5
72.9
71.2
Table Test accuracies of the zl-classifier on the same ColoredMNIST setup, where Chroma- VAE trained using different values of 8 in the ELBO objective (B = 1 recovers standard Chroma- VAE). Choosing 3 > leads to degraded performance, worsening with higher values of 3. Choosing 3 < does not have significant impact the test accuracies.
the latent mean, averaged across all inputs x in the dataset. In Figure 4c, observe that the dimensions corresponding 22 have smaller differences than dimensions 2 suggesting color (shortcut) is minimally contained in Z2.
Next, the poor performance of the naive baselines highlight our finding that generative models alone are insufficient for avoiding shortcuts. Neither the VAE classifier (Naive VAE-Class) nor the independent VAE classifier (Naive-Indepent) Independent) improved the ERM model (Naive-Class) While our work motivated VAES specifically from observations about the information bottleneck, we note that the community at-large has hoped [10] that simply having generative component might incentivize the model to learn a comprehensive representation the data-generating factors, and not just the minimum compression necessary for small training loss. Our results show that this is not true. This failure not limited to hybrid models - further results in Section C.I show that class-conditional generative classifiers fare just as poorly.
Ablations for hyperparameters dim(z) and Zp-- Table 3 shows partial reconstructions as function of varying dzm and zp. As we expect, reconstructions are poor when dim(z) is small but improve as latent capacity increases. For sufficiently large dim (z), smaller values of 2p are successful at isolating the shortcut color representation in ensuring that samples of X2 are color-invariant but retain the digit shape. As Zp increases, 2ı learns both color and digit, resulting in samples of x that closely resemble the original image. As such, we will expect the æy-classifier to perform poorly since 22 no longer contains meaningful representations for prediction. These results suggest that hyper- parameters should be tuned by first tuning dim(z) to ensure sufficient capacity for reconstruction.

---

CelebA
MF-Dominoes
Chest X-Ray
Blond/Gender
Attractive/mmiling
Method
Ave
Worst
Ave
Worst
Ave
Worst
Ave
Worst
Naive-Class
64.2
29.4
62.5
21.5
50.4
0.61
85.0
10.3
JTT
65.3
28.3
64.7
30.0
50.6
0.81
64.2
52.3
Chroma-VAE
82.0
54.4
66.9
53.1
78.5
73.8
59.9
57.8
Table 5: Average and worst-group test accuracy for the various large-scale benchmark tasks. Chroma- VAE beats all baselines worst-group accuracies, demonstrating that it is effective at mitigating shortcut learning at scale.
O
(a)
Figure 5: Partial reconstruction samples for the large-scale enchmark datasets. (a) CelebA example. The original example is Attractive but not Smiling. All samples of X (top right) are not smiling, whereas samples of X2 (bottom right) may may not be smiling, suggesting that z2 learns "smile- invariance". (h) (*1" digit), whereas samples of &2 (bottom right) retain the core feature (shirt image) but can have variable digit.
example. All samples of & (top right) contain the shortcut
before tuning Zp to ensure that the shortcut is learnt by z but not the true feature. As we observe in our experiments, small values of Zp typically work well.
Ablations for 3. Our method relates to thread of research aimed at learning (unsupervised)
disentangled latent representations using VAEs. The main intuition behind these methods is enforcing independence between dimensions of the marginal distribution 96 (z) := Ex-po,, I9 (z|x)). One such approach is 3- [15), which sets the hyperparameter 3 >> 1 (the coefficient on the Kullback- Liebler (KL) divergence term (I)) to encourage this independence. Table 4 shows the ablation where we train Chroma-VAR with different values of B. Counterintuitively. increasing B > 1 results in degraded performance on Dout- We postulate that unlike the original B-VAE, the VAE here is trained jointly with supervision. As such, the latent factors in the data (color and digit) are both highly correlated with the label, and therefore with each other, so they cannot be disentangled by 3-VAE.
5.2 Large-Scale Benchmarks
Setup. We consider two benchmark CelebA tasks: (1) predicting Blond where Gender is the spurious feature [38, 28), and (ii) predicting Attractive where Smiling the spurious feature [27]. We also consider the MF MF-Dominoes dataset [35], where input images consist of MNIST digits (0 or 1) concatenated with FashionMMIST objects (coat or dress), with the FashionMNIST object being the true label. For both of these datasets, use the harder setting from Lee al. [27], where the spurious feature is completely correlated with the label at train time (by filtering out the two minority groups). Chroma-VAL is naturally suited for this regime, as complete correlation allows the spurious feature to be simultaneously compressible and highly discriminative.
In addition, as example of a real-life domain where shortcut learning happens, we consider the prediction of pneumonia from frontal chest X-ray scans [52, 36, 21]. The shortcut is the background of the image, which contains artifacts specific to the machine that took the X-ray. As different hospitals use different machines and have differing rates of diagnosis, the background becomes spuriously correlated with the label. There is no benchmark dataset for this problem and different

---

authors have created their own datasets from existing public or private X-ray repositories. Likewise, we make a training set (N 20K) with roughly equal numbers of X-rays from the National Institutes of Health Clinical Center (NIH) dataset [48] and the Chexpert dataset [18], such that 90% of NIH images are from the negative class and 90% of Chexpert images are from the positive class.
Table summarizes the results. On CelebA and MF-Dominoes, Chroma-VAE has the highes average and worst-group accuracies. Figure shows the corresponding partial reconstructions, which visually confirm that the desired invariances have been learnt and the shortcuts correctly isolated. JTT underperforms in this complete correlation setting (no minority groups) as it is unable to leverage misclassification of minority examples to train the final model. This experiment highlights a significant difference between group robustness methods and Chroma- VAE. Methods like JTT perform poorly when there are fewer minority examples, as they rely on these examples to approximate the test distribution and implicitly teach the model what features are most useful. In contrast, Chroma-Van VAE learns explicit representations of the shortcut and true features. Chroma- VAE performs better when the correlation between the shorteut and label is stronger, as the shorteut is then more likely to be the minimal encoding that is predictive of the label.
On the Chest X-ray dataset, Chroma-VAE has the highest worst-group accuracy; however, we note that it suffers from lower average accuracy. The fact that Chroma-VAle improves on worst-group accuracy shows that the model works correctly to avoid shortcut learning by isolating the shortcut in Z1. However, the trade-off against average accuracy also highlights one limitation of our approach, which is the reliance on VAE architecture to model the data distribution well. To the best of our knowledge, we are the first to attempt to model Chest X-ray data with a deep generative model, as existing work for Chest X-ray shorteut learning rely on discriminative models - specifically, pre-trained ResNets [e.g. 36). As the VAE is unlikely to model the X-ray dataset perfectly, some meaningful predictive features will not be well-captured, resulting in poorer performance of the z2-classifier despite being shortcut-free. Reconstruction examples for the Chest X-ray dataset can be found in Appendix C.3.
Discussion
We empirically observe that shortcuts being the most efficient and predictive compression of the data are often preferentially encoded by deep neural networks. Inspired by this result, we propose Chroma-VAL which exploits VAE classifier to learn a latent representation of the image where the shortcut is abstracted away. This representation can be used to train a classifier that generalizes well to shortcut-free distributions. We demonstrate the efficacy of Chroma- VAE on several benchmark shortcut learning tasks.
limitation of our work is the reliance on VAEs to model the underlying data distribution well, which can be challenging for many natural image datasets, Extending Chroma-Vas VAE to incorporate stronger deep generative models, such diffusion- based models, could lead to stronger performance on many real-life datasets. Indeed, in principle our pipeline is not anchored to the VAE, and could be used with other families of deep generative models.
Moreover, many spurious correlations may not be easily described with low information. For example, the spurious feature can be the entire background (such as water or land in the Waterbirds dataset (38)), which can contain a relatively large amount of information. In the future, it would be exciting to generalize Chroma-VAL to learn richer latent representations that compartmentalize different types of features in the inputs, rather than simply segmenting low-information shortcuts from the rest of the image. This outcome could possibly be achieved by introducing priors that explicitly encourage different subspaces of z to correspond to different features.
Acknowledgements
We would like to thank Mare Finzi, Pavel Izmailov, and Wesley Maddox for helpful comments. This research is supported by NSF CAREER UIS-2145499 NSF I-DISRE 193471, NIH R01DA048764- 01A1, NSF 191-19,, NSF 1922658 NRT-HDR: FUTURE Foundations, Translation, and Respon- sibility for Data Science Meta Core Data Science, Google AI Research, Hig Biosciences, Capital One, and Amazon Research Award.
10

---

References
[1] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization, arXiv preprint arXiv: 1107:0990009999 2019.
[2] Marcus A Badgeley, John R Zech, Luke Rayner, Benjamin Glicksberg, Manway Liu, William Gale, Michael V McConnell, Bethany Percha, Thomas M Snyder, and Joel T Dudley. Deep learning predicts hip fracture using confounding patient and healthcare variables. NPJ digital medicine, 2(1): 1-10, 2019.
[3] Nicholas Baker, Hongjing Lu, Gennady Erlikhman, and Philip J Kellman. Deep convolutional networks do not classify based on global object shape. PLoS computational biology, 14(12): e1006613, 2018
[4] Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita, Proceedings of the European conference computer vision (ECCV), pages 456-473, 2018.
[5] Battista Biggio, Igino Corona, Davide Maiorca, Blaine Nelson, Nedim Smdicć, Pavel Laskov, Giorgio Giacinto, and Fabio Roli. Evasion attacks against machine learning at test time. In Joint European conference on machine learning and knowledge discovery in databases, pages 387 402. Springer, 2013.
[6] Joy Buolamwini and Timnit Gebru. Gender shades: commercial gender classification. In Conference on fairness, accountability and transparency, pages 77-91. PMLR, 2018.
Intersectional accuracy disparities in
[7] Elliot Creager, Meen-k Jacobsen, and Richard Zemel. Environment inference for invariant learning. In International Conference on Machine Learning, pages 2189-2200. PMLR, 2021.
[8] Yaroslav Ganin and Victor Lempitsky. Unsupervised domain adaptation by backpropagation. In International conference on machine learning, pages 1180-1189. PMLR, 2015.
[9] Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix Wichmann, and Wieland Brendel. Imagenet-trained enns are biased towards texture; increasing shape bias improves accuracy and robustness. arXiv preprint arXiv: /2231, 2018.
[10] Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix Wichmann. Shorteut learning deep neural networks. Nature Machine Intelligence, 2(1 1::65-673,, 2020.
[11] Patrick J Grother, Patrick Grother, Jonathon Phillips, and George Quinn. Report on the evaluation of 2D still-image face recognition algorithms. Citeseer, 2011.
[12] Suchin Gururangan, Swabha Swayamdipta, Omer Levy, Roy Schwartz, Samuel R Bowman, and Noah Smith. Annotation artifacts in natural language inference data. arXiv preprint 1803. 2018.
[13] 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.
[14] Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint arXiv:1903.122 12261,, 2019.
[15] 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, 2017.
[16) Dirk Hovy and Anders Søgaard. Tagging performance correlates with author age. In Proceedings of the 53rd annual meeting of the Association for Computational Linguistics and the 7th international joint conference on natural language processing ( (volume 2: Short papers), pages
183-488, 2015.

---

[17] Forrest landola, Matt Moskewicz, Sergey Karayev, Ross Girshick, Trevor Darrell, and Kurt Keutzer. Densenet: Implementing efficient convnet descriptor pyramids. arXiv preprint arXiv: 1404. 1869, 2014.
[18] Jeremy Irvin, Pranav Rajpurkar, Michael Ko, Yifan Yu, Silviana Ciurea-Ileus, Chris Chute, Henrik Marklund, Behzad Haghgoo, Robyn Ball, Katie Shpanskaya, et al. Chexpert: A large chest radiograph dataset with uncertainty labels and expert comparison. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 590-597, 2019.
[19] Pavel Izmailov, Polina Kirichenko, Mare Finzi, and Andrew Gordon Wilson. Semi-supervised learning with normalizing flows. In International Conference on Machine Learning, pages 4615-4630. PMLR, 2020.
[20] Pavel Izmailov, Polina Kirichenko, Nate Gruver, and Andrew Gordon Wilson. On feature learning in the presence of spurious correlations. arXiv preprint arXiv:2210.11366: 2022.
[21] Sarah Jabbour, David Fouhey, Ella Kazerooni, Michael W Sjoding, and Jenna Wiens. Deep learning applied to chest x-rays: Exploiting and preventing shortcuts. In Machine Learning for Healthcare Conference, pages 750-782. PMLR, 2020.
[22] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 142.6944, 2014.
[23] Diederik P Kingma and Max Welling, Auto-encoding variational bayes. arXiv preprint 1rXiv:1312.61 14, 2013.
[24] Diederik P Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling, Semi- supervised learning with deep generative models. In Advances in neural information processing systems, pages
2014.
[25] Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations. arXiv preprint arX1v:2204.02937,, 2022.
[26] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 (0:11)2778-2224, 1998.
[27] Yoonho Lee, Huaxiu Yao, and Chelsea Finn. Diversify and disambiguate: Learning from underspecified data. arXiv preprint arXiv:2202.034166 2022.
[28] Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan. Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, pages 6781-6799 PMLR, 2021.
(29) Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Large-scale celebfaces attributes (celeba) dataset. Retrieved August, 15(2018):11, 2018.
[30] Tom McCoy, Ellie Pavlick, and Tal Linzen. Right for the wrong reasons: Diagnosing syn- tactic heuristics in natural language inference. In Proceedings of the 57th Annual Meet- ing of the Association for Computational Linguistics, pages 3428-34224421 3448, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 1/P99 URL attps://aclanm: hwwphoooogyoooog/gP.g wtwo///-////1/0
[31] Mazda Moayeri. Phillip Pope. Yogesh Balaji, and Soheil Feizi. A comprehensive study of image classification model sensitivity to foregrounds, backgrounds, and visual attributes. arXiv preprint arXiv:2001.10766 2022.
[32] Eric Nalisnick, Akihiro Matsukawa, Yee Whye Teh, Dilan Gorur, and Balaji Lakshminarayanan. Hybrid models with deep and invertible features. In International Conference on Machine Learning, pages 3723-4733 PMLR, 2019.
[33] Junhyun Nam, Hyuntak Cha, Sungsoo Ahn, Jaeho Lee, and Jinwoo Shin. Learning from failure: Training debiased classifier from biased classifier. arXiv preprint arXiv:2007.025611 2020.
12

---

[34] Luke Oakden- Rayner, Jared Dunnmon, Gustavo Carneiro, and Christopher Ré. Hidden strat- ification causes clinically meaningful failures in machine learning for medical imaging. In Proceedings of the ACM conference on health, inference, and learning, pages 151-159, 2020.
[35] Matteo Pagliardini, Martin Jaggi, François Fleuret, and Sai Praneeth Karimireddy. Agree to disagree: Diversity through disagreement for better transferability. arXiv preprint arXiv:2202.04414. 2022.
[36] Aahlad Puli, Lily H Zhang, Eric K Oermann, and Rajesh Ranganath. Predictive modeling in the presence of nuisance-induced spurious correlations. arXiv preprint arXiv:2107.00520 2021.
[37] Benjamin Recht, Rehecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do cifar-10 classifiers generalize to 10far-10? arXiv preprint /rxiv: 2018.
[38] Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprini 151ivv/911.077533 2019.
[39] Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterization exacerbates spurious correlations. In International Conference Machine Learning, pages 8346-836 8356. PMLR, 2020.
[40] Lukas Schott, Jonas Rauber, Matthias Bethge, and Wieland Brendel. Towards the first adversar- ially robust neural network model mnist. arXiv preprint 1rXiv:1805.09190 2018.
[41] Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based localization. In Proceedings of the IEEE international conference computer vision, pages 618-622 006,626,220 2017.
[42] Ravid Shwartz-Ziv and Naftali Tishby. Opening the black box of deep neural networks via information. arXiv preprint arXiv: 1703:07000000100 2017.
[43] Fabian H Sinz, Xaq Pitkow, Jacob Reimer, Matthias Bethge, and Andreas S Tolias. Engineering a less artificial intelligence. Neuron, 103(6);967-977 2019.
[44] Nimit S Sohoni, Jared Dunnmon, Geoffrey Angus, Albert Gu, and Christopher Ré. No subclass left behind: Fine-grained robustness in coarse-grained classification problems. arXiv preprint arXiv:2011.12944 2020.
[45] Christian Szegedy, Wojeiech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfel- low, and Rob Fergus. Intriguing properties of neural networks. arXiv preprint 1312:633 2013.
[46] Damien Teney, Ehsan Abbasnedjad, and Anton van den Hengel. Learning what makes difference from counterfactual examples and gradient supervision. In European Conference on Computer Vision, pages 580-599. Springer, 2020.
[47] Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 IEEE Information Theory Workshop (I/W), pages 1-5. IEFE, 2015.
[48] Xiaosong Wang, Yifan Peng. Le Lu, Zhiyong Lu, Mohammadhadi Bagheri, and Ronald M Summers. Chestx-ray8: Hospital chest x-ray database and benchmarks weakly- supervised classification and localization common thorax diseases. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 2097-2106. 2017.
[49] Mika Westerlund. The emergence of deepfake technology: A review. Technology Innovation Management Review, 9(11), 2019.
[50] Kai Xiao, Logan Engstrom, Andrew Ilyas, and Aleksander Madry. Noise or signal: The role of image backgrounds in object recognition. arXiv preprint arXiv:2006.099 09994, 2020.
[51] Yadollah Yaghoobzadeh, Soroush Mehri, Remi Tachet, Timothy Hazen, and Alessandro Sordoni. Increasing robustness to spurious correlations using forgettable examples. arXiv preprint 1rXiv:1911.03861, 2019.
13

---

[52] John R Zech, Marcus A Badgeley, Manway Liu, Anthony B Costa, Joseph J Titano, and Eric Karl Oermann. Variable generalization performance of deep learning model detect pneumonia in chest radiographs: cross-sectional study. PLOS medicine, 15(1 I 101002683, 2018.
[53] Zhilu Zhang and Mert R Sabumcu. Generalized cross entropy loss for training deep neural networks with noisy labels. In 32nd Conference Neural Information Processing Systems (NeurIPS), 2018.
[54] Roland S Zimmermann, Lukas Schott, Yang Song, Benjamin A Dunn, and David A Klindt. Score-based generative classifiers. arXiv preprint arXiv:21 2021.
14

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes] The experiments Sections and support the claims made.
(b) Did you describe the limitations of your work? [Yes] Limitations are addressed in Section 5 and 6.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] Potential societal impact of our work discussed in Appendix D.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material Appendix B describe the experimental details necessary to reproduce our results. We will release our code accompanying the public release of our paper.
URL)? [Yes] Section 5 and
(b) Did you specify all the training details (c.g.. data splits, hyperparameters, how they were chosen)? [Yes] See Appendix B.
(c) Did you report error bars (e.g., with respect to the random seed after running exper- iments multiple times)? [No] We have generally found our results to be robust to repeated random runs.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [No] We note that the computational resources used typical of most research using the same models and datasets.
4. If you using existing assets (e.g., code, data, models) curating/releusing
assets... (a) If your work uses existing assets, did you cite the creators? [Yes] All datasets and models used were cited in Section 5 or Appendix B.
(b) Did you mention the license of the assets? [No] All these datasets and models are well-known and broadly used. We refer readers the original papers for licenses.
URL? [No] (d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A] This only applies to CelebA, discussed in the original paper. (c) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [N/A] This only applies to CelebA, discussed the original paper.
(c) Did you include any new assets either the supplemental material
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
15


Reviews:
[{"review_id": "jPM9i6HfEep", "reviewer": "Reviewer_aZBR", "summary": "The authors presented a method for mitigating shortcuts using generative classifiers. The main idea is to isolate the latent factors corresponding to the shortcut in the VAE latent space and use complementary space for learning the classifier. The underlying intuition is that the shortcuts are easy to compress features, so they will only be present/detectable in a small fraction of the latent space embedding. Experiments show the efficacy of the method on multiple datasets.", "questions": "1. It is not clear why gender would be a shortcut and not the hair color, if both are correlated with the output labels.\n2. In Figure 4. C, the dimensions of z_1 is more than z_2. Shouldn’t it be smaller since z_1 is correlated to shortcuts and thus more compressible? \n3. What does ‘invariant’ means in  Table2.\n4. Why do you think KNN gives better accuracy? \n5. “MLPs with smaller hidden layers have better accuracies, and our top performance is a k-nearest neighbor classifier with k = N^(½).” What does N denote here?\n", "limitations": "Yes, the authors have addressed the limitations and potential negative societal impact in the appendix. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### **Strengths**\n\n1. The paper is well written and easy to follow.\n2. The idea is well-motivated and explained using information bottleneck theory and GradCam. The method separates the shortcut in the latent space of VAE and uses the remaining space to train the classifier invariant to the shortcuts.\n3. The paper is trying to solve an important problem of mitigation shortcuts and spurious correlations.  \n4. The authors presented intensive experiments with detailed analysis. The methods show clear performance gains and robustness to shortcuts.\n\n### **Weakness**\n\nThe paper requires more benchmarks, in particular with causal learning methods. One proper benchmark would be the WILDS dataset [1]. Comparison with other methods such as [2] would be helpful.\n\n\n\n[1] Koh et al.,  WILDS: A Benchmark of in-the-Wild Distribution Shifts\n[2] Lee et al., Diversify and Disambiguate: Learning From Underspecified Data\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "dXkvIyUw_qB", "reviewer": "Reviewer_x5WA", "summary": "This paper proposes a method to mitigate the pervasive shortcut learning issue. The authors first make the observation that deep generative classifiers tend to encode spuriousness when the bottleneck is present. Based on this insight, authors propose a two-pronged approach where a generative classifier is trained while only a partition of the latent space is used for classification, and this part of the latent space absorbs the spuriousness. In the second stage, only the other latent code/representation is used to predict the label.", "questions": "- In the paper, the observation that the presence of the bottleneck affect how much the spuriousness will be picked up. Can this be shown as an ablation for the main experiments that when the z1-clf has larger bottleneck capacity (i.e., varying $z_p$), it carries less spuriousness and thus results in worse OOD performance? ", "limitations": "Yes.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The paper tackles an important long standing problem. The motivation is intuitive and the method is simple and easy for the community to reproduce and build on top.\n- The paper is laid out clearly and easy to follow, the logic flows smoothly\n- Empirical results show large benefit over the chosen baselines\n- The analyses are illuminant\n\nWeaknesses\n- The choice of baselines seems a bit slim. The main baselines are the vanilla ERM and the generative classifier. The only stronger baseline was JTT. The results will be more convincing if put into context w.r.t. other recent work that tries to mitigate spuriousness such as NURD (Out-of-distribution Generalization in the Presence of Nuisance-Induced Spurious Correlations; https://openreview.net/pdf?id=12RoR2o32T).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "D2E4Pg8dIlt", "reviewer": "Reviewer_qobE", "summary": "The authors present an approach that aims to prevent \"generative classifiers\" from learning spurious correlation, which the authors refer to as short-cuts. They achieve this by splitting the latent representation in two, z1 and z2. The latent, z1, is passed to a classifier, while z1 and z2 are used to reconstruct the image. The model is trained using gradients from the classifier and the ELBO loss. The motivation being that z1 should contain the short-cut info, while z2 should not.", "questions": "The authors describe short-cut learning as “using simple features to achieve low training loss”. Surely it is best for the model to learn the simplest explanation since this is more likely to generalise? It would be helpful here to make the connection between “short-cuts” and spurious correlations.\n\nLine 39-40: By “generative classifier” do you mean a classifier trained on representations learned via a generative model? Could you provide some (anecdotal) example of this?\n\nTo what extend is short-cut learning a failure to collect sufficiently diverse and unbiased data?\n\nLine 90-95: In an unsupervised representation learning setting there are no labels, do these short-cuts still appear here? And if not, why not just fine-tune a classifier on top of representations learned without supervision? \n\nFigure 2(a): Reconstructions look poor in all cases. The do **not** appear to be especially bad in the case with the blue square. \n\nFigure 2(b): These results are very interesting and do motivate the problem well. This figure could also be improved by clearly showing an x-axis with the model capacity increasing and a box around the square.\n\nThis approach can only be used to factor out the information specific to the classes you have labels for?\n\nIn Figure 3: what is the class label here? Is it hair colour? Please clarify in the caption.\n\nIn Figure 3(d)-(e) please show the original faces too. This figure does show that this approach is “working”. However, it also suggests that gender is being encoded? Is there bias in the data to blonde women?\n\nTable 2: Is it worth pointing out that the authors expect z1-clf to be poor on D_out? And that this further goes to show that z1 contains the spurious correlations?\n\nIt is not clear to me how using z1 for classification helps to ensure that z2 is better than a representation learned using a VAE without supervision (other than the fact that z1 and z2 share parameters, but to me this does not appear to be sufficient)? Is this achieved because z1 and z2 are needed for reconstruction and the assumption is that the information would not need to be encoded in both z1 and z2? This is very unlikely to hold in most cases. It may hold for a beta-VAE (or other disentangled representation). However, then there may not be sufficient information in z2 for classification? \n\nDoes the approach only work for single label classification?\n\nTable 3 and 4: why is there such high variation in accuracy between groups?\n\nTable 4: Include the dataset information in the caption.\n\n**I am happy to significantly increase my score if the authors can address the following:**\n(1) How does the current setup explicitly prevent \"short-cut\" information being encoded in z2.\n(2) Add a clear diagram with losses and gradients.\n(3) It is highly commendable that the authors include results on real world data and I do not want to discourage this. Can the authors please explain any benefits their approach is having on this dataset or why we should be optimistic about the results.", "limitations": "Limitations are not addressed in the main body of the paper.", "rating": 7, "confidence": 4, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality:\nThis work proposes using a classifier on only a sub-set of the latent representation rather than the whole representation.\n\nQuality:\nIt is very encouraging and highly commendable that the authors show results on a real world dataset, this is excellent! While their approach does much better on worst group accuracy, it still does very poorly on the average accuracy. Can the authors please elaborate more on this and explain why this result is still encouraging? Or if this suggests that their approach is not suitable for real world problems? Why do the other methods perform better?\n\nIt is not clear that there are sufficient pressures to ensure that the short-cut is not also encoded in z2, especially in non-trivial settings. This may explain why the model performs poorly on real world data.\n\nThere is very high variance between accuracy on different groups of data.\n\nClarity:\nAside from some terminology confusions mentioned in the box below, this work is exceptionally well motivated in the introduction.\n\nIt would be very helpful to have a diagram showing the encoding, decoding, classification with losses and gradients. The exact implementation and gradients are not 100% clear.\n\nSignificance:\nThe authors tackle an important and challenging problem. However, they fail to show convincing results on real world data.\n\nMain strengths:\n(1) The work is very well motivated.\n(2) The authors evaluate on real-world data.\n\nMain weaknesses:\n(1) It is not clear how the approach works to mitigate short-cut information being encoded in z2.\n(2) The approach does not work well on real-world data.\n(3) The positive results are on simple datasets.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "BAOomDofKqr", "reviewer": "Reviewer_ef3p", "summary": "This paper addresses the problem of shortcut learning in deep neural networks. First, the paper makes the observation that shortcut features tend to be the ones that are preferentially encoded when a limited amount of information can be encoded in one of the hidden layers of the network. The paper then builds on this observation to propose a two-stage approach for learning a classifier robust to shortcut solutions: first a VAE classifier is learned whereby the classification is only based on a small part of the latent, second an independent classifier is trained on the part of the latent space of the VAE that did not feed into the classifier. The efficacy of the method is demonstrated using several hard versions of ‘shortcut benchmarks’ in which the shortcut features are strongly or completely correlated with the targets.", "questions": "I would like to ask the authors to include evaluations of the proposed Chroma-VAE method on “softer” versions of the benchmarks in which the shortcut feature is only weakly correlated with the target and/or the authors should provide convincing evidence that the Chroma-VAE method can be useful in real-world settings (or discuss clearly why the currently reported real-world experiment has practical value).", "limitations": "If my main concern with this study holds (see above), the limitation that the proposed method depends on the shortcut feature being very strongly correlated with the target should be discussed.\n\nI think the potential negative societal impacts are satisfactorily addressed.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is well written, insightful and a pleasure to read. I am not very familiar with shortcut learning benchmarks, but I had the feeling I was able to follow the entire paper without too much effort.\n\nThe proposed solution for overcoming shortcut learning, the Chroma-VAE, is well motivated, clearly explained and to the best of my knowledge novel and original.\n\nThe performance of the Chroma-VAE is very strong on the toy dataset benchmarks. The paper provides a clear intuition for why the method is able to perform so well, which is backed up by insightful visualizations\n\nIn general I would like to support acceptance of this paper, but I do have one important concern that should be addressed first. My concern is that I suspect that the proposed Chroma-VAE method will only work well in ‘artificial’ situations where the shortcut feature that should be avoided is very strongly correlated with the target. If I’m not mistaken, in general, the part of the VAE’s latent space connected to the classifier should encode all features relevant for the classification task (i.e., both the shortcut feature and the true signal) while the unconnected part of the latent space should encode only irrelevant features. It seems to me that only when the shortcut features are very strongly correlated with the target, some other features with relevance for the classification problem will be left to be encoded in the unconnected part of the latent space. Perhaps by restricting the size of the part of the latent space connected to the classifier it is possible in practice to encourage that mostly the shortcut feature is encoded there, but on softer / realistic benchmarks I do not think this will be straightforward and I expect it might lead to low performance.\nTherefore, to investigate this, I suggest that the authors additionally include experiments on “softer” versions of the shortcut benchmarks in which the shortcut features are less strongly correlated with the target.\n\nThe authors do include one experiment on a “real-world”, X-ray dataset, but the results of the Chroma-VAE method on this experiment are not very convincing, as there is a large drop in average accuracy. I am not familiar with this benchmark, but it seems to me that this is not a result with any practical value. (If I’m mistaken, please correct me. But in what practical situation would we prefer the model with the lower average accuracy?) I therefore think that the claim that “our method is viable on a natural, real-world data set” must be moderated, or else better motivated.\n\nMinor points:\n- The result mentioned in L270-271 seems to be missing from Table 3.\n- It would be good to make the caption of Table 4 clearer (e.g., to include that these are the results for the X-Rays experiment).\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept