Original: G ROUP -D ISENTANGLING C ONDITIONAL S HIFT
Revised: G ROUP -D ISENTANGLING C ONDITIONAL S HIFT

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We propose a novel group disentanglement method called the Context-Aware Variational Autoencoder (CxVAE).
Revised: We propose a novel group disentanglement method called the Context-Aware Variational Autoencoder (CxVAE).

Original: Our model can learn disentangled representations on datasets with conditional shift.
Revised: Our model can learn disentangled representations on datasets with conditional shift.

Original: This phenomenon occurs when theconditional distribution of the instance-level latent variable z given the input observation x  changes from one group to another
Revised: This phenomenon occurs when the distribution of the instance-level latent variable z conditional on the input observation x , p ( z | x ) , changes from one group to another

Original: (i.e. p i ( z | x ) (cid:54) =
Revised: (i.e. p i ( z | x ) (cid:54) =

Original: p j ( z | x ) , where i, j are two different groups).
Revised: p j ( z | x ) , where i, j are two different groups).

Original: We show that existing methods fail to learn disentangled representations under this scenario because they infer the group u and instance z variables separately.
Revised: We show that existing methods fail to learn disentangled representations under this scenario because they infer the group u and instance z representations separately.

Original: CxVAE overcomes this limitation by conditioning the instance inference on the group variable q ( z | x , u ) .
Revised: CxVAE overcomes this limitation by conditioning the instance inference on the group variable q ( z | x , u ) .

Original: 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.
Revised: 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.

Original: Additionally, we  demonstrate empirically that conditional shift is the cause of our model’s improved performance.
Revised: Additionally, we use a fair comparisons task to demonstrate empirically that conditional shift is the cause of our model’s improved performance.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Group disentanglement is the goal of learning representations that separate group-level variation from instance-level variation.
Revised: Group disentanglement is the goal of learning representations that separate group-level variation from instance-level variation.

Original: Consider a dataset of observations organised into N groups of the form x n, 1: K n = { x n, 1 , ..., x n,K n } , n ∈ 1 : N .
Revised: Consider a dataset of observations organised into N groups of the form x n, 1: K n = { x n, 1 , ..., x n,K n } , n ∈ 1 : N .

Original: These could be pictures grouped by author, clinical outcomes grouped by the patient, or ﬁlm ratings grouped by user.
Revised: These could be pictures grouped by author, clinical outcomes grouped by the patient, or ﬁlm ratings grouped by user.

Original: We train a representation network r ( x n ) that encodes a group of observations { x n, 1 , . . .
Revised: We train a representation network r ( x n ) that encodes a group of observations { x n, 1 , . . .

Original: x n,K n } into one group code u n and a set of instance codes { z n, 1 , . . .
Revised: x n,K n } into one group code u n and a set of instance codes { z n, 1 , . . .

Original: z n,K n } , one for each observation.
Revised: z n,K n } , one for each observation.

Original: We want u to capture only the variation across groups and z only the variation within groups.
Revised: We want u to capture only the variation across groups and z only the variation within groups.

Original: The current state-of-the-art approaches for group disentanglement train the representation network r by using it as the variational latent posterior distribution in a Variational Autoencoder (Bouchacourt et al., 2018; Hosoya, 2019; Németh, 2020).
Revised: The current state-of-the-art approaches for group disentanglement train the representation network r by using it as the variational latent posterior distribution in a Variational Autoencoder (Bouchacourt et al., 2018; Hosoya, 2019; Németh, 2020).

Original: They assume a hierarchical generative model whereby the observation x n,k is generated by combining a group latent variable u n and an independent instance latent variable z n,k
Revised: They assume a hierarchical generative model whereby the observation x n,k is generated by combining a group latent variable u n and an independent instance latent variable z n,k

Original: (Figure 1-left).
Revised: (Figure 2-left).

Original: The standard setup involves training the variational latent posterior q ( u n , z n, 1: K n | x n ) by maximising a lower bound to the data likelihood (Kingma & Welling, 2014;
Revised: The standard setup involves training the variational latent posterior q ( u n , z n, 1: K n | x n ) by maximising a lower bound to the data likelihood (Kingma & Welling, 2014;

Original: Rezende et al., 2014).
Revised: Rezende et al., 2014).

Original: In our work, we show that the variational latent posterior, as conceived in existing models, is unsuited to datasets with conditional shift .
Revised: In our work, we show that the variational latent posterior, as deﬁned in existing models, is unsuited to datasets with conditional shift .

Original: 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
Revised: 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

Original: i ( z | x ) (cid:54) =
Revised: i ( z | x ) (cid:54) =

Original: p j ( z | x ) where i, j are two groups (Zhang et al., 2013; Gong et al., 2016).
Revised: p j ( z | x ) where i, j are two groups (Zhang et al., 2013; Gong et al., 2016).

Original: In our case, the conditional instance distribution for group i , which is p i ( z | x ) , corresponds to p ( z i,k
Revised: In our case, the conditional instance distribution for group i , which is p i ( z | x ) , corresponds to p ( z i,k

Original: | x i,k , u i ) where k is a given instance in the group.
Revised: | x i,k , u i ) where k is a given instance in the group.

Original: Conditional shift occurs in many real-world datasets that we would like to group-disentangle.
Revised: Conditional shift occurs in many real-world datasets that we would like to group-disentangle.

Original: For example, in the context of collaborative ﬁltering, if we want to infer the quality of an item based on a user rating, we should take into account how the same user rated other items (Koren et al., 2009).
Revised: For example, in the context of collaborative ﬁltering, if we want to infer the quality of an item based on a user rating, we should take into account how the same user rated other items (Koren et al., 2009).

Original: Different users might give the same rating to items of different quality based on their preferences, so inferring the item representation without conditioning on the user will produce widely varying estimates across users (Pinheiro & Bates, 2001).
Revised: Different users might give the same rating to items of different quality based on their preferences, so inferring the item representation without conditioning on the user will produce widely varying estimates across users (Pinheiro & Bates, 2001).

Original: Existing VAE-based methods fail to 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 1-middle): When deﬁning the variational latent posterior, existing works based on the Group VAE (GVAE) (Bouchacourt et al.,2018; Hosoya, 2019; Németh, 2020;
Revised: Existing VAE-based methods fail to 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 2-middle): When deﬁning the variational latent posterior, existing works based on the Group VAE (GVAE) (Bouchacourt et al.,

Original: Batmanghelich, 2020) make the assumption that the group and instance variables are conditionally independent given the observations (Figure 1-middle):
Revised: Batmanghelich, 2020) make the assumption that the group and instance variables are conditionally independent given the observations (Figure 2-middle):

Original: [Equation]
Revised: [Equation]

Original: The limitations of this assumption have not been identiﬁed so far in the literature because the datasets used to test disentanglement, such as Shapes3D (Kim & Mnih, 2018), SmallNORB (LeCun et al., 2004), dSprites (Higgins et al., 2017), Cars3D (Reed et al., 2015), MPI3D (Gondal et al., 2019), have the property that one image is always sufﬁcient to accurately infer its latent variables.
Revised: The limitations of this assumption have not been identiﬁed so far in the literature because the datasets used to test disentanglement, such as Shapes3D (Kim & Mnih, 2018), SmallNORB (LeCun et al., 2004), dSprites (Higgins et al., 2017), Cars3D (Reed et al., 2015), MPI3D (Gondal et al., 2019), have the property that one image is always sufﬁcient to accurately infer its latent variables.

Original: For example, we only require a single image from the MPI3D dataset to uniquely identify the colour, position, and rotation of the depicted object.
Revised: For example, we only require a single image from the MPI3D dataset to uniquely identify the colour, position, and rotation of the depicted object.

Original: We present the following contributions: We show that even a simple dataset with conditional shift can prohibit current methods from learning disentangled representations (Figure 2b).
Revised: 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.

Original: We approach the task of learning fair representations of students from different schools/socio-economic backgrounds. We propose a new model, the Context-Aware VAE (CxVAE), that has the ability to produce group-disentangled representations of datasets with conditional shift (Figure 2c). We show that  conditional shiftdirectly causes our model’s improvement in  performance over existing methods.
Revised: In the ﬁrst instance, we show that only our method is able to correctly disentangle between object-colour and spotlight-colour in the 3DIdent dataset (Zimmermann et al., 2021), illustrated in Figure 1. Then, we use the task of fair comparisons between student test-scores (Figure 3) to show that the amount of conditional shift in the dataset determines the performance gap between our model and the rest.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Group Disentanglement.
Revised: Group Disentanglement.

Original: This class of problems comes under different names: style-content disentanglement (Tenenbaum & Freeman, 2000), content-transformation disentanglement (Hosoya, 2019), and disentanglement with group supervision (Shu et al., 2020), to name a few.
Revised: This class of problems comes under different names: style-content disentanglement (Tenenbaum & Freeman, 2000), content-transformation disentanglement (Hosoya, 2019), and disentanglement with group supervision (Shu et al., 2020), to name a few.

Original: Recent work (Shu et al., 2020; Locatello et al., 2020) has contextualised GID as a subproblem of weaklysupervised disentanglement, where disentangled representations are learned with the help of nondatapoint supervision (e.g. grouping, ranking, restricted labelling).
Revised: Recent work (Shu et al., 2020; Locatello et al., 2020) has contextualised group disentanglement as a subproblem of weakly-supervised disentanglement, where disentangled representations are learned with the help of non-datapoint supervision (e.g. grouping, ranking, restricted labelling).

Original: Early work in this area focused on separating between visual concepts (Kulkarni et al., 2015; Reed et al., 2015).
Revised: Early work in this area focused on separating between visual concepts (Kulkarni et al., 2015; Reed et al., 2015).

Original: This area has received renewed interest after the theoretical impossibility result of Locatello et al.
Revised: This area has received renewed interest after the theoretical impossibility result of Locatello et al.

Original: and the identiﬁability proofs of Khemakhem et al.
Revised: and the identiﬁability proofs of Khemakhem et al.

Original: (2020) and Mita et al.
Revised: (2020) and Mita et al.

