Group-Disentangling Conditional Shift
Authors: Dan Andrei Iliescu
Damon Wischik
Venue: Submitted to ICLR 2023 (2023)
Keywords: group disentanglement
variational autoencoders
conditional shift

Abstract: We propose a novel group disentanglement method called the Context-Aware Variational Autoencoder (CxVAE). Our model can learn disentangled representations on datasets with conditional shift. This phenomenon occurs when the conditional distribution of the instance-level latent variable $\mathbf{z}$ given the input observation $\mathbf{x}$ changes from one group to another (i.e. $p_i(\mathbf{z}|\mathbf{x}) \neq p_j(\mathbf{z}|\mathbf{x})$, where $i,j$ are two different groups). We show that existing methods fail to learn disentangled representations under this scenario because they infer the group $\mathbf{u}$ and instance $\mathbf{z}$ variables separately. CxVAE overcomes this limitation by conditioning the instance inference on the group variable $q(\mathbf{z}|\mathbf{x},\mathbf{u})$. Our model has the novel ability to disentangle ambiguous observations (those with incomplete information about the generative factors), which we evaluate on the task of fair comparisons between student test scores. Additionally, we demonstrate empirically that conditional shift is the cause of our model's improved performance.

Under review as a conference paper at ICLR 2023
GROUP-DISENTANGLING CONDITIONAL SHIFT
Anonymous authors Paper under double-blind review
ABSTRACT
We propose novel group disentanglement method called the Context-Aware Vari- ational Autoencoder (CxVAE). Our model can learn disentangled representations datasets with conditional shift. This phenomenon occurs when the distribu- tion of the instance-level latent variable z conditional on the input observation x, p(z|x), changes from one group another (i.e. Pt (z(x) A Py (z(xx, where i, are two different groups). We show that existing methods fail to learn disentangled representations under this scenario because they infer the group and instance z representations separately. CxVAE overcomes this limitation by conditioning the instance inference on the group variable q(z) x, u). Our model has the novel ability to disentangle ambiguous observations (those with incomplete information about the generative factors), which we evaluate on an image dataset. Additionally, we use fair comparisons task to demonstrate empirically that conditional shift is the cause of our model's improved performance.
INTRODUCTION
Group disentanglement is the goal of learning representations that separate group-level variation from instance-level variation. Consider dataset observations organised into N groups of the form X,,1:K, = (x ..., K., n E 1 N. These could be pictures grouped by author, clinical outcomes grouped by the patient, or film ratings grouped by user. We train representation network r(x,) that encodes a group of observations 1x K., } into one group code u, and set of instance codes (In,I,.... Kn.K., J. one for each observation. We want u to capture only the variation across groups and z only the variation within groups
The current state-of-the-an approaches for group disentanglement train the representation network r by using it as the variational latent posterior distribution in a Variational Autoencoder (Bouchacourt cct al. a 2018 Hosoya 2019; Németh 2020). They assume a hicrarchical generative model whereby the observation x,,k is generated by combining group latent variable u, and an independent instance latent variable & (Figure(Hleff The standard setup involves training the variational latent posterior (U,,, Zn, I:K,, xn) by maximising a lower bound to the data likelihood (Kingma & Welling 2014
Rezende et al. 2014).
In our work, we show that the variational latent posterior, defined in existing models, is unsuited to datasets with conditional shift. This is a property of the data-generating process whereby the true conditional distribution of the instance latent variable z changes from one group to another p; (z|x) # Pj( (z) x) where i, are two groups (Zhang et al. 2013 Gong et al. 2016). In our case, the
conditional instance distribution for group i, which pi Z x), corresponds to p(Zi,k Xt,ke u;) where
k is given instance in the group.
Conditional shift occurs in many real-world datasets that we would like to group grout-ddiee For example, in the 3DIdent dataset (Zimmermann et al. 2021). if we want to infer the colour of the teapot 2n,k based on an image of that teapot X,,k, we should take into account the colour of the spotlight that illuminates the scene u, different coloured spotlights will make the same object appear different colours, as can be seen in Figure [ Existing wroup-disentangleee the instance variable (teapot colour) independently of the group (spotlight colour) fail to disentangle the two colours.
methods, which infer
Existing VAE-basee methods fail disentangle in the conditional shift setting because they make the assumption that the group and instance variables can be inferred independently of each other from the input observation (Figure middlee) When defining the variational latent posterior, existing works

---

Under review as a conference paper at ICLR 2023
Figure I: Conditional shift in the 3DIdent dataset (Zimmermann et al. 2021). The two images on the leftmost column appear show objects of the same colour, even though each image was generated by a different combination of object-colour, spotlight-colour: One teapot is actually orange, while the other is bright green. We can see this by looking at other examples of the same objects under different lighting. The set of images to the right depict different views of the same object on each corresponding row.
based on the Group VAE (GVAE) (Bouchacourt et al. 2018 Hosoya, 2019 Németh 2020 Chen & [Batmanghelich] 2020) make the assumption that the group and instance variables are conditionally independent given the observations (Figure middle):
guu, 2n Xn,1:K.)) qu 1:K,
q Xn,
The limitations of this assumption have not been identified so far in the literature because the datasets used test disentanglement, such Shapes3D (Kim & Mnih 2018), SmalINORB (LeCun al. 2004), dSprites (Higgins et al.) (017). Cars3D (Reed et al. 2015). MPI3D (Gondal et al. 2019). have the property that one image is always sufficient accurately infer its latent variables. For example, we only require single image from the MPI3D dataset to uniquely identify the colour, position, and rotation of the depicted object.
In this work, we show that conditioning the instance encoder on the group latent vector enables the model to learn disentangled representations on datasets with conditional shift.
1. In the first instance, we show that only our method is able to correctly disentangle be- tween object-colour and spotlight-colout in the 3DIdent dataset (Zimmermann et 202I). illustrated in Figure] 1
2. Then, that the amount of conditional shift in the dataset determines the performance gap between model and the other approaches.
the task of fair comparisons between student test-scores (Figure 3) show
RELATED WORK
Group Disentanglement. This class of problems comes under different names: style-content disentanglement (Tenenbaum Freeman 2000), content-transforme 2019), and disentanglement with group supervision (Shu et al.) (020), to name a few. Recent work (Shu al. 2020 Locatello al. 2020 has contextualised group disentanglement as subproblem of reakly-supervised disentanglement, where disentangled representations are learned with the help of non-datapoint supervision (e.g. grouping, ranking, restricted labelling). Early work in this area focused on separating between visual concepts (Kulkarni al. 2015 Reed et al.) 2015). This area has received renewed interest after the theoretical impossibility result of Locatello et al. (2019) and the identifiability proofs of Khemakhem al. (2020) and Mita et al. (2021). A key aspect of recent weakky-suuprvisee models the interpretation of the grouping as signal similarity between datapoints (Chen & Batmanghelichi 2020).
disentanglement (Hosoyal
2

---

Under review as a conference paper at ICLR 2023
kn
Ui
Zmnk
ui
Znnkk
ui
Xn.k
kE1:K
Xm.k
ke1:K
Xa.kk
KET:K
161:N
1:1:
161:N
(a) Group generative model
(b) Inference model existing method GVAE (osooa 2019)
(c) Our inference model (CxVAE)
Figure 2: Our variational distribution conditions the instance representation Zn.k the group representation Un.
Conditioning on the Group Variable. While conditioning the instance encoder on the group variable common in the areas of semi-supervised learning and fair representations (Kingma et al.) 2014; Louizos et al., 2016). we are the first to apply it to unsupervised group disentanglement where explicit group labels are not available. In the field of sequence disentanglement, state- of-the-art methods (Hsu et al., 2017 Denton & Birodkar] 2017, Li & Mandi, (018) infer the instance variable (capturing shorter timescales) conditionally on the group variable (capturing longer timescales). Recent works in weakly-supervee
disentanglement (Shu et al. (2200) [ocatello et al.) 2019. Roeder et al. 2019) also condition the instance variable on the group, but their group variable is discrete variable used for selection rather than representation. It marks which units of the instance representation are common within the group and which are free to vary. We argue that this is not sufficient to account for the variation in p(u[2r, (z u) produced by the conditional shift, we include AdaGVAE (Locatello et al.) 2020) in our evaluation for comparison (see Section 7
Conditional Shift. Our instance encoder conditioned on the group variable is a new strategy to deal with conditional shift. This problem has been studied extensively in the context of supervised learning (Zhang et al., 20033: Gong et al., 2016). However, we are the first to explore the effect of conditional shift on unsupervised learning. Methods for mitigating the effects of conditional shift typically focus on learning domain-invarian representations (Ben-Dae David et al. (009). However, Zhao et al. (2019) show that learning a domain-invarian representation is not sufficient for learning a correct mapping between instance variables from different groups.
Image Translation. Note that our variational latent posterior is different from the one used in COCO-FUNIT (Saito al.) 2020). The authors are motivated by the same limitations with existing works are, namely that unsupervised translation methods struggle to disentangle under condi- tional shift. However, because they train explicitly for translation rather than disentanglement, they arrive at different solution than ours. When performing translation, their approach is to condition the representation of the target group on the source image, thereby bypassing the need for an accurate instance representation. This mechanism produces impressive results image translation tasks, but it cannot be extended to models based on the GVAE which do not train explicitly for translation; in our case, there on the test score dataset and show that our model outperforms it both in terms of disentanglement and translation.
source and target groups in the training set. Regardless, evaluate COCO-FUNIT
3 BACKGROUND
The Group-Inste Instance Generative Model (Bouchacourt et al.) 2018 Hosoya| 2019) is multi-level model that uses two latent variables to generate grouped data: the instance variable Zn,k (O,0 1) controls the variation within groups, and the group variable u, ~ N/ 1) controls the variation across groups (figure(HHeft) The likelihood of a group Xn.I:K, is:
K
p(Xn,1:K.)) = Ep(un.)
Ep(an..)) [p(xn, [n 2n.k))
3

---

Under review as a conference paper at ICLR 2023
3.1 VARIATIONAL INFERENCE
Because the exact likelihood is intractable, the standard approach train the group-instance generative model with a Variational Autoencoder (Kingma & Welling 2014 Rezende et al. (2014) which performs optimisation by introducing variational latent posterior q( 2,,1:K, K, xx and maximizing the Evidence Lower Bound (Jordan et al. 2004):
K
log ((x,,1.K,) 2 E g tn. Den I.
Σ
log p p(x| z)
- -KL(g(u,, (nx In,1: I,, || Uny In,,
Existing methods use class of variational distributions that assume conditional independence between the latent variables.
Kn
quu, K. (n,1::k.) = q(un I:K.)
II
q aZnn Xn,k)
4 CONTEXT-AWARD VARIATIONAL AUTOENCODER
We propose a new model which can perform well on group-confounded problems. We call our model the Context-ANe Aware Variational Autoencoder (Ax. Its defining feature is variational latent posterior whose instance variable is conditioned on the group variable:
K. ((Zn, Xn, Un)
quun Zn,1:Kn, (n,11KK.) ((,, I.K.)
II k=1
Thus, the instance encoder is implemented network which takes input the concatenation of the observation X,, and previously sampled group representation un
(Znnk Xn,k: u,) = Nu,,o), (H, f(xn,k; un);
This form of the variational distribution reflects the correct factorisation of the generative model. It has the potential to learn the true generative latent posterior, which is disentangled by definition. The Evidence Lower Bound for our model is
ELBO Ey(u 1:K..
log p(xxx((xxxk Un, kn,
KK[g(u,, un Ie,, )bpp(,))
(3)
K..
Eey(u, (c..
KL [L Xx,,k un u un)ll(Zzn,))
5 EVALUATION METHODOLOGY
We show that by making the instance encoder conditional on the inferred group variable, we obtain considerable gain in translation accuracy and marked improvement in disentanglement. We also demonstrate that the gap performance between our model and other group disentanglement methods is caused by conditional shift in the data-generating process.
4

---

Under review as a conference paper at ICLR 2023
5.1 MODEL SETUP
We compare our conditional CxVAE with the state-of-the-an in group disentanglement, namely the Group VAE (Hosoyal 2019 Bouchacourt et al., 2018), COCO-FUNIT (Saito al. (020). and AdaGVAE (Locatello et al. 2019). As in Hosoya (2019), the group encoder is applied to each datapoint the group and then all the outputs are averaged.
For all experiments, our CxVAE will be modified GVAE such that the group variable u, is concatenated with the observation Xn.6 and fed into the instance encoder order to compute the instance variable kn For sampling the variational latent posteriors, we use the standard reparametrisation trick. We use an Adam optimiser with learning rate of le-4 with 31 = 0.9, 32 = 0.5.
For the 3DIdent dataset, we implement all networks (encoders and decoders) convolutional nets with hidden layers and filters each. Both latent variables have latent dimensions.
For the test-score dataset, we use MLPs with 3 hidden layer of 32 activations each. The group variable will have dimensions and the instance variable will have 2 dimension.
We train each model for 64 epochs, and use the last 10 epochs for evaluation. Additionally, we run the experiment for 100 different random seeds initialisations, both for the data generating process and the networks. Confidence intervals are computed by resampling train-test splits, weight initialisations and sampling seeds. We use the same 100 seeds each model. This gives 1000 measurements to plot in Table
5.2 EVALUATION METRICS
We compare the models with respect to 3 different criteria: 1) How well does the model fit the holdout data?, 2) How disentangled are the representations inferred by the encoder?, and 3) How well can the model answer the counterfactual question "What colour would appear on the screen if this object were lit by different spotlight?", We assess each criterion with a different evaluation metric.
Fitting the holdout data. As a general metric, we report the reconstruction error (MSE) on the holdout data for every experiment, commonly used as a proxy for the likelihood the holdout set.
Disentangled representations. We use the Mutual Information Gap (Chen et al.) 2018) to measure the quality of the disentanglement. We measure empirically the amount of mutual information between the inferred latent variables u, z and the ground-truth group-level factor u'. Consequently, the goal is to have maximal mutual information between the group variable u and the ground-truth u', and minimal mutual information between the instance variables z and the ground nound-truth u'. The gap between the two (normalized with the entropy of the ground-truth factors) is the metric of disentanglement:
MIG
H(u') (T(,; u(z;
Since the data-gent Gata-generating process is known, the mutual information between the inferred group variable and the ground-truth group variable [(u; u') is straightforward to implement by following the approach from Chen al. (2018).
We measure only the mutual information between the ground-trutt group factor u' and the latent variables because, as pointed out by Németh (2020), the common failure case we are trying to guard against in group disentanglement is that the instance variables z might learn information belonging to the ground-truth group factor u'.
Translation task. measure how well the learned representations can answer the question "What would the score of student & from school have been if they had attended the typical school (the school with scores distributed according to (O, Iz))?". This problem, also known as translation, is commonly used downstream task for disentangled representations (Tenenbaum & Freeman, (2000). We translate the score of student & to the typical school, and then take the mean squared error against the ground-truth translation, which is generated when the data is generated using the ground-truth
5

---

In order to obtain the translation, we first infer the instance variable of student & from school n and the group variable of the typical school. We then feed the two variables to the decoder. For evaluation, translate all the scores from each school and then measure the distance between the predicted translation and the ground-truth translation. We compute the total error as an average over all the translation errors. Under review as a conference paper at ICLR 2023
Table I: Our CxVAE outperforms existing methods at disentangling between object-colour and spotlight-colour in the 3DIdent (Zimmermann et al. (022)) and our test-score datasets. Our model's MIG scores much higher than the one produced by competing models, and the performance gap is greater than the 95% confidence interval for any one of these models. Confidence intervals are computed by resampling train-test splits, weight initialisations and sampling seeds.
MODEL
RECONSTRUCTION
MIG
TRANSLATION
3DIdent (immerrman et al. 2021
GVAE (Hosoya 2019 VdaGVAE (Locatello et al.
0,08 = 0.02
0.09 + 0.06
0.23 0.18
2020
0.09 = 0.05
0.11
0.15 + 0.05
COCO-FUNITY Saioo et CxVAE (ours)
2020
0.05 = 0.01
0.13 0.04
0.13 0.06
0.05 = 0.03
0.72 + 0.06
0.08 + 0.05
Student test scores
GVAE (Hosoya 2019) AdaGVAE (Locatello et
0.41 = 0.01
0.04 + 0.02
0.49 + 0.01
2020
0.41 = 0.01
0.05 0.02
0.49 + 0.01
COCO-FUNIT (Saito et al. 2020
0.40 + 0.02
0.04 0.01
0,49 + 0.02
CxVAE (ours)
0.35 + 0.02
0.44 0.08
0.37 0.02
lower is better
higher is better
lower is better
generative factors. For our dataset, the correct translation corresponds to the Earth-Mover distance between the multivariate normal distributions of scores in each school (Knott Smith 1984).
We can also use translation as an additional qualitative comparison between CxVAE and other group disentanglement methods, which can be seen in Figure
DISENTANGLING AMBIGUOUS OBSERVATIONS
We show that CxVAE able to group-disentangle in setting where conditional shift produces ambiguous observations. The dataset comprises images from the 3DIdent dataset ZZimmermann al. 202I) depicting teapots different colours being lit by spotlights of different colours. Within a group, the images have the same spotlight colour and only the colour of the teapot varies. The goal for the group representation to encode the spotlight colour and for the instance representation to encode the object colour.
The goal of disentangling between object-colour and spotlight- colour is useful for many real-worke applications, such as object recognition. Once we have separated the group-level variation from the instance level variation, we can use the instance representation as a low-level feature on which to train classifier that predicts the colour of the object. The better the disentanglement, the better performance this predictor would have at identifying the true colour of the object.
This is difficult problem because the data exhibits conditional shift. The same exact image could have been generated by different combination of spotlight colour and object colour, as can be seen in Figure [ This makes it difficult to identify the true colour of the object by just looking image. Indeed all previous methods, which infer the instance representation solely from the current image, fail to learn disentangled representations of this dataset.
single

---

Under review as a conference paper at ICLR 2023
Typ School Schos Sahor Schoo
Schol
Typ School
School Bovo Smoo A Typ
Schos
Schoo A-o Typ
Reading Text Score
Reading Test Score
Reading Text Score
#
Hatri
Soore
Matri
(a) Test scores grouped by school.
(b) Failed translation (GVAE).
(c) Successful translation (ours).
Figure 3: Our model correctly disentangles between student aptitude and school characteristics, whereas the GVAE (Hosoya 2019) fails. We ask the counterfactual question "What score would student & from school A have obtained if they had attended the rypical schoo!?". The task is to generate a set of test scores by translating the scores from school A onto the distribution of scores from the typical school. Each line shows the translation for an individual student. Our model preserves the relative positions of the scores, whilst also capturing the distribution of the typical school.
The results in Table show that model, the CxVAE, produces representations that disentangle between spotlight colour and object colour much more than a representative selection of existing models: GVAE, AdaGVAE, and COCO-FUNIT. Our model's MIG score is much higher than the one produced by competing models, and the performance gap is greater than the 95% confidence interval for any one of these models.
CONDITIONAL SHIFT CAUSES THE GAP IN PERFORMANCE
Our CxVAE produces considerable improvements over the competing methods in terms of fitting the holdout set, disentangled representations, translation accuracy and predicting the generative factors (Table Tp. While the scores of the existing methods cluster together, the gap between them and the CXVEET is larger than the 95% confidence interval of any method.
7.1 FAIR COMPARISONS BETWEEN STUDENTS
Consider the task of fair comparisons between students attending different schools based on their standardised test scores in maths and reading (Braun et al., 2006). The typical assumption in the literature is that each school has similar distribution of aptitude among its students, so our goal is to learn disentangled student-level (instance) and school-level (group) representations of the scores. The instance representation 2,, should reflect the aptitude of student k from school n independently of what school the student had attended.
This analysis is crucial for university admission boards aiming to judge students based on their aptitude regardless of the socio-economic cireumstances associated with attending one school or another, such as affluence, location, or curriculum (Raudenbush & Willms 1995 Braun et al. (006). Group disentanglement has the potential to reduce the computational costs this analysis compared to the current variational inference methods for mixed-effects models (Gelman & Hill 2006). The variational latent posterior of a GVAE can perform scalable inference of the student and school representations requiring just forward pass through the network for every additional student. Contrast this with running a separate optimisation routine for the latent variables of each student (Gelman al. 1995). GVAEs can also optimise highly non-linear generative models which would otherwise have to be designed explicitly for the problem in hand (Pinheiro & Bates 2001).
Current methods fail to learn disentangled representations this data. Figure 3b shows the GVAE model (Hosoyal 2019) incorrectly translating the scores from one school to another. Translation is a well-established downstream task for evaluating disentanglement (Tenenbaum & Freeman, 2000).
7

---

Under review as a conference paper at ICLR 2023
In the case of test scores, translation corresponds to the counterfactual question "What score would student from school A have obtained if they had attended the rypical school (i.e. a school whose scores are distributed according to (O, 12))?". We do this by generating new score by combining the instance representation of student & from school n with the group representation the typical school. [Raudenbush & Willms (1995) use translation to directly compare students from different schools. Notice that the unconditional GVAE scrambles the order of the scores and fails to capture the (1, 12) distribution of the typical school.
We propose a modification to the variational latent posterior in order to enable the GVAE to learn disentangled representations of test score data. Looking the data (Figure Baj. it clear that we are dealing with the conditional shift scenario. The same reading score could be obtained by either a high- achieving student from school C' low-achieving student from school B. Inferring the aptitude of the student requires knowing the distribution of scores within each school. In order to account for the variation across groups, we introduce the Context-Awar Variational Autoencoder (Ax) (Figure@rright),
whose instance encoder is conditioned on the group representation (Figure riight). This reflects the correct factorisation of the generative latent posterior pl u 2n,1:K, I,,I:K,,)) thus making no assumptions about the relationship between the group and the instance. We observe variational model successfully translating test scores in Figure 3c
7.2 DATASET
We generate our dataset of test scores using the classic "varying intercept, varying slope" mixed- effects model (Laird & Ware 1982 Pinheiro & Bates] [2001] Gelman & Hill (0066.. This is a well-established approach for modelling student scores K as function of individual aptitude a and school-level characteristics (b,,, C..) (e.g. affluence, curriculum, location, etc.) (Raudenbush & Willms 1995 Braun et al. 2006). We choose this model for its simplicity and for the wide variety of
phenomena to which it can be applied. All the scores and factors are 2-dimensional vectors, with one component for the maths score and another for the reading score:
Xn,k b, Cn. O an,k + 6n is the score of student in school n.
an,k NO,,, I2) is the aptitude of student & in school n
bn ~ NO, I2) is the mean score school n Cn Exp(1) is the standard deviation of scores in school n
(5)
6n,k ~ NO, 0.1 I2) is per-student error term.
For the evaluation procedure, For each school n, we generate M = 128 values for an,k- We then randomly select half of the schools assemble a training dataset with 2, 097, 152 scores split across 384 schools. We take the other half of schools to create the holdout dataset, so that every testing school and student are unseen during training.
the above model to generate N = 32, 768 values for b C.
We have chosen to evaluate on a synthetic dataset because it allows for fine control over the parameters of the data- generatingn process (especially relevant in Section 7.3) and also enables us measure the quality of disentanglement using the Mutual Information Gap (Chen et al. 2018).
7.3 CONDITIONAL SHIFT CAUSES THE GAP IN PERFORMANCE
We show, by modifying the data- gentratinn process ()) that conditional shift explains the increased performance of the VAF. We insert hyper- parameter À to control the strength of the conditional shift; X = means the conditional shift stays the same as in the previous experiment, while À = means there is no conditional shift.
Consider the case where the maths score only depends on the school, and the reading score only depends the student. In this situation, the two generative factors can be easily disentangled since you can infer the student aptitude from the reading score and the school profile from the maths score. We use this as extreme data- generating process that will continuously move between this case and the original case. Our modified data-generating process is
of lack of conditional shift, and insert hyper-pee syper-parameter À in
8

---

Under review as a conference paper at ICLR 2023
Reconstruction
Gap
0.41
0.40
0.35
0.37
0.42
CHAAE
0.35
0.30
0.2g
0g
9,g Rgg
4g
Rdg
Rg
9.7g
Amount of conditionll
Figure 4: Conditional shift in the dataset causes the relative performance gain of our CxVAE over the GVAE. We show performance on datasets generated with different values of the À hyper- parameter, controlling the amount of conditional shift. For low values of À the conditional distribution of student aptitude given a score changed very little from one group to another, and so the Cx and GVAE perform equally well. However, as À increases, the GVAE fails to disentangle.
A
1
Xn.k
b
a &
Cn
+ 6n,k
(6)
where denotes the elementwise power operation and the generative factors (an,k: bu, Caay fn.k) are sampled in (5)
This model has no conditional shift when X = 0 because each ground-truth factor controls separate component of the data. Inferring the student aptitude requires only the reading score and can ignore the school characteristics. When À = 1, the problem exhibits conditional shift in exactly the same way in
our hypothesis is correct, that the conditional shift causes the performance gap between VAE and other group disentanglement methods, then the gap should decrease À approaches The measurements displayed in Figure confirm our expectations. For low values of À the performance of our Cx VAE is evenly matched to the GVAE. As increases, CxVAE metrics remain stable while GVAE performance decreases substantially. It is clear that the degree of confounding in the dataset explains the performance gain that we see CxVAE.
CONCLUSIONS
In this work, we show empirically that conditioning the instance encoder on the group variable produces diuup-dienttngled representations on datasets with conditional shift. We also show that the strength of the conditional shift in the data-generating process determines the performance gap between our model and other group disentanglement methods. Our evaluation is run on the downstream task of extracting student aptitudes from dataset test scorse grouped by school, a problem which group-instance models have not been applied before.
The main limitation our work is that we perform evaluation on a synthetic dataset student scores rather than real data. Although this is a justifiable choice with respect to evaluation (it gives us access to the ground-trus truth values of the latent variables), future work should focus on evaluating on real-world datasets with conditional shift, such user-item ratings (Koren et al. 2009
REFERENCES
Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fermando C Pereira, and Jennifer Wort- man Vaughan. theory learning from different domains. Machine Learning, 79:151-175, 2009.

---

Under review as a conference paper at ICLR 2023
Diane Bouchacourt, Ryota Tomioka, and Sebastian Nowozin. Multi-level variational autoencoder: Learning disentangled representations from grouped observations. AAAI Conference on Artificial Intelligence, 2018.
Henry Braun, Frank Jenkins, and Wendy S. Grigg. Comparing private schools and public schools using hierarchical linear modeling. U.S. Department of Education, National Center for Education Statistics, Institute of Educational Sciences, 2006.
Junxiang Chen and K. Batmanghelich. Weakly supervised disentanglement by pairwise similarities. AAAI Conference on Artificial Intelligence, 2020.
Tian Qi Chen, Xuechen Li, Roger B. Grosse, and David Kristjanson Duvenaud. Isolating sources of disentanglement in variational autoencoders. NeurIPS, 2018.
Emily L. Denton and Vighnesh Birodkar. Unsupervised learning of disentangled representations from video. ArXiv, hts//705..1911 2017.
Andrew Gelman and Jennifer Hill. Data Analysis Using Regression and Multileve//ftierarchica Models. Cambridge University Press, 2006.
Andrew Gelman, John B. Carlin, Hal S. Stern, David B. Dunson, Aki Vehtari, and Donald B. Rubin. Bayesian data analysis. 1995.
Muhammad Walced Gondal, Manuel Wüthrich, Dorde Miladinovic, Francesco Locatello, Martin Breidt, Valentin Volchkov, Joel Bessekon Akpo, Olivier Bachem, Bernhard Schölkopf, and Stefan Bauer. On the transfer of inductive bias from simulation to the real world: a new disentanglement dataset. In NeurlPS, 2019.
Mingming Gong, Kun Zhang, Tongliang Liu, Dacheng Tao, Clark Glymour, and Bernhard Schölkopf. Domain adaptation with conditional transferable components. JMLR workshop and conference proceedings, 48:2839-2848 2016.
Irina Higgins, Loic Matthey, Arka Pal, Christopher P. Burgess, Xavier Glorot, Matthew M. Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. In ICLR. 2017.
Haruo Hosoya. Group-based learning of disentangled representations with generalizability for novel contents. In IJCAI, 2019.
Wei-Ning Hsu, Yu Zhang, and James R. Glass. Unsupervised learning of disentangled and inter- pretable representations from sequential data. In NIPS, 2017.
Michael I. Jordan, Zoubin Ghahramani, T. Jaakkola, and Lawrence K. Saul. An introduction to variational methods for graphical models. Machine Learning, 17:183-233, 2004.
Ilyes Khemakhem, Diederik P. Kingma, and Aapo Hyvärinen. Variational autoencoders and nonlinear ica: A unifying framework. ArXiv, abs/1907 0000004 2020.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. ArXiv, abs/1802.05983, 2018.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. CoRR, abs/1312.6114, 2014.
Diederik P. Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling. Semi-supervised learning with deep generative models. NeurlPS, 2014.
Martin Knott and C. S. Smith. On the optimal mapping of distributions. Journal of Optimization Theory and Applications, 43:39-49, 1984.
Yehuda Koren, Robert M. Bell, and Chris Volinsky. Matrix factorization techniques for recommender systems. Computer, 42, 2009.
Tejas D. Kulkarni, William Whitney, Pushmeet Kohli, and Joshua B. Tenenbaum. Deep convolu- tional inverse graphics network. In NIPS, 2015.
10

---

Under review as a conference paper at ICLR 2023
Nan M. Laird and James H. Ware. Random-effee -ffiettts models for longitudinal data. Biometrics, 38(4): 963-974, 1982.
Yann LeCun, Fu Jie Huang, and Léon Bottou. Learning methods for generic object recognition with invariance to pose and lighting. Proceedings of the 2004 IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2004, CVPR 2004., 2:11-104 Vol.2, 2004.
Yingzhen Li and Stephan Mandt. Disentangled sequential autoencoder. In ICML, 2018.
Francesco Locatello, Stefan Bauer, Mario Lucic, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions the unsupervised learning of disentangled represen- tations. Proceedings of the 36th International Conference Machine Learning. 2019.
Francesco Locatello, Ben Poole, Gunnar Rätsch, Berhhar Scholkopf, Olivier Bachem, and Michael Tschannen. Weakly-superrised disentanglement without compromises. Proceedings of the 37th International Conference on Machine Learning, 2020.
Christos Louizos, Kevin Swersky, Yujia Li, Max Welling, and Richard S. Zemel. The variational fair autoencoder. CoRR, abs/151 1 100830, 2016.
Graziano Mita, Maurizio Filippone, and Pietro Michiardi. An identifiable double vae for disentangled representations. ICML, 2021.
József Németh. Adversarial disentanglement with grouped observations. In AAAI, 2020.
Josae C. Pinheiro and Douglas M. Bates. Mixed-effects models in and s-plus. Technometrics, 43: 113 114, 2001.
Stephen Raudenbush and Douglas Willms. The estimation school effects. Journal of Educational Statistics, 20:307 335, 1995.
Scott E. Reed, Yi Zhang. Y. Zhang, and Honglak Lee. Deep visual maaogyy-mkingg In NIPS, 2015.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In ICML, 2014.
Geoffrey Roeder, Paul K Grant, Andrew Phillips, Neil Dalchau, and Edward Meeds. Efficient amor- tised hayesian inference for hierarchical and nonlinear dynamical systems. ArXiv, abs/1905. 2019.
Kuniaki Saito, Kate Saenko, and Ming- Yu Liu. Coco-funit: Few-shot unsupervised image translation with a content conditioned style encoder. ECCV, 2020.
Rui Shu, Yining Chen, Abhishek Kumar. Stefano Ermon, and Ben Poole. Weakly supervised disentanglement with guarantees. International Conference on Learning Representations, 2020.
Joshua Tenenbaum and William T. Freeman. Separating style and content with bilinear models.
Neural Computation, 12:1247- 2000.
Kun Zhang, Bernhard Schölkopf, Krikamol Muandet, and Zhikun Wang. Domain adaptation under target and conditional shift. ICML, 2013.
H. Zhao, Rémi Tachet des Combes, Kun Zhang, and Geoffrey J. Gordon. On learning invariant representations for domain adaptation. In ICML, 2019.
Roland S. Zimmermann, Yash Sharma, Steffen Schneider, Matthias Bethge, and Wieland Brendel. Contrastive learning inverts the data generating process. In Proceedings of the 38th International Conference on Machine Learning, pp. 12979-12990, 2021.
11


Reviews:
[{"review_id": "bl40Ae-6jn", "reviewer": "Reviewer_G1Wn", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper makes a simple modification to the parameterized posterior distribution to allow for learning group-representations and the within-group instance representation when they are dependent conditional on the observed features: condition the instance representations on both the features and the group representation. As the paper puts it, this can handle conditional shift where changing the group changes the instance representation for the same features.", "summary_of_the_review": "The paper is written well, but it remains to be seen whether the proposed method is useful for any real datasets.", "strength_and_weaknesses": "Strengths:\n\n1. Simple modification to the posterior affords good advantanges.\n2. Promising performance on synthetic data.\n\n\nWeaknesses:\n\nMy main concern with the paper is acknowledged by the authors but nonetheless remains important: \"The main limitation of our work is that we perform evaluation on a synthetic dataset of student scores rather than real data.\"\n\nI do not think such an evaluation can be avoided. Reconstruction error is the one metric that I can trust and that to me only sounds like a part of the story in the paper. Translation additionally seems to be important but I do not see it evaluated on real data.\n\nImportant questions include:\n\n1. What is the point out the translation metric if it cannot be evaluated on real data? The authors say \"Our model preserves the relative positions of the scores\" in figure 2. It this something we desire naturally or something that comes out of an assumption?\n2. How can we guarantee relative positions of the features when translating without restrictions on q?\n3. What the desiderata for disentanglement here without stating the method? How should one evaluate them?\n4. The definition of conditional shift seems to say \"changing the group changes the instance representation for the same features.\". This is a natural consequence of conditioning on the collider as in figure 1 first figure (assuming a causal graph). Why call it conditional shift when it's a consequence of the assumed data generating process?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "arTBwVj0e2L", "reviewer": "Reviewer_iRs6", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper tackles the problem of group disentanglement in presence of conditional shift. This work proposes a new group disentanglement method called the Context-Aware Variational Autoencoder. Experiments on toy datasets show that the proposed method can significantly improve over existing methods.", "summary_of_the_review": "Overall, the writing of the paper is very unclear and the proposed method is only evaluated on toy datasets. It is also unclear how the inherent underspecification of conditional shift is handled in the paper. ", "strength_and_weaknesses": "**Strength** \n- Paper tackles an important and relevant problem\n\n**Weaknesses** \n- Results are present only on the toy dataset in the paper. This is the biggest weakness of the paper. Moreover, since the data-generating process is also proposed in the paper, it is unclear if the dataset is specially designed that can show the failure modes of other methods and if those failure modes are present in other real-world datasets. \n- Since all the experiments are on toy datasets, the claims made in the abstract and introduction are overstated. For example, \"Our model has the novel ability to disentangle ambiguous observations\". There is no concrete evidence in the paper when this will hold and how general of a statement this is? \n- Tackling the problem of conditional shift is very general and ill-posed. It is unclear from the writing how the paper deals with inherent underspecification. \n- Method description in Section 4 is a bit skim. Equations 8-10 appear to be a bit out of the place and it is unclear how the text above these equations follows.  \n- Only the toy dataset is considered in the paper. Any description evaluation criterion is missing. It is hard to understand the tasks considered in Figures 3 and 4. \n- Reproducibility statement is not present and No code is provided as well. Authors can use the 9th page in the main paper and additional appendices to provide those details. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VzmEi_RpaSb", "reviewer": "Reviewer_KtA6", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed a novel group entanglement method under the concern of conditional shift in dataset. In the paper, the author argues that under the conditional shift, the group representation and instance representation cannot be inferred independently since the instance distribution is confounded by the group identity. The proposal is to control the group variable while learning the individual representation. This paper claims to be the first in unsupervised group disentanglement to condition on group variables while learning the individual variables. \n", "summary_of_the_review": "In sum, this paper provides an interesting perspective on the group entanglement under conditional shift. The solution is intuitive and easy to follow. However, the idea is not very novel since it has been explored in many tasks before. It also fails to demonstrate its strength and weakness under more realistic and high dimensional dataset. ", "strength_and_weaknesses": "\nThe major strength of this paper can be summarized in the following aspects:\n\n1. This paper correctly points out the weakness of existing methods in group disentanglement when the conditional shift exists. The idea that the group variables are confounders when inferring the individual variables is important to know. \n\n2. Learning disentangled representation to handle the conditional shift is relatively new since as the paper stated, most methods focus on learning invariant representation under the shift. \n\n3. The synthesis examples are easy to follow and it demonstrates the impact of conditional shift over existing group entanglement methods\n\nThe major weakness of this paper can be summarized in the following aspects:\n\n1. The idea of conditioning/controlling on group variables when inferring the individual variables is not novel. It follows naturally by the definition of conditional shift, i.e. the group conditional distribution of instances are different.  As the paper points that it is widely used in semi-supervised learning. The innovation point is its use in unsupervised learning and generative models. However, it is not sufficient to meet the bar for ICLR. \n\n\n2. One of main concerns for the conditioning method is when it deals with high dimensional problems. In high dimensional space, conditioning would restrict the set of samples that are available to the model in each group. Note that this method essentially learns a set of group specific models. In high dimensional setting, the generative model needs a lot more samples generated before a robust inference result is obtained. This is partially why the conditional independence assumption was used, since it would save a lot of time in data generation. \n\n\n3. That being said, the experiments are too simple. It is a low dimensional example, while the high dimensional data set such as image dataset are mentioned but not tried. It is important to demonstrate the strength and weakness of this method in high dimensional setting, esp the time for inference and the robustness of the inference. \n\n\n4. No code provided. Although this is not hard to implement, it is better to have some code to prove the reproducibility. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "G5RoyxZHM4v", "reviewer": "Reviewer_sLYB", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a context aware variational auto encoder which modified the structure of previous C-VAE. The evaluation is on synthetic data only. ", "summary_of_the_review": "I think the paper is lacking of intuition and details at this stage. In addition, the experiment is insufficient (only synthetic data used). ", "strength_and_weaknesses": "Strength: \n- the work touches a fundamental problem. \n\nWeakness: \n- Only synthetic experiments are conducted. \n- The VAE only tested with MLP.  \n- The data generated is in low dimensional and not very persuasive. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject