D ISENTANGLEMENT AND G ENERALIZATION U NDER C ORRELATION S HIFTS
A BSTRACT
Correlations between factors of variation are prevalent in real-world data.
However, often such correlations are not robust (e.g., they may change between domains, datasets, or applications) and we wish to avoid exploiting them.
Disentanglement methods aim to learn representations which capture different factors of variation in latent subspaces.
A common approach involves minimizing the mutual information between latent subspaces, such that each encodes a single underlying attribute.
However, this fails when attributes are correlated.
We solve this problem by enforcing independence between subspaces conditioned on the available attributes, which allows us to remove only dependencies that are not due to the correlation structure present in the training data.
We achieve this via an adversarial approach to minimize the conditional mutual information (CMI) between subspaces with respect to categorical variables.
We first show theoretically that CMI minimization is a good objective for robust disentanglement on linear problems with Gaussian data.
We then apply our method on real-world datasets based on MNIST and
CelebA, and show that it yields models that are disentangled and robust under correlation shift, including in weakly supervised settings.
I NTRODUCTION
Most research on disentanglement has assumed that the underlying factors of variation in the data are independent (e.g., that factors are not correlated).
However, this assumption is often violated in real-world settings: for example, in domain adaptation, the class distribution often shifts between domains (yielding a correlation between the class and domain); in natural images, there is often a strong correlation between the foreground and background (Beery et al., 2018), or between multiple foreground objects that tend to co-occur (e.g., a keyboard and monitor)
(Tsipras et al., 2020; Beyer et al., 2020).
Importantly, correlated data occurs in areas that affect people’s lives, including in healthcare (Chartsias et al., 2018) and fairness applications (Madras et al., 2018; Creager et al., 2019;
Locatello et al., 2019a), and correlation shifts in these applications are common (e.g., demographics are likely to differ from one hospital to another).
The goal of disentanglement is to encode data into independent subspaces that preferably match the ground truth generative factors.
A common approach (used in ICA, PCA, and VAEs) is to ensure that the latent subspaces share as little information as possible, for example by minimizing the mutual information (MI) between them.
However, recently it has been shown that this fails to disentangle correlated factors (Träuble et al., 2020).
Several works have sought to address this by introducing partial supervision (Träuble et al., 2020; Shu et al., 2019;
Locatello et al., 2020b).
Here, we show that even under full supervision, minimizing the MI can fail: it is impossible to encode generative factors into independent subspaces if they are correlated in the training data.
To address this, we propose minimizing the MI between subspaces conditioned on the correlated attributes.
The goal of this work is to identify and explain the behaviors of different objective functions for correlated and noisy data in a systematic fashion. [Table]
B ACKGROUND & R ELATED W ORK
Disentangled representation learning is often studied in the unsupervised setting, where the groundtruth factors of variation are unknown.
Widely-used approaches for this include variational autoencoders (VAEs) (Kingma & Welling, 2013) and their variants (beta-VAE (Higgins et al., 2017),
TC-beta-VAE (Chen et al., 2018), FactorVAE (Kim & Mnih, 2018), etc.).
However, it was shown by
Locatello et al.
(2019b) that purely unsupervised disentanglement may not be possible.
In this paper, we focus on comparing MI and CMI minimization in the fully-supervised setting.
The goal of domain adaptation and generalisation is closely related as it attempts to learn representations from multiple source domains that transfer to known (e.g., adaptation) or previously unseen (e.g., generalization) target domains.
This is done by either learning domain-invariant representations which discard domain information (Tzeng et al., 2017) or by learning disentangled representations, with latent subspaces that correspond to the domain and the class, respectively (Peng et al., 2019; Ilse et al., 2020; Liu et al., 2018).
For the latter approach, disentanglement is achieved by minimizing the mutual information between latent subspaces (Cheng et al., 2020; Gholami et al., 2020; Nemeth, 2020).
Additional related work is discussed in Appendix A.
D ISENTANGLING C ORRELATED V ARIABLES
In this section, we introduce a toy disentanglement problem where all quantities of interest can be computed analytically.
First, we show that the supervised loss alone does not yield robust disentangled representations, and we discuss why this is problematic.
Then, we show that additionally minimizing the unconditional MI forces the model to learn an even worse solution .
Finally, we show that minimizing the conditional MI yields appropriately disentangled representations that are robust to correlation shift.
This analysis motivates CMI as a good objective for achieving robust disentanglement.
Problem Statement.
Suppose we observe noisy data x ∈ R m obtained from an (unknown) generative process x = g ( s ) where s =
( s 1 , s 2 , . .
, s K ) are the underlying factors of variation , also called source variables or attributes, which may be correlated with each other.
We wish to find a transformation f : R m → R n to a latent space f ( x ) =
z =
( z 1 , z 2 , . . .
, z K ) such that each of the original attributes s k can be recovered from the corresponding subspace z k by a linear mapping
R k , e.g., ˆ s k
= R k z k .
We denote by z − i the set of subspaces { z 1 , . . .
z i − 1 , z i +1 , . . .
, z K } .
We consider three different objectives for learning the latent subspaces: 1) minimizing a supervised loss
L (e.g., mean squared error or cross-entropy), (cid:80) K i =
1 L (ˆ s i , s i ) , denoted “Base” ; 2) minimizing the unconditional mutual information between subspaces in addition to the supervised loss, (cid:80) i L (ˆ s i , s i )+
I ( z 1 , . . .
, z K ) , denoted “Base+MI” ; and 3) minimizing the conditional mutual information between subspaces conditioned on observed attributes , in addition to the supervised loss, (cid:80) i L (ˆ s i , s i ) +
I ( z i ; z − i | s i ) denoted “Base+CMI”
3. F ULL S UPERVISION D OES N OT Y IELD D ISENTANGLEMENT
Consider a linear generative model with correlated Gaussian source variables s , given by:
[Equation]
where A is the mixing matrix and C s and C n are the covariance matrices for the source and noise variables, respectively.
We assume that x is observed and wish to disentangle the underlying source variables s ; this corresponds to finding the mapping A − 1 that inverts the data generating process.
When we have access to the source variables, a natural approach is to minimize a supervised loss to ensure that each subspace contains information about its attribute.
The optimal linear regression solution, both in the least squares sense and with respect to maximum likelihood, is given by the posterior mean: ˆ s ( x ) =
E [ s | x ]
=
C sx C − 1 x x where C sx = E (cid:2) s ( As + n ) ⊤ (cid:3) =
C s A ⊤ and
C x =
AC s
A ⊤ + C n .
The least-squares optimal mapping C sx C − 1 x x is not equal to the inverse
A −of the generative model, as it is biased by the correlation structure C s and C n towards directions of maximal signal-to-noise ratio.
Thus, regression is sensitive to noise and does not disentangle the underlying sources.
In Table 1, we see that in the two-dimensional case where s = ( s 1 , s 2 ) for
A =
I , σ = 0 .
1 and corr ( s 1 , s 2 )
= 0 .
8 , ˆs explains 91 .
9% of the variance in s (column “Base” ).
However, when s 1 and s 2 are uncorrelated, performance drops to 87 .
This drop occurs because the estimator ˆs tries to make use of the assumed correlation between s 1 and s 2 to counteract the information lost due to noise, but this correlation is no longer present in the test data.
3. U NCONDITIONAL D ISENTANGLEMENT F AILS U NDER C ORRELATION S HIFT
In the 2D linear case, we have:
[Equation]
where the matrix W encodes the observation into the latent space.
In standard supervised objectives, there is no constraint preventing a subspace z k from containing information about other source variables than s k .
A common approach to counter the lack of generalization is to minimize the
MI between the latent subspaces
z 1 and z 2 (Chen et al., 2018; Peng et al., 2019).
In the Gaussian case, random variables are independent if
and only if they are uncorrelated .
The optimal linear regression weights that yield I ( z 1 ; z 2 ) =
(e.g., W such that Cov ( z ) is diagonal) can be computed by whitening x and rotating the result by an angle ϕ opt which leads to maximal VE ( ϕ opt = − π / 4 for positive correlations and A = I ).
See Figure 6 in Appendix C for details.
However, the resulting model no longer performs well on in-distribution-data (Table 1, column “Base+MI” ).
There is correlation between the sources s 1 and s 2 and therefore I ( s 1 ; s 2 ) > 0 .
By enforcing independence, at least one of the subspaces cannot contain all relevant information about its target value and therefore will have poor predictive performance (see Proposition 1 in Appendix E).
3. C ONDITIONAL D ISENTANGLEMENT IS R OBUST TO C ORRELATION S HIFT
We have seen that enforcing unconditional independence between the latent spaces does not solve the disentanglement problem.
However, from the graphical model (Appendix C, Fig. 5), it is clear that z 1 and z 2 are independent conditioned on either of s 1 or s 2 : assuming a common cause for the correlation between s 1 and s 2 , there is a connection in the graphical model between zand z 2 introducing the statistical dependence.
Observing either s 1 or s 2 disconnects z 1 and z 2 .
Enforcing independence conditioned on each of the source variables is also sufficient to yield a robust disentangled representation.
For our 2D example, this corresponds to I ( z 1 ; z 2 | s 1 )
=
and I ( z 1 ; z 2 | s 2 )
= 0 .
If s 1 and s 2 are correlated and if we can predict s 1 from z 1 , then it must be the case that z 1 contains information about s 2 .
To improve robustness to shifts, we wish to ensure that z 1 and z 2 share as little information as possible, namely exactly the information which is necessary to account for the correlation between the sources.
This can be enforced via I ( z 1 , z 2 | s 2 ) , which states that if we know s 2 , then knowing z 1 does not give us more information about z 2 .
The same logic applies for the need of I ( z 1 , z 2 | s 1 ) .
The optimal solution under the constraint [Figure] of conditional independence is achieved by the mapping W =
A − 1 , successfully recovering the underlying generative model.
We find that for this solution, the performance of the model is not affected by the correlation shift (Table 1, column “Base+CMI” ).
This demonstrates the usefulness of minimizing CMI and motivates us to investigate CMI minimization for larger-scale tasks.
M INIMIZING CMI
In simple cases such as linear regression, we can compute and minimize the MI/CMI analytically; however, for most tasks (including classification), there is no closed form for the mutual information.
In this section, we describe our approach to minimize the CMI in general classification settings.
We formally describe the algorithms for the baselines ( Base and Base + MI ) in Appendix B.
Suppose we have a dataset D
= { ( x ( i ) , a ( i ) ) } Ni =1 where x ( i ) is an example and a ( i ) is a vector of attribute labels ( a ( i ) k is the label for the k th attribute of the i th example).
We consider discrete attribute values, a ( i ) k ∈ N .
Let f θ : x (cid:55)→ z denote an encoder function parameterized by θ that maps examples x ∈ R m to latent representations z ∈ R n .
We aim to learn one latent subspace per attribute, such that each subspace is independent from all other subspaces conditioned on the encoded attribute.
To obtain samples from p ( z 1 , .
, z K | a k ) and p ( z k | a k ) p ( z − k | a k ) , we loop over values of a k , and for each condition { a k = 0 , a k = 1 , . . .
} , we select examples from the minibatch that satisfy the condition, giving us samples from p ( z 1 , . . .
, z K | a k ) ; then we shuffle the latent subspaces z j , ∀ j ̸ = k jointly batchwise (e.g., combining z k from one example with z − k from another) to obtain samples from p ( z k | a k ) p ( z − k | a k ) .
To enforce p ( z 1 , . .
z K | a k ) =
p ( z k | a k ) p ( z − k | a k ) , we train the encoder f adversarially against a discriminator trained to distinguish between these two distributions.
The discriminator takes as input a representation and predicts whether it is “real” (e.g., drawn from the joint distribution) or “fake” (e.g., drawn from the product of marginals).
One discriminator is trained for each attribute a k , which receives samples from the two distributions and the attribute value it is conditioned on.
In practice, we use a conditional discriminator, effectively sharing parameters between the discriminators for each of the attributes.
This approach is architectureagnostic.
The process is illustrated in Figure 1 and the corresponding algorithms are in Appendix B.
E XPERIMENTS
First, we present results for synthetic regression and classification tasks.
Next, investigate a correlated version of CelebA under fully and weakly supervised settings, and show that minimizing CMI can largely eliminate the gap in performance under correlation shift.
Experimental details and additional experiments are in Appendix C and D. [Figure]
Toy Linear Regression.
In the linear regression problem in Section 3, we evaluated the performance under correlation shift for one specific noise level and correlation strength.
Here, we varied these parameters and found that our findings hold over all noise levels and non-zero correlation strengths.
As shown in Figure 2a, the performance of Base drops most severely under correlation shift for strong correlations and intermediate noise levels.
In this regime, the advantage of Base+CMI is most clear.
Toy Multi-Attribute Classification.
We show in a controlled setting that our findings hold for classification tasks with multiple attributes.
Here, the binary source attributes a k = ± 1 , ∀ k ∈ { 1 , . . .
, K } generate the observed data via x = Aa + n (we set A
= I for simplicity) with normally distributed noise n ∼ N ( 0 , C n ) .
We induce pairwise correlations between the attributes a k .
As for the regression task, we find that Base+CMI leads to robustness under correlation shift (Figure 2b).
CelebA. Finally, we consider a realistic setting using the CelebA faces dataset (Liu et al., 2015).
We considered two attributes that we knew a priori were not causally related, Male and Smiling , and we created subsampled datasets that satisfied specific correlations between attributes.
We investigated a range of training correlations { 0 , 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 } , and evaluated our models on both anti-correlated and uncorrelated test sets (Figures 3a and 3b).
We found that minimizing CMI has a larger effect for medium-to-high correlation; however, CMI minimization does not hurt performance at low correlation strengths (Figure 3c).
Note that while the unconditional model appears to have good performance on the anti-correlated test set, its performance is poor on the validation set (that has the same correlation structure as the training set), so this model does not perform well on in-distribution data.
In contrast, the conditional
MI model performs well on both in-distribution data and shifted test distributions.
Also note that the problem of disentangling correlated attributes does not occur only under correlation shift,
[Figure]
domain where certain attribute combinations will reliably be treated incorrectly.
For example, Base fails to recognize the rare non-smiling male faces in 49% of the cases, while Base+CMI fails only in 25% of the cases.
Extension to Weakly Supervised Settings.
Our method can be applied directly to weakly supervised settings.
Importantly, it is not necessary to have labels for multiple attributes for a single data point.
We find that when reducing the number of labels, Base+CMI outperforms the other objectives under correlation shift (Figure 4).
C ONCLUSION
Correlations are prevalent in real-world data, yet pose a substantial challenge for disentangled representation learning.
Standard approaches learn to rely on these correlations, and when the attributes are not causally related, this leads to poor performance under test-time correlation shift.
We establish CMI minimization as a more appropriate alternative to mutual information minimization, which sets the stage for the development of more powerful objective functions for disentanglement. A CKNOWLEDGEMENTS
We thank Jörn-Henrik Jacobsen for his valuable contributions in the early stage of this work.
We thank Steffen Schneider, Dylan Paiton, Lukas Schott, Elliot Creager and Frederik Träuble for helpful discussions.
We thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Christina Funke.
Paul Vicol was supported by a JP Morgan AI Fellowship.
We acknowledge support from the German Federal Ministry of Education and Research (BMBF) through the Competence Center for Machine Learning (FKZ 01IS18039A) and the Bernstein Computational Neuroscience Program Tübingen (FKZ: 01GQ1002), the German Excellence Initiative through the Centre for Integrative Neuroscience Tübingen (EXC307), and the Deutsche Forschungsgemeinschaft (DFG; Projektnummer 276693517 – SFB 1233). Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute www.vectorinstitute.ai/partners. A E XTENDED R ELATED W ORK
Estimating & Optimizing Mutual Information.
Many approaches have been proposed for MI and CMI estimation and optimization.
The Mutual Information Neural Estimator (MINE) (Belghazi et al., 2018) uses a lower-bound of the MI based on the Donsker-Varadhan dual representation of the
KL divergence (Donsker & Varadhan, 1983).
Poole et al.
(2019) provide an overview of variational bounds that can be used to estimate MI; most are lower bounds , which are useful in principle for maximizing MI, but which have also been used to minimize MI (even though minimizing a lower bound is not guaranteed to decrease MI).
CLUB (Cheng et al., 2020) introduced a variational upper bound of MI, providing a more principled objective for minimizing MI.
Several CMI estimators have been proposed, including conditional-MINE (Molavipour et al., 2020a), C-MI-GAN (Mondal et al., 2020), CCMI (Mukherjee et al., 2020), and an approach based on nearest neighbors (Molavipour et al., 2020b).
Many approaches to MI minimization are based on batchwise shuffling of latent subspaces, sometimes referred to as metameric sampling (Belghazi et al., 2018; Nemeth, 2020;
Feng et al., 2018; Park et al., 2020; Peng et al., 2019).
The approach we use in Section 4 follows this paradigm of latent-space shuffling.
Correlations Between Features.
With roots in ICA, most research on disentanglement focuses on data that was generated by independent factors, including synthetic benchmarks such as dSprites (Matthey et al., 2017), Shapes3D (Burgess & Kim, 2018), Cars3D (Reed et al., 2015),
SmallNORB (LeCun et al., 2004), or MPI3D (Gondal et al., 2019).
In real-world datasets on the other hand, factors are often correlated (Welinder et al., 2010; Lin et al., 2014).
Träuble et al.
(2020) pointed out the challenges that arise when attempting to learn disentangled representations on correlated data, and performed a large-scale empirical evaluation of the effect of correlations on widely-used VAE-based disentanglement models.
Our work shows that even under full supervision, correlated attributes are problematic when enforcing independence between latent subspaces.
Causally-informed modeling (Zhang et al., 2020) is another approach to learning disentangled representations and extracting invariant features; an example of this is Invariant Risk Minimization (IRM) (Arjovsky et al., 2019).
To investigate the effect of correlations systematically, it is common to modify existing datasets to induce correlations, for example by subsampling the data, or by generating synthetic datasets with the desired properties (Dittadi et al., 2020; Cimpoi et al., 2014; Jacobsen et al., 2018; Locatello et al., 2019b).
We follow this approach in our experiments.
Fairness.
An important application of disentanglement is fairness.
Typically, this can be addressed by modifying the training data to be unbiased or by adding a regularizer (e.g. based on mutual information) that quantifies and minimizes the degree of bias (Kamiran & Calders, 2009; Kamishima et al., 2011; Zemel et al., 2013; Hardt et al., 2016; Cho et al., 2020).
Mutual Information.
The mutual information (MI) between two random variables x and y , denoted
I ( x ; y ) , is the KL divergence between the joint distribution p ( x , y ) and the product of the marginal distributions p ( x ) p ( y ) :
I ( x ; y )
=
D KL [ p ( x , y ) ||
p ( x ) p ( y )] .
MI minimization is at the heart of many approaches to disentanglement.
The conditional mutual information (CMI) is defined as:
[Equation]
CMI measures the dependency between two variables given that we know the value of a third variable.
For example, there is a dependency between a country’s number of Nobel laureates per capita and chocolate consumption per capita (Prinz, 2020).
However, this dependency is largely explained by the wealth of a country, thus I ( nobel ; chocolate | wealth ) <
I ( nobel ; chocolate ) .
In general, the
CMI can be smaller or larger than the unconditional MI.
B A LGORITHMS
In this section, we provide formal descriptions of the algorithms we consider in this paper.
Algorithm 1 describes the classification-only baseline, that trains separate linear classifiers to predict attributes a k from the corresponding latent subspaces z k (corresponds to “Base” ).
Algorithmand Algorithm 3 describe the encoder and discriminator training loops for the unconditional disentanglement baseline (corresponds to “Base+MI” ).
This objective adversarially minimizes the discrepancy between samples from the joint distribution p ( z 1 , . .
z k ) and the product of marginals p ( z 1 ) · · · p ( z k ) .
Algorithm 4 and Algorithm 5 describe the encoder and discriminator training loops for conditional disentanglement (corresponds to “Base+CMI” ).
1: Input: { ϕ 1 , . . .
, ϕ K } , initial parameters for K linear classifiers C 1 , . . .
, C K 2: Input: θ , initial parameters for the encoder F
: Input: α, β learning rates for training the encoder and linear classifiers 4: while true do 5: ( x , { a k } Kk =1 ) ∼ D Train ▷ Sample a minibatch of data with attribute labels 6: z ← F
[Table]
( x ) ▷ Forward pass through the encoder 7: { z
[Table]
Kk
[Table]
← SplitSubspaces ( z , k ) ▷ Partition the latent space into k subspaces 8:
L ← (cid:80)
[Table] [Table] [Table]
, a
[Table]
▷ Cross-entropy for each attribute 9: θ ←
[Table]
L ▷ Update encoder parameters 10: ϕ k ← ϕ k − β ∇ ϕ k L , ∀ k ∈ { 1 , . .
, K } ▷ Update classifier parameters 11: end while
1: Input: { ϕ 1 , . . .
, ϕ K } , initial parameters for K linear classifiers C 1 , . . .
, C K 2: Input: θ , initial parameters for the encoder F
: Input: α, β learning rates for training the encoder and linear classifiers 4: while true do 5: ( x , { a k
[Table]
Kk = [Table] ) ∼ D Train ▷ Sample a minibatch of data with attribute labels 6: z ←
[Table]
▷ Forward pass through the encoder 7: {
[Table]
SplitSubspaces (
[Table]
, k ) ▷ Partition the latent space into k subspaces 8:
L ← (cid:80) Kk
=
L cls ( C k ( z k ; ϕ k ) , a k ) ▷ Cross-entropy for each attribute 9: z ′ ∼ p ( z 1 ) p ( z 2 ) · · · p ( z k )
▷ Samples w/ batchwise-shuffled subspaces 10:
L ← L + log (1 − D ω ( z ′ ))
+ log ( D ω ( z ))
▷ Add adversarial loss 11: θ ← θ − α ∇ θ L ▷ Update encoder parameters 12: ϕ k ← ϕ k − β ∇ ϕ k L , ∀ k ∈ { 1 , . .
, K } ▷ Update classifier parameters 13: end while
1: Input: ω , initial parameters for the discriminator D 2: Input: γ , learning rate for training the discriminator 3: while true do 4: ( x , { a k } Kk =1 ) ∼ D Train ▷ Sample a minibatch of data with attribute labels 5: z ← F θ ( x ) ▷ Forward pass through the encoder 6: { z k } Kk =1 ← SplitSubspaces ( z , k ) ▷ Partition the latent space into k subspaces 7: z ′ ∼ p ( z 1 ) p ( z 2 ) · · · p ( z k )
▷ Samples w/ batchwise-shuffled subspaces 8:
L ← L + log ( D ω ( z ′ ))
+ log (1 − D ω ( z ))
▷ Add adversarial loss 9: ω ← ω − γ ∇ ω L ▷ Update discriminator parameters 10: end while
1: Input: { ϕ 1 , . . .
, ϕ K } , initial parameters for K linear classifiers R 1 , . . .
, R K 2: Input: θ , initial parameters for the encoder f 3: Input: α, β learning rates classifiers 4: while true do 5: ( x , { a k } Kk =1 ) ∼ D Train ▷ Sample a minibatch of data with attribute labels 6: z ← f θ ( x ) ▷ Forward pass through the encoder 7: { z k } Kk =1 ← SplitSubspaces ( z , K ) ▷ Partition the latent space into K subspaces 8:
L ← (cid:80) Kk
=
L cls ( R k
( z k ; ϕ k ) , a k ) ▷ Cross-entropy for each attribute 9: for k ∈ { 1 , . . .
, K } do ▷ For each attribute/subspace 10: z ′ ∼ p ( z 1 , . . .
z K | a k )
▷ Samples from the joint distribution 11: z ′′ ∼ p ( z k | a k ) p ( z − k | a k )
▷ Samples w/ batchwise-shuffled subspaces 12:
L ← L + log (1 − D ω ( z ′′ ))
+ log ( D ω ( z ′ ))
▷ Add adversarial loss 13: end for 14: θ ← θ − α ∇ θ L ▷ Update encoder parameters 15: ϕ k ← ϕ k − β ∇ ϕ k L , ∀ k ∈ { 1 , . .
, K } ▷ Update classifier parameters 16: end while
▷ Sample a minibatch of data with attribute labels 5: z ← F θ ( x ) ▷ Forward pass through the encoder 6: { z k } Kk =1 ← SplitSubspaces ( z , k ) ▷ Partition the latent space into K subspaces 7:
L ←▷ L will accumulate the losses over all subspaces 8: for k ∈ { 1 , . .
, K } do 9: z ′ ∼ p ( z 1 , . . .
z K | a k )
▷ Samples from the joint distribution 10: z ′′ ∼ p ( z k | a k ) p ( z − k | a k )
▷ Samples w/ batchwise-shuffled subspaces 11:
L ← L + log ( D ω ( z ′′ ))
+ log (1 − D ω ( z ′ ))
▷ Add adversarial loss 12: end for 13: ω ← ω − γ ∇ ω L ▷ Update discriminator parameters 14: end while
C E XPERIMENTAL D ETAILS
Graphical Model.
In Figure 5 we provide the graphical model discussed in Section 3.3.
[Figure]
Compute Environment.
Our experiments were implemented using PyTorch (Paszke et al., 2019), and were run on our internal clusters.
The toy 2D experiments were run on a single NVIDIA RTX 2080 TI GPU, and took approximately 48 hours for all the results presented.
The MNIST and
CelebA experiments were run on NVIDIA Titan Xp GPUs.
Each run of the multi-digit MNIST and
CelebA tasks for a given method and correlation strength (and noise level in the MNIST case) took approximately 12 hours, and these were run in parallel.
Toy Linear Regression.
Here, we describe the optimal solution for the Base+MI objective for the linear regression task with A
= I .
Optimal linear regression with zero mutual information between z 1 and z 2 can be obtained by taking the singular value decomposition of x followed by whitening and rotation by − 45 ◦ (see Figure 6).
xxsszzsssvd
VE=0.
zzss...
+ whitening
VE=0.
[Figure]
[Figure]
by
[Figure]
opt [Figure] our case this is easy: the principal components of x are x 1 + x 2 and x 1 − x 2 .
The optimal regression loss with minimal MI is then given by whitening and rotating the result by angle ϕ opt which leads to maximal VE ( ϕ opt = − π / 4 for positive correlations and A = I ).
[Figure]
Toy Multi-Attribute Classification.
The following figure illustrates the data x for different correlation strength and noise levels (Figure 7) in the case of two attributes.
In the 2D case increasing the correlation strength means that data points with a 1 = a 2 are increasingly more common relative to a 1 ̸
= a 2 .
In the multi-attribute setting, the correlation strength refers to the pairwise correlation between all attributes.
We used a PacGAN-style setup (Lin et al., 2018) for our toy experiments, where the discriminator takes as input a concatenation of 50 samples.
• Base : We used Adam (Kingma & Ba, 2014) with a learning rate of 0.01.
• Base + MI : We used Adam to optimize the encoder, linear classifiers, and discriminators.
After each step of optimizing the discriminator and encoder, we optimized the linear classifiers ( R ) for 10 steps.
The disentanglement loss term was weighted by a factor ofrelative to the classification loss.
In preliminary tests, we found that the optimal learning rate depended on noise level, correlation strength, and number of attributes.
The results in Figure 2b were obtained using one of the following learning rates for the discriminator
{ 1 e − 4 , 2 e − 4 , 5 e − 4 , 1 e − 3 , 5 e − 3 } .
The learning rate of the generator and linear classifiers was chosen to be 10 times smaller than the discriminator learning rate.
• Base + CMI :
For A =
I , no optimization was necessary, as we already know the optimal solution to be W
=
A −
= I .
We confirmed experimentally that the discriminator could not get above chance performance for this solution.
CelebA. For all experiments, we used minibatch size 100, and latent dimension D = 10 .
As the encoder model, we used a three-layer MLP with 50 hidden units per layer and ReLU activations.
Similarly to the MNIST setup, we trained for 200 epochs, using Adam to optimize the encoder, linear classifiers, and discriminators.
For each method, we performed a grid search over learning rates { 1 e − 5 , 1 e − 4 , 1 e − 3 } separately for each of the encoder, discriminator(s), and linear classification heads; we selected the best learning rates based on validation accuracy.
Weakly Supervised Setting.
For the fully supervised CelebA experiment, labels for both attributes were available for all 10260 images.
For the weakly supervised setting, we reduced the number of labels to 5130 (50% of the labels of the fully supervised dataset), 2565 (25%), 1026 (10%), or(5%) for each attribute.
This implies that some images had both labels, some had only one label and some images had no labels (for example when using 50% of the labels the distinction is as follows: 25% of the images had both labels; 25% had only labels for attribute 1; 25% had only labels for attribute 2; and 25% had no labels).
The three objectives can be applied to these weakly supervised settings.
For Base , the cross-entropy loss for each attribute was computed only for the images that had labels for the corresponding attribute.
For Base+MI no labels are required for the unconditional shuffling; thus this objective can be applied even for the images without labels.
For Base+CMI , our method shuffles only images that have the same value for a given attribute.
This also works if the labels of the other attribute are missing.
We used the same training parameters as for the supervised experiment, except for increasing the number of training epochs (up to 1200 epochs) and adapting the minibatch size to the number of labels.
D E XTRA E XPERIMENTS
Extended Analysis for Toy Linear Regression.
Here, we provide intuitive explanations of the behaviors of the different objectives.
Comparing the correlation of target s and data x with the correlation of the predictions ˆ s of the different models can help us understand the findings of Section 3.
As shown in Figure 8, the predictions of optimal linear regression ( Base+MI ) (corr (ˆ s 1 , ˆ s 2 ) = 0 . 85 ) are stronger correlated than the data (corr ( x 1 , x 2 ) = 0 .
This shows that the correlation present in the training data is used to compensate for the noise.
Enforcing unconditional independence ( Base+MI ) on the other hand, leads to uncorrelated predictions (corr (ˆ s 1 , ˆ s 2 ) = 0 ).
Undoubtedly, this cannot be the correct solution, as the targets are correlated.
Additionally, when the correlations change, the independence constraint does not hold anymore.
This can lead to interesting effects under correlation shift.
While for most noise levels, the performance on the test data is poor, for some noise levels the performance can be even higher than training performance (Figure 2a).
In these cases the model can "accidentally" exploit the correlation in the test data to make the correct predictions.
The correlation of the predictions (corr (ˆ s 1 , ˆ s 2 ) = 0 . 73 ) matches the correlation of the data only for
Base+CMI .
Multi-Digit Occluded MNIST.
As an additional experiment, we designed a larger-scale task to systematically investigate whether the properties found in the toy classification task hold in a more complex setting.
We created a dataset by concatenating two MNIST digits side-by-side, where the aim is to predict both the left- and right-hand-side labels.
We used a subset of MNIST consisting of classes 3 and 8 (which are visually similar and can become ambiguous under occlusions).
We generated ss
Target
Corr = 0.
xx
Data
Corr = 0.
ˆsˆ s
Predictions of Base
Corr = 0.
ˆsˆ s
Predictions of Base + MI
Corr
= -0.
ˆsˆ s
Predictions of Base + CMI
Corr = 0.
Training (Corr = 0.8) ss
Corr
= -0.
xx
Corr
= -0.
ˆsˆ s
Corr = 0.
ˆsˆ s
Corr
= -0.
ˆsˆ s
Corr
= -0.
Test (Corr = 0) occlusion masks using the procedure used by Chai et al.
(2021); examples from our synthetic dataset under a range of noise settings are shown in Figure 9a.
This mimics the multiple-object classification setting in a way that allows us to control the correlation strength and noise level (via the amount of occlusion), allowing for systematic analysis.
While this would also be possible for colored MNIST and dSprites, one advantage of our choice is the symmetry of our task, which allows us to exclude potential side-effects: here, the attributes have the same type (the digit identity), whereas the attributes in colored MNIST (digit identity and color) and dSprites (shape, size, position, etc.) are more diverse.
Similarly to the toy tasks, we train an encoder to map images onto a D -dimensional latent space, which is partitioned in two equal-sized subspaces corresponding to the two digits; we train a linear classifier on each subspace to predict the respective class labels.
We consider different correlation strengths between the left and right digits in the training set (where strong correlation means that the digits match most of the time, e.g., 3-3 or 8-8 are more common than 3-8 or 8-3 ).
We evaluate each model on test data with correlation strengths ranging from [ − 1 , 1] .
The results are shown in
Figure 9b.
We found that the conclusions from the toy experiments hold in this setting: supervised learning with only the classification loss, as well as with unconditional MI minimization, fails under test-time correlation shift, while the model minimizing conditional MI is more robust.
[Figure]
Disentanglement Metrics.
Previous work (Locatello et al., 2020a) has shown that common disentanglement metrics are not suitable for the correlated setting.
For this reason, we focused on comparing performance under correlation shift, which we consider more suitable for correlated data: if a model cannot predict a factor of variation well for certain values of another factor, then the model did not successfully disentangle these factors of variation.
However, one can still make use of the disentanglement metrics by evaluating them on uncorrelated data using the models trained on correlated data.
We performed this analysis for two of our datasets and found in both cases that Base+CMI reached better scores compared to the other objectives for almost all metrics.
Disentanglement results for the toy classification task with ten attributes are shown in Table 2.
The disentanglement results for
CelebA are shown in Table 3.
Since the disentanglement metrics require that the factors of variation are each encoded in one-dimensional subspaces, we set latent dimension D = 2 for this experiment.
[Table]
[Table]
E P ROOF OF P ROPOSITION RELATED TO S ECTION 3.
In Section [Table] 2 we argue that by enforcing independence, at least one of the subspaces cannot contain all relevant information about its target value and therefore will have poor predictive performance.
We make this precise in the following proposition:
Proposition 1 If I ( s 1 ; s 2 ) > 0 , then enforcing I ( z 1 ; z 2 ) = 0 means that I ( z k ; s k ) <
H ( s k ) for at least one k .
Proof.
Assume that I ( s 1 ; s 2 ) > 0 and at the same time
I ( z k ; s k )
=
H ( s k ) (i.e., we are proving by contradiction).
Since I ( z 1 ; s 1 )
=
H ( s 1 ) , we have H ( s 1 | z 1 )
= 0 and with H ( s 1 |
z 1 )
=
H ( s 1 | z 1 , s 2 )
+ I ( s 1 ; s 2 | z 1 ) (both non-negative), it follows that H ( s 1 | z 1 , s 2 )
=
I ( s 1 ; s 2 | z 1 ) = 0 .
Since for the interaction information, by definition I ( s 1 ; s 2 ; z 1 ) =
I ( s 1 ; s 2 )
− I ( s 1 ; s 2 | z 1 ) , and
I ( s 1 ; s 2 | z 1 ) = 0 , we have I ( s 1 ; s 2 ; z 1 ) =
I ( s 1 ; s 2 ) > 0 .
Since we also assume H ( s 2 | z 2 ) = 0 , we also have I ( s 1 ; s 2 ; z 2 )
=
I ( s 1 ; s 2 ) > 0 .
We can use this to compute the fourth order interaction information I ( s 1 ; s 2 ; z 1 ; z 2 ) .
By definition, we have I ( s 1 ; s 2 ; z 1 ; z 2 ) =
I ( s 1 ; s 2 ; z 1 ) − I ( s 1 ; s 2 ; z 1 | z 2 ) .
We just showed that I ( s 1 ; s 2 ; z 1 ) =
I ( s 1 ; s 2 ) , and therefore we have I ( s 1 ; s 2 ; z 1 | z 2 )
=
I ( s 1 ; s 2 | z 2 ) .
Together it follows that:
[Equation]
On the other hand, we know that
= H ( s 1 | z 1 ) =
H ( s 1 | z 1 ; z 2 )
+ I ( s 1 , z 2 | z 1 )
and therefore I ( s 1 , z 2 | z 1 ) = 0 .
Therefore, the interaction information I ( s 1 ; z 2 ; z 1 )
=
I ( s 1 ; z 2 ) − I ( s 1 ; z 2 | z 1 )
=
I ( s 1 ; z 2 ) ≥ 0 .
At the same time, we assumed that I ( z 1 ; z 2 ) =
and hence I ( z 1 ; z 2 ; s 1 )
+ I ( z 1 ; z 2 | s 1 ) = 0 , which shows that I ( z 1 ; z 2 ; s 1 ) ≤ 0 .
Together, we see that I ( z 1 ; z 2 ; s 1 ) =
I ( s 1 ; z 2 ) = 0 .
Now we can decompose I ( s 1 ; s 2 ; z 1 ; z 2 ) in a different way: I ( s 1 ; s 2 ; z 1 ; z 2 )
=
I ( s 1 ; z 1 ; z 2 ) −
I ( s 1 ; z 1 ; z 2 | s 2 ) .
We know that I ( s 1 ; z 1 ; z 2 ) =
I ( s 1 ; z 2 ) and therefore I ( s 1 ; z 1 ; z 2 | s 2 ) =
I ( s 1 ; z 2 | s 2 ) > 0 and that I ( s 1 ; z 1 ; z 2 ) = 0 .
Therefore, it follows that:
[Equation]
which is a contradiction with I ( s 1 ; s 2 ; z 1 ; z 2 )
=
I ( s 1 ; s 2 ) > 0 .
Therefore, if I ( s 1 ; s 2 ) > 0 and
I ( z 1 ; z 2 ) = 0 , it must hold that I ( z k ; s k ) <
H ( s k ) for at least one k , which we wanted to show.