Original: A key aspect of recent weakly-supervised models is the interpretation of the grouping as a signal of similarity between datapoints (Chen & Batmanghelich, 2020). 
Revised: A key aspect of recent weakly-supervised models is the interpretation of the grouping as a signal of similarity between datapoints (Chen & Batmanghelich, 2020). [Figure]

Original: Conditioning on the Group Variable.
Revised: Conditioning on the Group Variable.

Original: While conditioning the instance encoder on the group variable is common in the areas of semi-supervised learning and fair representations (Kingma et al., 2014;
Revised: While conditioning the instance encoder on the group variable is common in the areas of semi-supervised learning and fair representations (Kingma et al., 2014;

Original: Louizos et al., 2016), we are the ﬁrst to apply it to unsupervised group disentanglement where explicit group labels are not available.
Revised: Louizos et al., 2016), we are the ﬁrst to apply it to unsupervised group disentanglement where explicit group labels are not available.

Original: In the ﬁeld of sequence disentanglement, stateof-the-art methods (Hsu et al., 2017; Denton & Birodkar, 2017; Li & Mandt, 2018) infer the [Figure] a d i n g T
Revised: In the ﬁeld of sequence disentanglement, stateof-the-art methods (Hsu et al., 2017; Denton & Birodkar, 2017; Li & Mandt, 2018) infer the instance variable (capturing shorter timescales) conditionally on the group variable (capturing longer timescales).

Original: [Figure]

Original: s t S c

Original: [Figure]

Original: r e [Figure] instance variable (capturing shorter timescales) conditionally on the group variable (capturing longer timescales). Recent works in weakly-supervised disentanglement (Shu et al., 2020; Locatello et al., 2019; Roeder et al., 2019) also condition the instance variable on the group, but their group variable is a discrete variable used for selection rather than a representation.
Revised: Recent works in weakly-supervised disentanglement (Shu et al., 2020; Locatello et al., 2019; Roeder et al., 2019) also condition the instance variable on the group, but their group variable is a discrete variable used for selection rather than a representation.

Original: It marks which units of the instance representation are common within the group and which are free to vary.
Revised: It marks which units of the instance representation are common within the group and which are free to vary.

Original: We argue that this is not sufﬁcient to account for the variation in p ( v | x, u ) produced by the conditional shift, so we include
Revised: We argue that this is not sufﬁcient to account for the variation in p ( v | x, u ) produced by the conditional shift, so we include

Original: AdaGVAE (Locatello et al., 2020) in our evaluation for comparison (see Section 6).
Revised: AdaGVAE (Locatello et al., 2020) in our evaluation for comparison (see Section 7).

Original: Conditional Shift Our instance encoder conditioned on the group variable is a new strategy to deal with conditional shift.
Revised: Conditional Shift. Our instance encoder conditioned on the group variable is a new strategy to deal with conditional shift.

Original: This problem has been studied extensively in the context of supervised learning (Zhang et al., 2013;
Revised: This problem has been studied extensively in the context of supervised learning (Zhang et al., 2013;

Original: Gong et al., 2016).
Revised: Gong et al., 2016).

Revised: However, we are the ﬁrst to explore the effect of conditional shift on unsupervised learning.

Original: Methods for mitigating the effects of conditional shift typically focus on learning domain-invariant representations (Ben-David et al., 2009).
Revised: Methods for mitigating the effects of conditional shift typically focus on learning domain-invariant representations (Ben-David et al., 2009).

Original: However,
Revised: However,

Original: Zhao et al.
Revised: Zhao et al.

Original: (2019) show that learning a domain-invariant representation is not sufﬁcient for learning a correct mapping between instance variables from different groups.
Revised: (2019) show that learning a domain-invariant representation is not sufﬁcient for learning a correct mapping between instance variables from different groups.

Original: Image Translation.
Revised: Image Translation.

Original: Note that our variational latent posterior is different from the one used in
Revised: Note that our variational latent posterior is different from the one used in

Original: COCO-FUNIT (Saito et al., 2020).
Revised: COCO-FUNIT (Saito et al., 2020).

Original: The authors are motivated by the same limitations with existing works as we are, namely that unsupervised translation methods struggle to disentangle under conditional shift.
Revised: The authors are motivated by the same limitations with existing works as we are, namely that unsupervised translation methods struggle to disentangle under conditional shift.

Original: However, because they train explicitly for translation rather than disentanglement, they arrive at a different solution than ours.
Revised: However, because they train explicitly for translation rather than disentanglement, they arrive at a different solution than ours.

Original: When performing a 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.
Revised: When performing a 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.

Original: This mechanism produces impressive results on 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 are no source and target groups in the training set.
Revised: This mechanism produces impressive results on 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 are no source and target groups in the training set.

Original: Regardless, we evaluate COCO-FUNIT on the test score dataset and show that our model outperforms it both in terms of disentanglement and translation.
Revised: Regardless, we evaluate COCO-FUNIT on the test score dataset and show that our model outperforms it both in terms of disentanglement and translation.

Original: B ACKGROUND 3. F AIR C OMPARISONS B ETWEEN S TUDENTS

