Chroma-VAE: Mitigating Shortcut Learning with Generative Classifiers
Anonymous Author(s)
AffiliationAddress email
Abstract
Deep neural networks are susceptible to shortcut learning, using simple features toachieve low training loss without discovering essential semantic structure.
Contraryto prior belief, we show that generative models alone are not sufficient to preventshortcut learning, despite an incentive to recover a more comprehensive representation of the data than discriminative approaches.
However, we observe that shortcutsare preferentially encoded with minimal information, a fact that generative modelscan exploit to mitigate shortcut learning.
In particular, we propose Chroma-VAE,a two-pronged approach where a VAE classifier is initially trained to isolate theshortcut in a small latent subspace, allowing a secondary classifier to be trainedon the complementary, shortcut-free latent subspace.
In addition to demonstratingthe efficacy of Chroma-VAE on benchmark and real-world shortcut learning tasks,our work highlights the potential for manipulating the latent space of generativeclassifiers to isolate or interpret specific correlations
Introduction
As science fiction writer Robert Heinlein quipped in his 1973 novel Time Enough for Love , progressis made not by “early risers”, but instead by “lazy men trying to find easier ways to do something”.
Indeed, we can accuse modern machine learning models of emulating the same behaviour.
There arenotable examples of models that wind up learning the wrong things.
For example, images of cowsstanding on anything other than grass fields are commonly misclassified, because the combinationof cows and grass fields is so prevalent in training data that the model simply learns to rely on thebackground as a predictive signal [4].
A more concerning example involves predicting pneumoniafrom chest X-ray scans.
Hospitals in the training data have differing rates of diagnosis, a fact thatthe model easily exploits by learning to detect hospital-specific metal tokens in the scans rather thansignals relevant to pneumonia itself [50].
Deep neural networks can learn brittle, unintended signals under empirical risk minimization (ERM)[11, 16, 13], a well-known phenomenon observed and studied by various communities.
These signalsoften possess two key attributes: (i) they are spuriously correlated to the label and are thereforestrongly predictive [6, 48, 29], 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 [41, 10].
A recentunifying effort by Geirhos et al.
[10] coins the term shortcut to describe such a signal.
Networks thatlearn shortcuts will fail to generalize to relevant or challenging distribution shifts.
Prior work has sought to alleviate this problem under various formalizations, most commonly in thesettings of group robustness [e.g. 37, 26, 7] or adversarial robustness [e.g. 5, 43].
In this paper, wemotivate a different approach to mitigate undesirable shortcuts, where we seek instead to learn ashortcut-invariant representation of the data — that is, “everything but the shortcut”.
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 is strongly predictive of training labels.
This helps toexplain why shortcut learning is so prevalent amongst discriminative classifiers.
However, this same preference can be exploited in a generative classifier to learn a shortcut-invariantcompression.
Specifically, we consider the Variational Auto-Encoder (VAE) [22], capable of learninglatent representations of the data.
Our key insight is to back-propagate the classifier’s loss through alatent subspace , while reconstructing with the entire latent space .
The model minimizes classificationerror by encoding the shortcut in this subspace.
Since the shortcut representation is isolated, thecomplementary subspace is free to learn a shortcut-invariant representation through reconstruction.
The final classifier is then trained on this invariant representation.
We dub our approach ChromaVAE , inspired by the technique of chromatography, which separates the components of a chemicalmixture travelling through the mobile phase.
Figure 1 summarizes our approach.
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 the VAE classifier (Section 3.2).
• These observations allow us to propose Chroma-VAE (Section 3.3), an approach designedto learn a classifier on a shortcut-invariant representation of the data.
We demonstrate theeffectiveness of Chroma-VAE on several benchmark and real-world datasets (Section 5).
Background and Notation
Even though shortcuts are present in many domains within deep learning [4, 48, 2, 32, 28, 12], werestrict our discussion to image classification tasks, as (i) VAEs are widely applied to modeling imagedistributions, and (ii) shortcut learning is ubiquitous in the vision domain.
Let D tr = { x i , y i } Ni =1 denote i.i.d.
training data.
A variational auto-encoder (VAE) [22] ( E ϕ , D θ )models x with a latent variable: p ( x , z )
= p ( z ) p ( x | z ) .
The prior p ( z ) is typically N ( 0 , I ) and thelikelihood p ( x | z ) is implicitly modeled by the decoder network D θ ( z )
= (cid:0) µ θ ( z ) , diag (cid:0) σ 2 θ ( z ) (cid:1) (cid:1) as
N (cid:0) x | µ θ ( z ) , diag (cid:0) σ 2 θ ( z ) (cid:1) (cid:1) .
Maximizing p ( x ) is intractable due to the integral over z .
Instead, we usean encoder network E ϕ
( x ) = (cid:16) µ ϕ ( x ) , diag (cid:16) σ 2 ϕ ( x ) (cid:17)(cid:17) as an (amortized) variational approximationq ϕ (cid:0) z |
µ ϕ ( x ) , diag ( σ 2 ϕ ( x )) (cid:1) and maximize the Evidence Lower BOund (ELBO):
[Equation]
by computing unbiased estimates of its gradients.
In a hybrid VAE classifier ( E ϕ , D θ , C φ ) , we model D tr as p ( z , x , y )
= p ( z ) p ( x | z ) p ( y | x ) .
p ( y | x )is modeled by a classifier network C φ that takes the latent mean µ ϕ ( x ) as input and outputs class probabilities, i.e. p φ,ϕ ( y | x ) = CE (cid:16) y, C φ (cid:0) µ ϕ ( x ) (cid:1)(cid:17) , where CE is the cross-entropy loss.
Again,maximizing p ( x , y ) is intractable and the training objective is:
[Equation]
where λ is a scalar multiplier widely used in practice to account for the fact that p ( x ) is typicallymagnitudes larger than p ( y | x ) on image datasets.
Group Terminology Let s ∈ S denote the shortcut (spurious) label.
The group of an input x isg =
( s, y ) , the combination of its shortcut and true labels.
Majority groups refer to examples with thedominant correlation between s and y on the training data, while minority groups refer to the smallnumber of examples with the opposite correlation that ERM models typically misclassify.
Grouprobustness refers to the ability of models to generalize to distribution shifts where shortcuts are nolonger predictive.
Methods to improve group robustness can make use of group annotations, however,our method does not assume access to group labels at train or test time.
Chroma-VAE: Separating Shortcuts Generatively
The key intuition behind our approach is to exploit shortcut learning for representation learning, byusing a VAE to learn a shortcut-invariant representation of the data.
Two central ideas, supported byempirical observations, motivate our method: (i) deep neural networks preferentially encode shortcutsunder finite and limited representation capacity, and (ii) the shortcut representation can be sequesteredin a 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 inductivebiases (the sum 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 approach this from an information-theoretic lens.
Deep neural networks are commonly thoughtof as representation learners that optimize the information bottleneck (IB) trade-off, i.e. they aimto learn a maximally compressed representation (minimally sufficient statistic) that fits the labelswell [45, 40].
Our key empirical observation is that in datasets where shortcuts exist, they are oftenefficiently compressed, and the compressed information is predictive of labels .
As such, deepmodels preferentially encode shortcuts, especially under limited representation capacity.
We consider an experiment on the CelebA dataset [27], where the task is predicting hair colour(“blond” or not).
We inject a synthetic shortcut in the form of a 10 × 10 blue “patch”, superimposedonto the image with 0.9 probability for the positive class (see Figure 2b).
We first train a discriminativeclassifier on the labels, and then independently learn a decoder on the (fixed) hidden layer of thetrained discriminative model.
Figure 2a shows the reconstructions of the decoder.
The shortcut is accurately replicated in sharpdetail where it exists, but other regions of the reconstructions are not as faithful to the originals.
In particular, note that reconstructions are even less accurate for examples of the positive class(with the shortcut).
Since the decoder was learned separately, this implies that the classifier’s latentcompression primarily encoded for the shortcut’s presence.
This result is further supported by visualizing the activated regions of the input images duringtraining.
We implement Grad-CAM [39], which computes a linear combination of activation mapsweighted by the gradients of output logits (or any upstream parameters).
This can be visualized as aheatmap superimposed on the original image, showing regions of the image with the greatest positivecontributions to the network’s prediction.
In Figures 2b-e, we visualized these heatmaps on models with a bottleneck hidden layer of differentsizes.
As we can see, the smaller the bottleneck size, the stronger the activations on the shortcutpatch.
With additional model capacity, the model is more likely to encode other relevant (or spurious)features in addition to the shortcut itself.
[Figure]
The smaller the bottleneck, the more likely it is that the model only focuses on the shortcut patch.
With more capacity, the model focuses on other regions of the image.
From these experiments, we see that deep neural networks tend to compress input data in a mannerwhich preserves the shortcut when it is compressible with little information and is simultaneouslypredictive of training labels.
This tendency is exacerbated under limited latent capacity and is thereason why discriminative classifiers are susceptible to shortcut learning under ERM.
3.2 VAE classifiers can sequester shortcuts into specific latent dimensions
What presents a problem for discriminative modeling can be exploited into an advantage using agenerative model.
The key insight here is that we can train a VAE classifier where the latent spaceis partitioned into two disjoint subspaces — however, only one subspace feeds into a classifierand is backpropagated through using labels.
This subspace encodes the shortcut precisely becausethis information maximizes the classifier’s performance (i.e. simultaneously exhibits high mutualinformation with the labels and small mutual information with the data).
The remaining subspacetherefore encodes a partial representation of the image without the shortcut.
Formally, we modify the standard VAE classifier described in Section 2 as follow: we partition thelatent space as z =
( z 1 , z 2 ) .
The classifier C φ uses only z 1 as input, i.e. for a given data point x ,the predicted output is C φ (cid:0) µ ϕ ( x ) 1 (cid:1) .
As such, z 1 is used for both reconstruction and classification,whereas z 2 is used for reconstruction only.
We visualize the results of this experiment in two ways: (i) Grad-CAM, where we use gradientweights both of µ ϕ ( x ) 1 and of µ ϕ ( x ) 2 , and (ii) sampling partial reconstructions ˜ x from the VAE:given an encoding µ ϕ ( x ) , we replace each half µ ϕ ( x ) i with standard Gaussian noise and samplefrom the decoder:
[Equation]
We use the same predictive task and synthetic shortcut as described in Section 3.1.
The visualizationsare displayed in Figure 3.
As can be seen from Figures 3b-c, Grad-CAM shows that µ ϕ ( x ) 1 is mostsensitive to pixels around the blue patch, whereas µ ϕ ( x ) 2 is sensitive to pixels spread out in otherregions of the image (including certain regions of the celebrity’s blonde hair).
This highlights that therepresentation of the patch is largely isolated in z 1 .
Partial reconstructions in Figures 3d-e support this observation.
In Figure 3d, where µ ϕ ( x ) 1 is fixedand we perturb only z 2 ∼ N ( 0 , I ) , the decoded samples ˜ x 1 represent a variety of faces that do notresemble the original image.
However, all samples contain the shortcut patch in the bottom-rightcorner.
Conversely, in Figure 3e where µ ϕ ( x ) 2 is fixed, the reconstructed samples ˜ x 2 greatly resemblethe original image, but may or may not contain the shortcut.
In other words, µ ϕ ( x ) 1 is stronglycorrelated with the shortcut’s presence, whereas µ ϕ ( x ) 2 is uncorrelated with the shortcut.
Together, Grad-CAM and the partial reconstructions suggest that the latent representation of theshortcut patch has largely been sequestered into z 1 .
Intuitively, this is believable — as the classification loss on C φ only back-propagates through z 1 , it will learn representations that are most useful
[Figure]
Algorithm
Input: data D tr , VAE classifier ( E ϕ , D θ , C φ ) , z 2 -classifier C (2) φ ′ , optimization hyperparameters //
Training initial VAE classifier initialize ( θ, ϕ, φ ) using Adam, minimize Equation (2), noting that C φ only takes µ ϕ ( x ) 1 as input //
Training z 2 -classifier initialize φ ′ compute the (fixed) input µ ϕ ( x ) 2 using E ϕ using Adam, minimize L ( φ ′ ) =
CE (cid:16) y, C (2) φ ′ (cid:0) µ ϕ ( x ) 2 (cid:1)(cid:17) return C (2) φ ′ for prediction.
This representation is dominated by the shortcut patch, as we showed in Section 3.1.
Since the model is also simultaneously learning to reconstruct the image, most other informationdescribing the image is partitioned into z 2 , resulting in ˜ x 2 samples being very similar to the original.
3.3 Chroma-VAE
These results suggest a simple approach: we only need to train a separate, secondary classifier C (2)
φ ′on z 2 after the initial VAE classifier ( E ϕ , D θ , C φ ) has been trained, i.e. z 2 is a fixed input into C (2) φ ′that is not back-propagated through.
As shown in Section 3.2, a well-trained VAE classifier learns z 2in a way that minimizes shortcut representation but also contains other predictive features (by virtueof z 2 learning to reconstruct).
As such, C (2) φ ′ will be shortcut-invariant while remaining predictive.
We name this approach Chroma-VAE .
A full description is provided in Algorithm 1.
Ourmethod has two hyperparameters: the dimensionality of z , dim( z ) , and the partition fraction,z p = dim( z 1 ) dim( z 1 )+dim( z 2 ) ∈ (0 , 1) , which controls the relative sizes of z 1 and z 2 .
From here on, we will refer to C φ as the z 1 -classifier ( z 1 - clf ; not used for prediction) and C (2) [-' 
-]φ ′ asthe
z 2 -classifier ( z 2 - clf ; shortcut-invariant classifier used for prediction).
Since z 2 is a vector, thez 2 -classifier need not be a deep neural network.
Indeed, we noticed in our experiments that simplermodels like k -nearest neighbors can sometimes perform better; see Section 5 for more details.˜ x 2 -classifier As we observed in Section 3.2, the partial reconstructions ˜ x 2 are uncorrelated withthe shortcut.
As such, instead of training C (2) φ ′ on z 2 , we can simply generate partial reconstructionsand train a classifier on samples of ˜ x 2 directly in X -space.
This allows us to exploit deep modelssince the inputs are now images.
We term this the ˜ x 2 -classifier ( ˜ x 2 - clf ).
Related Work
Generative Classifiers Depending on how we decompose p θ ( x , y ) , generative classifiers canbe divided into two categories: (i) class-conditional models p θ ( y ) p θ ( x | y ) or (ii) hybrid modelsp θ ( x ) p θ ( y | x ) .
The majority of previous work on deep generative classifiers [e.g. 38, 52] focus onthe former category, e.g. for semi-supervised learning [23, 19] or adversarial robustness [38].
Themost notable work in the latter category is Nalisnick et al.
[30], which applies hybrid models to OODdetection and semi-supervised learning.
Our work uses VAEs as the deep generative component,since we require explicit latent representation.
Due to the need to encode specific predictive features(the shortcut) into the latent space, our work focuses on hybrid models only.
Group Robustness Out-of-distribution (OOD) generalization is a broad area of study [e.g. 8, 14,35], depending on what assumptions are made on the relationship between train and test distributionsas well as the information known at train time.
In this work, we primarily consider distributionshifts arising from the presence of shortcuts or spurious correlations, where signals predictive inthe train distribution are no longer correlated to the label in the test distribution.
Prior work hasgenerally approached this from the group robustness perspective, where the objective is to maximizeworst-group accuracy (typically the minority groups) while retaining strong average accuracy.
To thebest of our knowledge, Chroma-VAE is the first method to tackle shortcut learning from a supervisedrepresentation learning approach using generative classifiers.
Methods in this space can be distinguished by the assumptions they make.
Some work rely on havinggroup labels for training data [1, 36, 51, 37], for example, Sagawa et al.
[36] optimize worst-groupaccuracy directly.
However, as group annotations can be expensive to acquire, other approachesrelax this requirement [26, 7, 42, 49, 31].
For example, Liu et al.
[26], which we compare to, treatmisclassified examples by an initial model as a proxy for minority groups ; these samples are upweighted 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.
Compared to existing work, Chroma-VAE does not require training group labels.
Furthermore, wefound that Chroma-VAE hyperparameters are not sensitive to groups — results in Section 5 weretuned using average accuracy on validation data.
We hypothesize the reason for this robustness: theeffectiveness of Chroma-VAE depends on (i) the shortcut being successfully isolated in z 1 and (ii)
z
2having sufficient capacity to learn the non-shortcut predictive features.
(ii) is achieved by choosingdim( z ) appropriately to ensure good reconstruction — this can be evaluated without group labels.
Given such an appropriate dim( z ) , most non-extreme values of z f will ensure that z 1 has enoughcapacity to encode the shortcut.
As such, the z 2 -classifier will be shortcut-invariant and can be tunedwithout needing group labels.
Table 1 summarize the group label requirements.
[Table]
Experiments
In Section 5.1, we present results on the ColouredMNIST benchmark, a proof-of-concept which wealso use to highlight some key observations and comparisons.
In Section 5.2, we present results on twoadditional benchmark datasets: CelebA (gender/hair colour) and Dominoes ( MNIST - FashionMNIST ).
Finally, in Section 5.3, we demonstrate Chroma-VAE on a real-life problem — namely, pneumoniaprediction using chest X-ray scans.
In addition, Appendix B.1 contains further results on the CelebAsynthetic patch proof-of-concept that was presented earlier in Section 3.
Baselines and Metrics Detailed experimental setups are found in Appendix A. For
ColouredMNIST , we evaluate predictive accuracy on the various test distributions.
For the otherdatasets, we evaluate average and worst-group accuracies.
We consider the following baselines: Method
D in D out [Table]
UB [Table] [Table] [Table]
[Table] [Table] [Table]
[Table] [Table] [Table]
=
[Figure]
Naive discriminative classifier (Naive-Class): ERM classifier2.
Naive VAE classifier (Naive-VAE-Class): standard VAE classifier
Naive VAE + classifier (Naive-Independent): standard VAE is first trained on unlabelleddata { x i } Ni =1 , the classifier is then separately trained on the latent projection { µ ϕ ( x ) , y } Ni =14.
Just Train Twice (JTT) [26]: classifier is trained with a limited number of epochs, misclassified training points are upweighted and a second classifier is trained again
For fairness, we do not compare to methods that use group labels [36, 34], multiple training environments [1], or counterfactual examples [44].
5.1 ColouredMNIST
We follow the same experimental design as Arjovsky et al.
[1]: (i) we first binarize MNIST [24] labelsˆ y (with digits 0-4 and 5-9 as the two classes), (ii) we obtain actual labels y
(used in training andevaluation) by flipping ˆ y with probability p d = 0 .
25 , and finally (iii) we obtain colour labels c byflipping y with variable probability p c , where c is used to colour each image green or red.
In thetraining distribution D in , p c = 0 .
1 , hence colour is the shortcut as c is more strongly correlated to ythan ˆ y is.
In the adversarial OOD test distribution D out , p c = 0 .
9 , i.e. every digit is more likely to beshaded the other colour.
Note from Table 2 that the invariant classifier (trained on black-and-whiteimages) attains similar performance on both D in and D out , confirming that degraded performance on
D out is solely due to shortcut learning.
The main takeaway from our results in Table 2 is that Chroma-VAE vastly outperforms all othermethods under the adversarial distribution D out , suggesting that the z 2 -classifier is learningfrom a largely shortcut-invariant representation of the data.
Accordingly, the z 1 -classifierperforms just as poorly as the naive methods, suggesting the shortcut is heavily present in z 1 .
Thisis further confirmed by the partial reconstructions shown in Figure 4b (further examples can befound in Appendix B).
Samples of ˜ x 1 are colour-invariant whereas samples of ˜ x 2 are digit-invariant,suggesting that the bulk of colour representation has been isolated in z 1 .
In addition, both the z 2 -classifier and ˜ x 2 -classifier perform well.
In particular, we noticed that simplermodels work better as the z 2 -classifier — MLPs with smaller hidden layers have better accuracies,and our top performance is a k -nearest neighbor classifier with k = √ N .
[Table]
[Figure]
As further evidence that the shortcut is represented by z 1 , we create D ′ tr , a dataset that is identical tothe training set D tr except that every image has its colour flipped.
We then measure log | z ( i ) − z ( i ) |for every dimension i of the latent space.
In Figure 4c, observe that the dimensions of z 2 have asmaller difference than dimensions of z 1 , suggesting colour (shortcut) invariance is preserved in z 2 .
Next, this experiment highlights our finding that generative models alone are insufficient.
Neitherthe VAE classifier (Naive-VAE-Class) nor the independent VAE + classifier (Naive-Independent)improved on the ERM model (Naive-Class).
While our work motivated VAEs specifically fromobservations about the information bottleneck, we note that the community at-large has hoped [10]that simply having a generative component might incentivize the model to learn a comprehensiverepresentation of the data-generating factors, and not just the minimum compression necessary forsmall training loss.
Our results show that this is not true.
This failure is not limited to hybrid models— further results in Section B.1 show that class-conditional generative classifiers fare just as poorly.
Detailed comparison with JTT JTT has two hyperparameters: (a) the number of epochs T withwhich the initial classifier was trained, and (b) the factor α by which initially mis-classified examplesare upweighted.
We tested various combinations of these two hyperparameters (see Appendix A)and Figure 4d shows the results of this ablation.
In particular, we can see that while it is possibleto choose hyperparameters such that JTT performs better than Chroma-VAE on one of the two testdistributions, JTT does not outperform Chroma-VAE on both .
5.2 Benchmark Datasets
We consider two benchmark CelebA tasks: (1) predicting Blond where Gender is the spuriousfeature [36, 26], and (ii) predicting Attractive where Smiling is the spurious feature [25].
Wealso consider the MF-Dominoes dataset [33], where input images consist of MNIST digits (0 or 1)concatenated with FashionMNIST objects (coat or dress), with the FashionMNIST object being thetrue label.
In all of these experiments, we use the harder setting from Lee et al.
[25], where thespurious feature is completely correlated with the label at train time (by filtering out the two minoritygroups).
Chroma-VAE is naturally suited for this regime, as complete correlation allows the spuriousfeature to be simultaneously compressible and highly discriminative.
The results are summarized in Table 3.
All results were tuned on average accuracy.
Chroma-VAEoutperforms the naive discriminative model, as well as JTT, on both average and worst-group
[Table]
accuracy.
Similarly, the fact that the z 1 -classifier does worse than a naive discriminative modelsuggests that the shortcut has been preferentially encoded by z 1 .
Partial reconstructions for a testexample from either dataset, shown in Figure 5, confirm these results.
Detailed Comparison with JTT On the CelebA dataset, Table 3 shows that JTT underperformed.
This is primarily due to the fact that these experiments are performed in the complete correlationsetting (no minority groups).
As such, JTT is unable to leverage on misclassification of minorityexamples to train the final model.
This experiment highlights a significant difference between grouprobustness methods and Chroma-VAE.
Group robustness methods like JTT tend to perform morepoorly when there fewer minority examples, as they rely on these examples to teach the model whatto learn.
In contrast, Chroma-VAE leverages representation learning of shortcuts by the classifier.
Stronger correlation between the shortcut and true signals means that it is more likely for the shortcutto be preferentially encoded by the VAE classifier, since the shortcut is now a perfect predictor of thelabel.
As such, Chroma-VAE will do better in learning a shortcut-invariant representation.
5.3 Chest X-Rays
We consider the prediction of pneumonia from frontal chest X-ray scans, a well-known exampleof shortcut learning [50, 34, 20].
The shortcut is the background of the image, which is spuriouslycorrelated with the hospital as different hospitals have different machines, lightning conditions, etc. —resulting in different background artifacts.
This poses a problem for datasets containing X-rays frommultiple hospitals with differing rates of pneumonia diagnosis.
There is no benchmark dataset and different authors have created their own datasets from existingpublic or private collections of X-ray repositories.
Likewise, we make a training set ( N = 20 K) withroughly equal numbers of X-rays from the National Institutes of Health Clinical Center (NIH) dataset[46] 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 4 shows the results.
Examples of partial reconstructions and a qualitative analysis can befound in Appendix B.
The main takeaway from this experiment is that our method is viable ona natural, real-world data set .
Compared to the other approaches, Chroma-VAE has the highestworst-group accuracy.
However, we note that it does suffer from lower average accuracy.
Discussion and Conclusion
We propose Chroma-VAE, inspired by the empirical observation that shortcuts, being the mostefficient and predictive compression of the data, are often preferentially encoded by deep neuralnetworks.
Our model makes use of a generative model to train a second classifier on the latentrepresentation of the image where the shortcut is abstracted away.
We demonstrate the efficacy of
Chroma-VAE on several benchmark and real-world shortcut learning tasks.
Finally, we concludewith a few promising avenues for future work:
Learning diverse ensembles with ≥ 2 classifiers In this work, we only considered separatingtwo sets of features (shortcut and true signal).
However, Chroma-VAE could be extended to settingswhere there are multiple relevant predictive features.
To this end (and truer to the chromatographmetaphor), Chroma-VAE can be extended to iteratively learn a diverse ensemble of predictors.
Prior on z Imposing a prior on z can help to further disentangle the features learnt by z 1 and z 2 .
One suggestion is to enforce an orthogonal loss component on z , i.e. by minimizing ⟨∇ z 1 y, ∇ z 2 y ⟩ . Checklist For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes] The experiments in Section 3 and Section 5 support the
claims we made.
(b) Did you describe the limitations of your work?
[Yes] Limitations are addressed in
experimental findings in Appendix 5.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] Potential
societal impact of our work is discussed in Section C.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
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]
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 or as a URL)?
[Yes] Section 5 and
Appendix A describe the experimental details necessary to reproduce our results.
We
will release our code accompanying the public release of our paper.
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Appendix A.
(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., typeof GPUs, internal cluster, or cloud provider)?
[No] We note that the computational
resources we used are typical of most research using the same models and datasets.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] All datasets andmodels used were cited in Section 5 or Appendix A.
(b) Did you mention the license of the assets?
[No] All these datasets and models are
well-known and broadly used.
We refer readers to the original papers for licenses.
[Table]
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
[Table]
using/curating?
[N/A]
This only applies to CelebA , discussed in the original paper.
[Table]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
[Table]
information or offensive content?
[N/A]
This only applies to CelebA , discussed in the
[Table]
original paper.
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable?
[N/A]
[Table]
(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 to participants and the total amount
spent on participant compensation?
[N/A]
A Experimental Details
Section 3: CelebA (Synthetic Patch)
We use standard CNN architecture.
The encoder contains4 convolutional layers, followed by fully-connected layers to µ ϕ and σ 2 ϕ .
The decoder contains afully-connected layer, followed by 4 deconvolutional layers.
The Adam optimizer [21] is used with alearning rate of 0.0005.
The hyperparameters are dim( z )
= 128 and z p = 0 .
Section 5: ColouredMNIST We use standard CNN architecture.
The encoder contains 4 convolutional layers, followed by fully-connected layers to µ ϕ and σ 2 ϕ .
The decoder contains a fully-connectedlayer, followed by 4 deconvolutional layers.
The Adam optimizer is used with a learning rate of 0.001.
The hyperparameters are dim( z ) =
48 and z p = 0 .
We performed a hyperparameter sweep for
JTT with T ∈ { 1 , 3 , 5 , 10 } and α ∈ { 2 , 5 , 50 , 100 } .
Section 5: CelebA and MF-Dominoes We use standard CNN architecture.
The encoder contains4 convolutional layers, followed by fully-connected layers to µ ϕ and σ 2 ϕ .
The decoder contains afully-connected layer, followed by 4 deconvolutional layers.
The Adam optimizer is used with alearning rate of 0.0005.
For CelebA , the hyperparameters are dim( z )
= 128 and z p = 0 .
For
MF-Dominoes , the hyperparameters are dim( z ) = 8 and z p = 0 .
We performed a hyperparametersweep for JTT with T ∈ { 1 , 3 , 5 , 10 } and α ∈ { 2 , 5 , 50 , 100 } .
Section 5: Chest X-rays
Following prior work [50], we use a Densenet (pretrained on ImageNet)[17] for the encoder, along with a standard CNN decoder (containing a fully-connected layer,followed by 4 deconvolutional layers).
The Adam optimizer is used with a learning rate of 0.0001.
The hyperparameters are dim( z )
= 128 and z p = 0 .
B Additional Results
[Table]
B.1 Synthetic Patch on CelebA
[Table]
We consider the same task as in Section 3.1, where we predict blond hair, with a synthetic shortcut(blue patch superimposed onto the positive class with probability 0.9).
We evaluate performance onthree test environments:
• Training Distribution ( D tr ): the patch is added to all members of the positive class
• Neutral Distribution ( D neut ): no patches are added
• Adversarial Distribution ( D anti ): the patch is added to all members of the negative class
Table 5 summarizes the results.
Chroma-VAE vastly outperforms all other methods under theadversarial distribution .
The z 2 -classifier has the same performance as the invariant classifier onthe in-distribution D tr , implying that it did not benefit (unfairly) from shortcut representations, unlikethe naive baselines.
However, the fact that the z 2 -classifier is still outperformed by the invariantclassifier on D anti and D neut suggests that not all desired features were captured by z 2 .
This implies that the maximally-useful compression of the data in z 1 primarily encodes the shortcut, but also otheruseful correlations.
Figure 6 shows partial reconstruction samples.
The first image represents an example of the negativeclass (non-blonde hair), the second image is an example of the positive class with the shortcut, and thethird image is an example of the positive class without the shortcut (which happens with probability0.1 in the training set).
Similar to Figure 3, samples of ˜ x 1 do not resemble the original image, butalways encode the shortcut patch if it exists in the original image.
Conversely, samples of ˜ x 2 beargreat resemblance to the original image but may or may contain the shortcut — confirming that theyare shortcut-invariant.
Comparison with JTT
While JTT does well on D neut , it performs just as poorly as the naivemethods on D anti .
This highlights a key difference in the approaches of JTT and Chroma-VAE.
JTTrelies on having sufficient examples in the minority group to upweigh and correct the second classifier.
In the case of the adversarial distribution, doing so is not possible because there are simply noexamples in the training set with a negative label and the shortcut present.
Conversely, Chroma-VAErelies on learning a shortcut-invariant representation of the data, and is therefore impervious toadversarial distribution shifts as such.
Ablations on β and relationship to disentanglement Our method relates to the idea of learningdisentangled latent representations.
β -VAE, proposed by Higgins et al.
[15], introduces the hyperparameter β (the coefficient on the Kullback-Liebler (KL) divergence term in (1)) as a means of learningdisentangled representations.
We ablate for this by learning the z 2 -classifier with different valuesof β .
Table 6 shows the results of this ablation.
Increasing β > 1 result in degraded performanceon D anti similar to naive baselines, which seems counterintuitive as further restricting the latentcapacity of the VAE should encourage it to encode the shortcut.
We postulate that this could bebecause increasing β > 1 restricts both z 1 and z 2 , and might instead have encouraged z 2 to encodethe shortcut when it normally would not have.
Method
[Table]
B.2 ColouredMNIST: Partial Reconstructions
Figure 7 shows partial reconstruction samples on three additional images.
We observe the samepatterns as in Figure 4, where samples of ˜ x 1 preserve the colour of the original image but not thedigit, whereas samples of ˜ x 2 preserve the digit of the original shape but is colour-invariant.
[Figure]
Figure 8 shows the reconstruction as well as partial reconstruction samples on a test example.
Similarto partial reconstructions on the CelebA dataset, samples of ˜ x 1 are more diverse and less likely toresemble the original image, whereas samples of ˜ x 2 resemble the original image more, especiallyaround the center region of the image (where true predictive signals for pneumonia are located).
Thisconfirms that z 1 is more likely to have encoded the shortcut (X-ray background) whereas z 2 is morelikely to have retained true predictive signals for pneumonia.
C Broader Societal Impact
Positive Sources of Impact By learning shortcut-invariant predictive signals, Chroma-VAE isrobust to distribution shifts and performs well when deployed on OOD test cases, especially in highstakes domains where errors are costly, such as the pneumonia prediction application in Section 5.3[50].
Furthermore, the partial reconstructions produced by Chroma-VAE can also indicate if shortcutlearning is occurring in the first place, as they can be examined to identify traits or features that areinvariant.
Potential for Abuse Deep generative models can be used to produce deepfakes [47],
which aresamples generated by the model that are realistic enough to resemble real-life examples, e.g. fakefaces produced by a generative model trained on CelebA .
Chroma-VAE can be abused in a similar(and potentially more dangerous) way, as partial reconstructions can be generated to keep certaintraits while changing others, by concatenating z 1 and z 2 from different samples.
For example,partial reconstructions produced by Chroma-VAE trained on gender labels might be used to generate“counterfactual deepfakes”, e.g. a person with the opposite gender.
Combating deepfakes is an activearea of research, and such techniques can be used to detect Chroma-VAE deepfakes as well.
D Chroma-VAE: Diagram of Model
Figure 9 depicts Chroma-VAE and its training procedure.
In the diagrams below, the forwardfacing solid arrows represent the forward pass and the backwards-facing dotted arrows represent thebackpropagation of gradients.
(a) In the first stage of the training procedure, the encoder, decoder, and classifier are all trained.
The loss function consists of two terms: the cross-entropy loss and the reconstruction loss.
(b) The classifier head is attached to z 1 only ( not the entire latent embedding z ).
As such, the cross-entropy loss term is only backpropagated through z 1 .
(c) The decoder is attached to the entire latent embedding z .
As such, the reconstruction loss term is backpropagated through both z 1 and z 2 .
Consequently, putting both terms of the loss function together, z 1 learns to encode shortcut features whereas z 2 encodes salient non-shortcut features of the input image.