Original: 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).

Original: The typical assumption in the literature is that each school has a 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.

Original:  The instance representation z n,k should reﬂect 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 circumstances associated with attending one school or another, such as afﬂuence, location, or curriculum (Raudenbush & Willms, 1995; Braun et al., 2006). Group disentanglement has the potential to reduce the computational costs of 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 one forward pass through the network for every additional student.
Revised:  B ACKGROUND

Original: Contrast this with running a separate optimisation routine for the latent variables of each student (Gelman et al., 1995).

Original: 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).

Original: Current methods fail to learn disentangled representations on this data.

Original: Figure 2b shows the GVAE model (Hosoya, 2019) incorrectly translating the scores from one school to another.

Original: Translation is a well-established downstream task for evaluating disentanglement (Tenenbaum & Freeman, 2000).

Original: In the case of test scores, translation corresponds to the counterfactual question “What score would student k from school A have obtained if they had attended the typical school (i.e. a school whose scores are distributed according to N (0 , I 2 ) )?”.

Original: We do this by generating a new score by combining the instance representation of student k from school n with the group representation of the typical school.

Original: Raudenbush & Willms (1995) use translation to directly compare students from different schools.

Original: Notice that the unconditional GVAE scrambles the order of the scores and fails to capture the N (0 , I 2 ) distribution of the typical school.

Original: We propose a modiﬁcation to the variational latent posterior in order to enable the GVAE to learn disentangled representations of test score data.

Original: Looking at the data (Figure 2a), it is clear that we are dealing with the conditional shift scenario.

Original: The same reading score could be obtained by either a highachieving student from school C or a low-achieving student from school B .

Original: Inferring the aptitude of the student requires knowing the distribution of scores within each school.

Original: In order to account for the variation across groups, we introduce the Context-Aware Variational Autoencoder (CxVAE), whose instance encoder is conditioned on the group representation (Figure 1-right).

Original: This reﬂects the correct factorisation of the generative latent posterior p ( u n , z n, 1: K n | x n, 1: K n ) , thus making no assumptions about the relationship between the group and the instance.

Original: We observe our variational model successfully translating test scores in Figure 2c.

Original: 3. G ROUP -I NSTANCE G ENERATIVE M ODEL

Original: The Group-Instance Generative Model (Bouchacourt et al., 2018; Hosoya, 2019) is a multi-level model that uses two latent variables to generate grouped data: the instance variable z n,k ∼ N (0 , 1) controls the variation within groups, and the group variable u n ∼ N (0 , 1) controls the variation across groups (Figure 1-left).
Revised: The Group-Instance Generative Model (Bouchacourt et al., 2018; Hosoya, 2019) is a multi-level model that uses two latent variables to generate grouped data: the instance variable z n,k ∼ N (0 , 1) controls the variation within groups, and the group variable u n ∼ N (0 , 1) controls the variation across groups (Figure 2-left).

Original: The likelihood of a group x n, 1: K n is:
Revised: The likelihood of a group x n, 1: K n is:

Original: [Equation]
Revised: [Equation]

Original: 3. V ARIATIONAL I NFERENCE
Revised: 3. V ARIATIONAL I NFERENCE

Original: Because the exact likelihood is intractable, the standard approach to train the group-instance generative model is with a Variational Autoencoder (Kingma & Welling, 2014; Rezende et al., 2014) which performs optimisation by introducing a variational latent posterior q ( u n , z n, 1: K n | x ) and maximizing the Evidence Lower Bound (Jordan et al., 2004):
Revised: Because the exact likelihood is intractable, the standard approach to train the group-instance generative model is with a Variational Autoencoder (Kingma & Welling, 2014; Rezende et al., 2014) which performs optimisation by introducing a variational latent posterior q ( u n , z n, 1: K n | x ) and maximizing the Evidence Lower Bound (Jordan et al., 2004):

Original: [Equation]
Revised: [Equation]

Original: Existing methods use a class of variational distributions that assume conditional independence between the latent variables.
Revised: Existing methods use a class of variational distributions that assume conditional independence between the latent variables.

Original: [Equation]
Revised: [Equation]

Original:  C ONTEXT -A WARE V ARIATIONAL A UTOENCODER
Revised:  C ONTEXT -A WARE V ARIATIONAL A UTOENCODER

Original: We propose a new model which can perform well on group-confounded problems.
Revised: We propose a new model which can perform well on group-confounded problems.

Original: We call our model the Context-Aware Variational Autoencoder (CxVAE).
Revised: We call our model the Context-Aware Variational Autoencoder (CxVAE).

Original: Its deﬁning feature is a variational latent posterior whose instance variable is conditioned on the group variable:
Revised: Its deﬁning feature is a variational latent posterior whose instance variable is conditioned on the group variable:

Original: [Equation]
Revised: [Equation]

Original: Thus, the instance encoder is implemented as a network which takes as input the concatenation of the observation x n,k and previously sampled group representation u n .
Revised: Thus, the instance encoder is implemented as a network which takes as input the concatenation of the observation x n,k and previously sampled group representation u n .

Original: [Equation]
Revised: [Equation]

Original: This form of the variational distribution reﬂects the correct factorisation of the generative model.
Revised: This form of the variational distribution reﬂects the correct factorisation of the generative model.

Original: It has the potential to learn the true generative latent posterior, which is disentangled by deﬁnition.
Revised: It has the potential to learn the true generative latent posterior, which is disentangled by deﬁnition.

Original: The
Revised: The

Original: Evidence Lower Bound for our model is
Revised: Evidence Lower Bound for our model is

Original: [Equation]
Revised: [Equation]

Original:   M ETHODOLOGY
Revised:  E VALUATION M ETHODOLOGY

Original: We show that by making the instance encoder conditional on the inferred group variable, we obtain a considerable gain in translation accuracy and a marked improvement in disentanglement.
Revised: We show that by making the instance encoder conditional on the inferred group variable, we obtain a considerable gain in translation accuracy and a marked improvement in disentanglement.

Original: We also demonstrate that the gap in performance between our model and other group disentanglement methods is caused by conditional shift in the data-generating process.
Revised: We also demonstrate that the gap in performance between our model and other group disentanglement methods is caused by conditional shift in the data-generating process.

Original: 5.  S ETUP
Revised: 5. M ODEL S ETUP

Original: We compare our conditional CxVAE with the state-of-the-art in group disentanglement, namely the Group VAE (Hosoya, 2019; Bouchacourt et al., 2018), COCO-FUNIT (Saito et al., 2020), and
Revised: We compare our conditional CxVAE with the state-of-the-art in group disentanglement, namely the Group VAE (Hosoya, 2019; Bouchacourt et al., 2018), COCO-FUNIT (Saito et al., 2020), and

Original: AdaGVAE (Locatello et al., 2019).
Revised: AdaGVAE (Locatello et al., 2019).

Original: We implement all networks (encoders and decoders) as MLPs with 3 hidden layer of 32 activations each.

Original: As in Hosoya (2019), the group encoder is applied to each datapoint in the group and then all the outputs are averaged.
Revised: As in Hosoya (2019), the group encoder is applied to each datapoint in the group and then all the outputs are averaged.

Original: The group variable will havedimensions (2 for b n and 2 for c n ) and the instance variable will have 2 dimension (for a n,k ).

Original: For all experiments, our CxVAE will be a modiﬁed GVAE such that the group variable u n is concatenated with the observation x n,k and fed into the instance encoder in order to compute the instance variable z n,k .
Revised: For all experiments, our CxVAE will be a modiﬁed GVAE such that the group variable u n is concatenated with the observation x n,k and fed into the instance encoder in order to compute the instance variable z n,k .

Original: 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
Revised: For sampling the variational latent posteriors, we use the standard reparametrisation trick. We use an Adam optimiser with learning rate of 1e-4 with β 1 = 0 . 9 , β 2 = 0 .

Original: [Figure] [Figure] GVAE AdaGVAE COCO-FUNIT
Revised: For the 3DIdent dataset, we implement all networks (encoders and decoders) as convolutional nets with 4 hidden layers and 64 ﬁlters each.

Original: CxVAE 0.

Original: [Figure]

Original: [Figure]

Original: GVAE

Original: AdaGVAE COCO-FUNIT CxVAE 0. [Figure] [Figure]
Revised: Both latent variables have 16 latent dimensions.

Original: [Figure]

Original: [Figure] and the networks. We use the same 100 seeds in each model. This gives 1000 measurements to plot in Figure 3.
Revised: For the test-score dataset, we use MLPs with 3 hidden layer of 32 activations each.

Original: Training and testing take approximately 2 hours to run on an Intel i9 CPU. 5. D ATASET We generate our dataset of test scores using the classic “varying intercept, varying slope” mixedeffects model (Laird & Ware, 1982; Pinheiro & Bates, 2001; Gelman & Hill, 2006). This is a well-established approach formodelling student scores x n,k as a function of individual aptitude a n,k and school-level characteristics ( b n , c n ) (e.g. afﬂuence, curriculum, location, etc.)
Revised: The group variable will have 4 dimensions (2 for b n and 2 for c n ) and the instance variable will have 2 dimension (for a n,k ).

Original: (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:
Revised: We train each model for 64 epochs, and use the last 10 epochs for evaluation.

Original: [Equation] score of student k [Equation] school n.
Revised: Additionally, we run the experiment for 100 different random seeds initialisations, both for the data generating process and the networks.

Original: [Equation] [Equation] is the aptitude of [Equation]
Revised: Conﬁdence intervals are computed by resampling train-test splits, weight initialisations and sampling seeds.

Original: standard [Equation] For the evaluation procedure, we use the above model to generate N = 32 , 768 values for b n , c n . For each school n , we generate M = 128 values for a n,k .
Revised: We use the same 100 seeds in each model.

Original: We then randomly select half of the schools to assemble a training dataset with 2 , 097 , 152 scores split across 16 , 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. We have chosen to evaluate on a synthetic dataset because it allows for ﬁne control over the parameters of the data-generating process (especially relevant in Section 6.1) and it also enables us to measure the quality of disentanglement using the Mutual Information Gap (Chen et al., 2018).
Revised: This gives 1000 measurements to plot in Table 1.

Original: 5.  M ETRICS
Revised: 5. E VALUATION M ETRICS

Original: We compare the models with respect to 3 different criteria: 1) How well does the model ﬁt the holdout data?, 2) How disentangled are the representations inferred by the encoder?
Revised: We compare the models with respect to 3 different criteria: 1) How well does the model ﬁt the holdout data?, 2) How disentangled are the representations inferred by the encoder?

Original: , and 3) How well can the model answer the counterfactual question "What score would a student have obtained if they had gone to a different school?".
Revised: , and 3) How well can the model answer the counterfactual question "What score would a student have obtained if they had gone to a different school?".

Original: We assess each criterion with a different evaluation metric.
Revised: We assess each criterion with a different evaluation metric.

Original: Fitting the holdout data.
Revised: Fitting the holdout data.

Original: 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 of the holdout set.
Revised: 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 of the holdout set.

Original: Disentangled representations.
Revised: Disentangled representations.

Original: We use the Mutual Information Gap (Chen et al., 2018) to measure the quality of the disentanglement.
Revised: We use the Mutual Information Gap (Chen et al., 2018) to measure the quality of the disentanglement.

Original: We measure empirically the amount of mutual information between the inferred latent variables u , z and the school characteristics ( b , c ) which represent the ground-truth group factor.
Revised: We measure empirically the amount of mutual information between the inferred latent variables u , z and the school characteristics ( b , c ) which represent the ground-truth group factor.

Original: Consequently, the goal is to have maximum mutual information between the group variable u and the ground-truth, and minimum mutual information between the instance variables z and the ground-truth.
Revised: Consequently, the goal is to have maximum mutual information between the group variable u and the ground-truth, and minimum mutual information between the instance variables z and the ground-truth.

Original: The gap between the two (normalized with the entropy of the ground-truth factors) is the metric of disentanglement:
Revised: The gap between the two (normalized with the entropy of the ground-truth factors) is the metric of disentanglement:

Original: [Equation]
Revised: [Equation]

Original: Since the data-generating process is known, the mutual information between the inferred group variable and the ground-truth group variable I ( u ; b , c ) is straightforward to implement by following the process of Chen et al.
Revised: Since the data-generating process is known, the mutual information between the inferred group variable and the ground-truth group variable I ( u ; b , c ) is straightforward to implement by following the approach from Chen et al.

Original: We measure only the mutual information between the ground-truth group factor 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 factors ( b , c ) .
Revised: We measure only the mutual information between the ground-truth group factor 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 factors ( b , c ) .

Original: Counterfactual question.
Revised: Translation task.

Original: We measure how well the learned representations can answer the question “What would the score of student k from school n have been if they had attended the typical school (the school with scores distributed according to N (0 , I 2 ) )?”.
Revised: We measure how well the learned representations can answer the question “What would the score of student k from school n have been if they had attended the typical school (the school with scores distributed according to N (0 , I 2 ) )?”.

Original: This problem, also known as translation, is a commonly used downstream task for disentangled representations (Tenenbaum & Freeman, 2000).
Revised: This problem, also known as translation, is a commonly used downstream task for disentangled representations (Tenenbaum & Freeman, 2000).

Original: We translate the score of student k 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 generative factors.
Revised: We translate the score of student k to the typical school, and then take the mean squared error against [Table] the ground-truth translation, which is generated when the data is generated using the ground-truth generative factors.

Original: 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).
Revised: 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).

Original: In order to obtain the translation, we ﬁrst infer the instance variable of student k from school n and the group variable of the typical school.
Revised: In order to obtain the translation, we ﬁrst infer the instance variable of student k from school n and the group variable of the typical school.

Original: We then feed the two variables to the decoder.
Revised: We then feed the two variables to the decoder.

Original: For evaluation, we translate all the scores from each school n and then measure the distance between the predicted translation and the ground-truth translation.
Revised: For evaluation, we translate all the scores from each school n and then measure the distance between the predicted translation and the ground-truth translation.

Original: We compute the total error as an average over all the translation errors.
Revised: We compute the total error as an average over all the translation errors.

Original: We can also use translation as an additional qualitative comparison between CxVAE and other group disentanglement methods, which can be seen in Figure 2.
Revised: We can also use translation as an additional qualitative comparison between CxVAE and other group disentanglement methods, which can be seen in Figure 3.

Revised:  D ISENTANGLING A MBIGUOUS O BSERVATIONS

Revised: We show that our CxVAE is able to group-disentangle in a setting where conditional shift produces ambiguous observations.

Revised: The dataset comprises images from the 3DIdent dataset (Zimmermann et al., 2021) depicting teapots of different colours being lit by spotlights of different colours.

Original: R ESULTS
Revised: Within a group, the images have the same spotlight colour and only the colour of the teapot varies. The goal is 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-world 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 a classiﬁer that predicts the colour of the object.

Revised: The better the disentanglement, the better performance this predictor would have at identifying the true colour of the object.

Revised: This is a difﬁcult problem because the data exhibits conditional shift.

Revised: The same exact image could have been generated by different combination of spotlight colour and object colour, as can be seen in

Revised: Figure 1.

Revised: This makes it difﬁcult to identify the true colour of the object by just looking at one single image.

Revised: Indeed all previous methods, which infer the instance representation solely from the current image, fail to learn disentangled representations of this dataset.

Revised: [Figure]

Revised: e a d i n

Revised: [Figure]

Revised: T

Revised: [Figure]

Revised: s t S c

Revised: o r e

Revised: [Figure]

Revised: T

Revised: [Figure]

Original: Our CxVAE produces considerable improvements over the competing methods in terms of ﬁtting the holdout set, disentangled representations, translation accuracy and predicting the generative factors (Figure 3).
Revised: (c [Figure] Successful translation (ours). The results in Table 1 show that our 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.

Revised: Our model’s MIG score is much higher than the one produced by competing models, and the performance gap is greater than the 95% conﬁdence interval for any one of these models.

Revised:  C ONDITIONAL S HIFT C AUSES THE G AP IN P ERFORMANCE

Revised: Our CxVAE produces considerable improvements over the competing methods in terms of ﬁtting the holdout set, disentangled representations, translation accuracy and predicting the generative factors (Table 1).

Original: While the scores of the existing methods cluster together, the gap between them and the
Revised: While the scores of the existing methods cluster together, the gap between them and the

Original: CxVAE is larger than the 95% conﬁdence interval of any method.
Revised: CxVAE is larger than the 95% conﬁdence interval of any method.

Original: 6. C ONDITIONAL S HIFT C AUSES THE G AP IN P ERFORMANCE
Revised: 7. F AIR C OMPARISONS B ETWEEN S TUDENTS

Revised: 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).

Revised: The typical assumption in the literature is that each school has a 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.

Revised: The instance representation z n,k should reﬂect the aptitude of student k from school n independently of what school the student had attended.

Revised: This analysis is crucial for university admission boards aiming to judge students based on their aptitude regardless of the socio-economic circumstances associated with attending one school or another, such as afﬂuence, location, or curriculum (Raudenbush & Willms, 1995; Braun et al., 2006).

Revised: Group disentanglement has the potential to reduce the computational costs of this analysis compared to the current variational inference methods for mixed-effects models (Gelman & Hill, 2006).

Revised: The variational latent posterior of a GVAE can perform scalable inference of the student and school representations requiring just one forward pass through the network for every additional student.

Revised: Contrast this with running a separate optimisation routine for the latent variables of each student (Gelman et al., 1995).

Revised: 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).

Revised: Current methods fail to learn disentangled representations on this data.

Revised: Figure 3b shows the GVAE model (Hosoya, 2019) incorrectly translating the scores from one school to another.

Revised: Translation is a well-established downstream task for evaluating disentanglement (Tenenbaum & Freeman, 2000).

Revised: In the case of test scores, translation corresponds to the counterfactual question “What score would student k from school A have obtained if they had attended the typical school (i.e. a school whose scores are distributed according to N (0 , I 2 ) )?”.

Revised: We do this by generating a new score by combining the instance representation of student k from school n with the group representation of the typical school.

Revised: Raudenbush & Willms (1995) use translation to directly compare students from different schools.

Revised: Notice that the unconditional GVAE scrambles the order of the scores and fails to capture the N (0 , I 2 ) distribution of the typical school.

Revised: We propose a modiﬁcation to the variational latent posterior in order to enable the GVAE to learn disentangled representations of test score data.

Revised: Looking at the data (Figure 3a), it is clear that we are dealing with the conditional shift scenario.

Revised: The same reading score could be obtained by either a highachieving student from school C or a low-achieving student from school B .

Revised: Inferring the aptitude of the student requires knowing the distribution of scores within each school.

Revised: In order to account for the variation across groups, we introduce the Context-Aware Variational Autoencoder (CxVAE), whose instance encoder is conditioned on the group representation (Figure 2-right).

Revised: This reﬂects the correct factorisation of the generative latent posterior p ( u n , z n, 1: K n | x n, 1: K n ) , thus making no assumptions about the relationship between the group and the instance.

Revised: We observe our variational model successfully translating test scores in Figure 3c.

Revised: 7. D ATASET

Revised: We generate our dataset of test scores using the classic “varying intercept, varying slope” mixedeffects model (Laird & Ware, 1982; Pinheiro & Bates, 2001; Gelman & Hill, 2006).

Revised: This is a well-established approach for modelling student scores x n,k as a function of individual aptitude a n,k and school-level characteristics ( b n , c n ) (e.g. afﬂuence, curriculum, location, etc.)

Revised: (Raudenbush &

Revised: Willms, 1995; Braun et al., 2006).

Revised: We choose this model for its simplicity and for the wide variety of phenomena to which it can be applied.

Revised: All the scores and factors are 2-dimensional vectors, with one component for the maths score and another for the reading score:

Revised: [Equation]

Revised: For the evaluation procedure, we use the above model to generate N = 32 , 768 values for b n , c n .

Revised: For each school n , we generate M = 128 values for a n,k .

Revised: We then randomly select half of the schools to assemble a training dataset with 2 , 097 , 152 scores split across 16 , 384 schools.

Revised: We take the other half of schools to create the holdout dataset, so that every testing school and student are unseen during training.

Revised: We have chosen to evaluate on a synthetic dataset because it allows for ﬁne control over the parameters of the data-generating process (especially relevant in Section 7.3) and it also enables us to measure the quality of disentanglement using the Mutual Information Gap (Chen et al., 2018).

Revised: 7. C ONDITIONAL S HIFT C AUSES THE G AP IN P ERFORMANCE

Original: We show, by modifying the data-generating process (11), that conditional shift explains the increased performance of the CxVAE.
Revised: We show, by modifying the data-generating process (6), that conditional shift explains the increased performance of the CxVAE.

Original: We insert a hyper-parameter λ to control the strength of the conditional shift; λ = 1 means the conditional shift stays the same as in the previous experiment, while λ =means there is no conditional shift.
Revised: We insert a hyper-parameter λ to control the strength of the conditional shift; λ = 1 means the conditional shift stays the same as in the previous experiment, while λ =means there is no conditional shift.

Original: Consider the case where the maths score only depends on the school, and the reading score only depends on the student.
Revised: Consider the case where the maths score only depends on the school, and the reading score only depends on the student.

Original: 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 proﬁle from the maths score.
Revised: 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 proﬁle from the maths score.

Original: We use this as an extreme case of lack of conditional shift, and insert a hyper-parameter λ in our data-generating process that will continuously move between this case and the original case.
Revised: We use this as an extreme case of lack of conditional shift, and insert a hyper-parameter λ in our data-generating process that will continuously move between this case and the original case.

Original: Our modiﬁed data-generating process is
Revised: Our modiﬁed data-generating process is

Original: [Figure]
Revised: [Figure]

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: (17) where ˆ · denotes the elementwise power operation and the generative factors ( a n,k , b n , c n , (cid:15) n,k ) are sampled as in (11).
Revised: where ˆ · denotes the elementwise power operation and the generative factors ( a n,k , b n , c n , (cid:15) n,k ) are sampled as in (6).

Original: This model has no conditional shift when λ = 0 because each ground-truth factor controls a separate component of the data.
Revised: This model has no conditional shift when λ = 0 because each ground-truth factor controls a separate component of the data.

Original: Inferring the student aptitude requires only the reading score and can ignore the school characteristics.
Revised: Inferring the student aptitude requires only the reading score and can ignore the school characteristics.

Original: When λ = 1 , the problem exhibits conditional shift in exactly the same way as in (11).
Revised: When λ = 1 , the problem exhibits conditional shift in exactly the same way as in (6).

Original: If our hypothesis is correct, that the conditional shift causes the performance gap between CxVAE and other group disentanglement methods, then the gap should decrease as λ approaches 0.
Revised: If our hypothesis is correct, that the conditional shift causes the performance gap between CxVAE and other group disentanglement methods, then the gap should decrease as λ approaches 0.

Original: The measurements displayed in Figure 4 conﬁrm our expectations.
Revised: The measurements displayed in Figure 4 conﬁrm our expectations.

Original: For low values of λ the performance of our CxVAE is evenly matched to the GVAE.
Revised: For low values of λ the performance of our CxVAE is evenly matched to the GVAE.

Original: As λ increases, CxVAE metrics remain stable while
Revised: As λ increases, CxVAE metrics remain stable while

Original: GVAE performance decreases substantially.
Revised: GVAE performance decreases substantially.

Original: It is clear that the degree of confounding in the dataset explains the performance gain that we see in the CxVAE.
Revised: It is clear that the degree of confounding in the dataset explains the performance gain that we see in the CxVAE.

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: In this work, we show empirically that conditioning the instance encoder on the group variable produces group-disentangled representations on datasets with conditional shift.
Revised: In this work, we show empirically that conditioning the instance encoder on the group variable produces group-disentangled representations on datasets with conditional shift.

Original: 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.
Revised: 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.

Original: Our evaluation is run on the downstream task of extracting student aptitudes from a dataset of test scorse grouped by school, a problem on which group-instance models have not been applied before.
Revised: Our evaluation is run on the downstream task of extracting student aptitudes from a dataset of test scorse grouped by school, a problem on which group-instance models have not been applied before.

Original: The main limitation of our work is that we perform evaluation on a synthetic dataset of student scores rather than real data.
Revised: The main limitation of our work is that we perform evaluation on a synthetic dataset of student scores rather than real data.

Original: Although this is a justiﬁable choice with respect to evaluation (it gives us access to the ground-truth values of the latent variables), future work should focus on evaluating on real-world datasets with conditional shift, such as user-item ratings (Koren et al., 2009).
Revised: Although this is a justiﬁable choice with respect to evaluation (it gives us access to the ground-truth values of the latent variables), future work should focus on evaluating on real-world datasets with conditional shift, such as user-item ratings (Koren et al., 2009).

Original: A A PPENDIX
Revised: A A PPENDIX

Original: You may include other additional sections here.
Revised: You may include other additional sections here.