Contrastive Mixture of Posteriors for Counterfactual Inference, Data Integration and Fairness
Authors: Adam Foster
Árpi Vezér
Craig A Glastonbury
Páidí Creed
Sam Abujudeh
Aaron Sim
Venue: NeurIPS 2021 Submitted (2021)
Keywords: Contrastive learning
variational autoencoders
counterfactual inference
data integration
fairness
batch effect correction
scRNA-seq
RNA-seq
genomics
omics
transcriptomics
invariant representations
TCGA
CCLE
single-cell

Abstract: Learning meaningful representations of data that can address challenges such as batch effect correction, data integration and counterfactual inference is a central problem in many domains including computational biology. Adopting a Conditional VAE framework, we identify the mathematical principle that unites these challenges: learning a representation that is marginally independent of a condition variable. We therefore propose the Contrastive Mixtures of Posteriors (CoMP) method that uses a novel misalignment penalty to enforce this independence. This penalty is defined in terms of mixtures of the variational posteriors themselves, unlike prior work which uses external discrepancy measures such as MMD to ensure independence in latent space. We show that CoMP has attractive theoretical properties compared to previous approaches, especially when there is complex global structure in latent space. We further demonstrate state of the art performance on a number of real-world problems, including the challenging tasks of aligning human tumour samples with cancer cell-lines and performing counterfactual inference on single-cell RNA sequencing data. Incidentally, we find parallels with the fair representation learning literature, and demonstrate CoMP has competitive performance in learning fair yet expressive latent representations.

31 present formal account of these challenges and show that, a great extent, they can be seen as different aspects of the same underlying problem, namely, that of learning a representation Contrastive Mixture of Posteriors for Counterfactual Inference, Data Integration and Fairness
Anonymous Author(s)
Affiliation Address email
Abstract
Learning meaningful representations of data that can address challenges such batch effect correction, data integration and counterfactual inference a central problem in many domains including computational biology. Adopting Condi- tional VAE framework, we identify the mathematical principle that unites these challenges: learning representation that is marginally independent of condition variable. We therefore propose the Contrastive Mixture of Posteriors (CoMP) method that uses novel misalignment penalty to enforce this independence. This penalty is defined in terms of mixtures of the variational posteriors themselves, unlike prior work which uses external discrepancy measures such MMD ensure independence in latent space. We show that CoMP has attractive theoretical properties compared previous approaches, especially when there complex global structure in latent space. We further demonstrate state of the art performance on a number of real-world problems, including the challenging tasks of aligning human tumour samples with cancer cell-lines and performing counterfactual in- ference on single-cell RNA sequencing data. Incidentally, we find parallels with the fair representation learning literature, and demonstrate CoMP has competitive performance in learning fair yet expressive latent representations.
2
4
6
8
10
11
13
14
15
17
Introduction
18
19
Large scale datasets describing the molecular properties of cells, tissues and organs state of health and disease are commonplace in computational biology. Referred to collectively as 'omics data,
20
21 thousands features are measured per sample and, as single-cell methodologies have developed,
is now typical to measure such features across 109-106 samples (1 Given these two properties of
23 'omics data, the need for scalable algorithms to learn meaningful low-dimensional representations that capture the variability of the data has grown. As such, Variational Autoencoders (VAFs) 34
24
25 have become an important tool for solving range of modelling problems in the biological sciences 55 6. 7 9 [0]. One such problem is utilising representations for counterfactual inference,
27 e.g. predicting how a certain cell or cell-type, observed only in the control, would have behaved when exposed drug [9] 10, II). Another key problem is removing batch effects spurious shifts in observations due to differing experimental conditions from data in order to integrate or compare multiple datasets
29
15).
33 that is marginally independent of condition variable (e.g- experimental batch, stimulated vs. control). Figure [CoMP] illustrates what this looks like in practice: the complete overlap of the cell 35 populations from different conditions in the latent space. This directly addresses batch correction, and in the case where also have generative model that maps from latent space back to the original
Submitted to 35th Conference on Neural Information Processing Systems (NeurIPS 2021). Do not distribute.

---

VAE
CVAE
trVAE
CoMP
control
stimulated
Figure 1: Latent representations of single-cell gene expression dataset under conditions (see
Section 6.2) From fully disjointed (VAE) to a well-mixed pair of distributions (CoMP).
37 data space, methods that solve this problem can also be applied to counterfactual inference [10]
38 This same mathematical requirement for independence also occurs in the fair representation learning literature, in which we seek representation that removes a sensitive attribute, e.g. gender.
39
Neither the VAE nor the conditional VAE (CVAE) [66] are typically successful at learning repre- sentations that achieve this desired independence, as shown in Figure | Despite the CVAE being theoretically able to remove batch effects, there is no constraint that prevents it from from separating different conditions in latent space. Existing methods use a penalty to encourage the CVAE to learn representations that overlap correctly in latent space, with Maximum Mean Discrepancy (MMD) [7] being the most common penalty, applied the VFAE [8] and the more recent tr VAE IO These methods, however, suffer from number of drawbacks: conceptually, they introduce an extraneous discrepancy measure that is not part of the variational inference framework; practically, they re- quire the choice of, and hyperparameter tuning for, MMI) kernel; empirically, whilst tr rAE is sigmificant improvement over an unconstrained CVAE, Figure [rrAE] shows that it can still fail to exactly align different conditions in latent space.
41
43
45
46
47
49
50
51 To overcome these difficulties, we introduce Contrastive Mixture of Posteriors (CoMP), a new method for learning aligned representations in CVAE framework. Our method features the novel CoMP misalignment penalty that forces the CVAE to remove batch effects. Inspired by contrastive learning [99 20), the penalty encourages representations from different conditions to be close, whilst representations from the same condition should be spread out. To achieve this, we approximate the requisite marginal distributions using mixtures of the variational posteriors themselves, leading to penalty that does not require an extraneous discrepancy measure or a separately tuned kernel. We prove that the CoMP penalty is stochastic upper bound weighted sum of KL divergences, so minimising the penalty minimises well-established statistical divergence measure. We analyse the training gradients of the CoMP and MMD penalties, finding key differences that help explain why CoMP gradients are generally more stable and better suited to datasets with complex global structure.
53
55
57
59
80
61
As shown in Figure [[ our method can achieve visually perfect alignment number of real-world biological datasets. We apply CoMP to two challenging biological problems: 1) aligning gene expression profiles between tumours and their corresponding cell-lines, as tackled in [21] and 65 2) estimating the gene expression profile of an unperturbed cell as if it had been treated with a chemical perturbation (counterfactual inference) [9]. We show that CoMP outperforms existing methods, achieving state-of-the performance on both tasks. Finally, given the connections to fair representation learning, we apply CoMP to the problem of learning a representation that is independent of gender in the UCI Adult Income dataset (22). showing that we can learn representation that is fully independent of the protected attribute whilst maintaining useful information for other prediction tasks. CoMP represents a conceptually simple and empirically powerful method for learning aligned representation, opening the door to answering high-value questions in biology and beyond.
83
67
88
89
70
71
72
Background
73
74
2.1
Variational Autoencoders and extensions
We begin by assuming that we have observations XI, XI,..., Xn. of an underlying data distribution. Variational autoencoders (VAEs)
76
S 4 explain the high-dimensional observations x, using low
dimensional representations z, The standard VAE places standard normal prior p(z)

---

Figure 2: Structural Equation Model for observation x under known condition e with unobserved latent variable In this model, and are independent in the prior.
78 the latent variable, and learns generative model (x(x)z) that reconstructs x using alongside
79 an inference network ta (z]x) that encodes x to z. Both and Φ are trained jointly by maximising
80 the ELBO, lower bound on marginal likelihood given by log (x) > (x(x) [log 2))
KL [g (x]x))[p(()) This can be maximised using stochastic optimisers [23 3 Various extensions
81
82
of the VAE have been proposed, such as the 3-VAE (24), which scales the KL term of the ELBO
by a hyperparameter B. Because the isotropic normal prior may limit the expressivity of the model
84 251, various authors have considered alternative priors. For example, (26) proposed the Variational
Mixture of Posteriors (VaMP) prior, that replaces the isotropic Gaussian with mixture of posteriors from the encoder network itself, evaluated at a number of learned greudo-inputs.
86
87
So far, we have assumed that the only data available are the observations XI,...,X,, X,,, but in many practical applications we may have additional information such as a condition label for each ob- servation. For example, in gene knock studies, we have information about which gene was targeted for deletion in each cell; in multi-batch experiments we have information about which exper- imental batch each samples was collected in. Thus, we augment our data by considering data pairs (X1, C1) ..., (Xn: G.) where x the same high-dimensional observation, and e a label indicating the condition experimental batch that x was collected under.
86
89
90
91
92
93
94
Whilst VAEs are theoretically able to model the pairs (x,, ci), it makes sense to build a model that explicitly distinguishes between the x and c. The simplest model is the Conditional VAE (CVAE)
95
96
16.. In this model, a conditional generative model po (x|2,, z, c) and a conditional inference network [z/2 c) are trained using a modified ELBO. A key observation for our work is that the CVAE has many different ways to model the data. For example, it can completely ignore the condition in pe
97
96
and yo. reducing to the original VAE. Assuming that x is not independent of c, this failure mode of
99
100
the CVAE would be apparent on a visualization of the representations. For example, different values
of might be visible as separate latent clusters, as shown in Figure ICVAE].
102
2.2
Counterfactual inference
103
(x, Q) represents an RNA transcript and the gene knock-out applied to the cell, a natural question to ask is "How would the transcript have differed if a different knock-out c' had been applied?" In general, counterfactual inference attempts to answer questions of the form "How would the data have changed if c; had been replaced by c'?" Answering counterfactual questions is notoriously difficult task, because they naturally refer to unobservable data [27). A principled approach to such questions is adopt the framework of Structural Equation Models (28 27. For example, could assume that the data generating process is given as in Figure 2 If this model is correct, counterfactual inference
105
107
109
110
in the Pearl framework 277 can then be performed by: 1. abduction: inferring the latent from and e using p(z) x, c), 2. action: swap e for c', 3. prediction: use p(x|2, c') to obtain a predictive
111
112
distribution for the counterfactual. Thus, the counterfactual distribution of x, observed with condition
113
e but predicted for condition c' is given by
p px Xx, = c))(x,,))((x dz.
(I)
In order to make use of this relationship. must fit latent variable model (29 such CVAE that 115 will estimate the encoding distribution p(z)x;,c) G) and the generative distribution p(x|2,c').
116 3 Unifying counterfactual inference, data integration and fairness
117 We have seen that batch effect correction, data integration and counterfactual inference are central problems of interest for the application of latent variable models in computational biology.
3

---

119
For counterfactual inference, latent variable models such as the CVAE are increasingly popular choices (2]] The failure mode in which different values of e form separate latent clusters, however, can be catastrophic for this application. When this happens, simply switching to c' is not correct, we have to account for the shift between clusters 9 Mathematically, the latent space clustering phenomenon violates the assumption z II. c that is required by the model in Figure 2 Thus, whilst it
120
122
123
124
is not always possible to know when we have found the right causal model [30], we can immediately say that a model in which z and are dependent is not correct.
126
Another key challenge for computational biology is data integration. Suppose our data (x1, ci).... (Xa.C) in which c; indicates the experimental batch, exhibits batch affects-thes are changes in the observation x; due to the experimental conditions rather than true changes in the underlying biology. One approach to dataset integration is to create a representation z z(x, c) that 'subtracts' the batch effects. Downstream tasks can then work with z in place of without learning signal based on misleading batch effects. To know when we have successfully subtracted batch effects, we might assume that there are no population-leve differences between batches. In other words, the marginal distribution of z should be the same for each value of the condition
127
128
129
130
132
134
Thirdly, this same notion of building representation that cannot be used recover has been studied widely in recent literature on fairness [1 32 33]. In particular, if we wish to make a predictive rule based on x that does not diseriminate between individuals in different conditions c, we can use a fair representation z, one which cannot be used to recover c, as an intermediate feature and train our model using z. Such a representation clearly needs to contain information from x, but without containing any information that could be used to recover c.
135
136
137
138
139
140
To connect these three notions of 'alignment in representation space" we recall the key components of the CVAE-the encoder go (z)xx, e) and decoder pe (x z, c)-and we now drop the 0, ø subscripts for conciseness. The marginal distribution of representations within condition c C is q(z|c)
141
142
143
Ep(xle) [x(z)x, c)), and the marginal distribution of over all conditions not equal to c is denoted
a(z) = (oeec.44o(c))((2
(2)
Exec,e p(ce)
144
The following Theorem brings together key notions in counterfactual inference, data integration and fair representation learning. See Appendix [ for the proof.
145
148
Theorem 1. The following are equivalent: 1) z TL under distribution q. 2) every c, E C,
147
(z(c) = qzzl'', 3) for every c e C. q(zic) = q() 4) the mutual information 1(z, c) under
148
distribution q. 5) z cannot predict better than random guessing.
4
Contrastive Mixture of Posteriors
149
150
We have seen that counterfactual inference, data integration and fair representation learning can be understood through the unified concept of learning a representation such that the latent variable z is independent of the condition e under the distribution q. so that the latent clusters with different values of are perfectly aligned. Building off the CVAE, which rarely achieves this in practice, a number of authors have attempted to use a penalty term to reduce the dependence of upon e during training.
151
152
153
154
155
The most successful methods, such as VAE [10, are based on a Maximum Mean Discrepancy (MMD) I7]. We discuss this and other common methods in Section 5 Whilst VAE and related methods can work well, they require an MMD kernel, not part of the original model, to be specified and its parameters to be carefully tuned. Experimentally, we observe that MMD-based methods can often struggle when there is complex global structure in the latent space. We also analyse the gradients of MMD penalties, showing that they have several undesirable properties.
156
158
159
160
181 We propose a novel method to ensure the conditions of Theorem do hold in CVAE model. Our penalty is based on posterior distributions obtained from the model encoder itself. That is, we do not introduce any external discrepancy measure, rather we propose a penalty term that arises naturally from the model itself. Taking our inspiration from contrastive learning 20 and the VaMP prior we suggest novel penalty to enforce equation condition 3) of Theorem I This equation requires the equality of the marginal distribution q(z(c) and q(2]-c) for each c E C. In practice,
162
163
164
185
26
186
167
these marginal distributions can be approximated by finite mixtures. To encourage greater overlap between (z(cc) and q(z) -c), we can encourage points with the condition e to be in areas of high density under the representation distribution for other conditions, i.c. areas in which g(z) -c) is also
188

---

170
high. To encourage this, we can add the penalty term Po(z,, c:) = log cc) to the objective for the data pair G)). When we minimise Po. this brings the representations of samples under condition towards regions of high density under g)] Since the density g(z|-c) c) is not known in closed form, we approximate ((z|=e) using other points in the same training batch as (x, c²). Indeed, suppose we have a batch (x1, .... (xB, CB). We let To denote the subset of indices for which Cj = c and L-c denote its complement. We use the approximation
171
172
173
174
175
log logq(z|1-c,) log
1
Σ jET-
((112),c))
178
and we will show in Theorem 2 this approximation in fact leads a valid stochastic bound.
177
may happen that the penalty Po causes points to become too tightly clustered. Indeed, the penalty encourages latent variables gravitate towards high density regions of q (cz) Inspired by contrastive learning, we include a second term which promotes higher entropy of the marginal, thereby avoiding tight clusters of points. Combined with Po. this leads
179
second penalty Pi (24,cz = log ((,](c)) - log (-) Again, the density ((2(c) is not known in closed form, approximate it using points within the same training batch in similar fashion to (3)
181
182
but
189
Combining both approximations and taking the mean over the batch gives our Contrastive Mixture of
184
Posteriors (CoMP) misalignment penalty
B
CoMP penalty
log
Σ (())) JeLe.
1 loog [T=.]]
Σ 5EI
a(z; (()
(4)
B
[Ta
1=1
185
where XIB, CI.B. 1.B
I1 B p(x,, ())(2/));/))))))) x c.) is random training batch of size B, Io denotes
the subset of (I...., B} with condition e and I-c = (l...., B} \ le. Our method therefore utilises training penalty for CVAE-type models that encourages the conditions of Theorem I to hold by
186
188
using mixtures of the variational posteriors themselves to approximate q(zle) and q(2| c). We do
not introduce an additional kernel or hyperparameter-heay discrepancy measures.
190
As hinted at by the definition of Pi, CoMP can be seen as approximating symmetrised KL- divergence between the distributions g c) and (z)]-c). In fact, the following theorem shows that the CoMP misalignment penalty is stochastic upper bound on a weighted sum of KL-divergences.
192
Theorem 2. The CoMP misalignment penalty satisfies
B
E THe ((x
Σ
log
Σ ((2x(x))0c) - log [L-o jela
Σ (12)11):.i)) jEI-c.
B
[Ll
cEC
194
and the bound becomes tight as xc.
The proof is presented in Appendix B Our result reveals that our penalty directly enforces
condition 3) of Theorem by reducing the KL divergence between each pair q(z)c) q(z]-c) weighted by p(c). As with standard contrastive learning, add the CoMP misalignment penalty to the familiar 3-VAE objective to give our complete training objective for batch of size B
196
method benefits from larger batch sizes. We
198
C
CoMP
1
Σ
18
log logp(x./2;.cc ci) + log p(z.) (Z, X;,
log
[T-o] Ejen ((, (x),) ) 1
EjeL (221))3)))) (i) T.T
(5)
B
with one new hyperparameter y that controls the strength of the regularisation apply enforce
the requirements IL Theorem 2 shows that, if C is the standard 3-VAE objective, then maximising E
[CS [CMMP]
<
C
Eesc p(c) KL [c))
4.1
Analysing CoMP gradients
204 Before presenting empirical results on the performance of CoMP, we attempt to understand how differs from existing penalties in the literature. Specifically, compare CoMP with Gaussian
5

---

206
posterior family with MMD using a Radial Basis Kernel [4]. In Appendix C we show that both
207
methods can be interpreted as applying a penalty to each element 2; C. of the training batch. We show further that, under certain conditions, the gradient of the MMD penalty for zi, cs takes the form
208
Vz, x
PMiM ci)
2
1121 2,, (2) zi) [T Teg
4
12 =,1" (2) 2;), (6)
jEle;
jEL
209
whilst the CoMP penalty gradient takes the form
2 2IjEr, e --llz--
Hoy (Mz, zi)
2 2EjEL-., e -lei-uayl 11² (Hee, - 2i)
Vz, PCoMp(Z: Gi =
(7)
B BEJET, e 12,-11,11
B BEj -lla--H,, May 112
where Uz, is the variational mean for Zj. One important feature of the MMD gradients is that, 210
211
lz, - z, /1 is large for all j * i, for instance when the point z, is part of an isolated cluster, then the gradient to update the representation z, will be small. So if z; is already very isolated from the
212
213
distribution (2)]-c)). then the gradients bringing it closer to points with condition - G. will be small, In comparison the MMD gradient, it can be seen that gradients for CoMP are self-normalixed. This means that the gradient through 2; will be large, even when Z, is very far away from any points with condition C₁. This, in turn, suggests that that CoMP is likely to be preferable to MMD when we have a number of isolated clusters or interesting global structure in latent space, something which often occurs with biological data. The CoMP approach also bears relationship with wearest-neighbour
214
215
216
217
218
219
approaches [35]. Indeed, for a Gaussian posterior as a - O, the - term of the gradient places all its weight on the nearest element of the batch under condition C..
220
5
Related Work
222
The problem of batch correction in data integration has been addressed using linear 12 13 and nonlinear methods D14 15] that perform transformations of the original feature space. In both cases, the goal is to transform the feature space so that information related to the scientific question of interest is retained while dependence on the batch (or nuisance covariate) is reduced. Methods based on representation learning learn a low-dimensional representation, z = q(x), which independent of nuisance factors while also being faithful representation of the original data [18] 5 36 10 37).
224
226
228
Of these, the work that is most similar to ours are the VFAE [18]. in which the authors introduce an MMD 17 penalty to encourage the marginal distributions of under different values of e to be close, and the tr VAE [0]. where the MMD penalty is applied to the output of the first layer of the decoder, rather than to directly. Representation learning algorithms for counterfactual inference have been shown to benefit from a penalty enforcing distributional similarity between the representations of the
230
232
233
treated and untreated samples [12]. Elsewhere, authors have applied the variational autoencoder to
234
inference on causal graphs 38
39 40
Experiments
235
236
We perform experiments on three datasets; 1) Tumour / Cell Line: bulk expression profiles of tumours and cancer cell-lines across 39 different cancer types; 2) Single-cell PBMCs: single-cell gene expression (scRNA-seq) profiles of interferon (IFN)-8 stimulated and untreated peripheral blood mononuclear cells (PBMCs) [I]; 3) UCI Adult Income: personal information relating to education, marriage status, ethnicity, self-reported gender of census participants and a binary high / low income label ($50,000 threshold) (2). All experiments used a 90/10 training/validation split.
237
238
239
241
242
The two broad objectives across our experiments are 1) to demonstrate the extent to which the two random variables z; and o are independent, and 2) to quantify useful information retained in z, To benchmark CoMP on the first objective, we use the following pair of k nearest-neighhor metrics: the metric used to evaluate batch correction methods in biology, and local Silhouette 4 Sk.c. In both cases a low value close to zero would indicate good local mixing of sample representations. As for the second objective, if we assume the existence of an additional discrete label d; that represents information one wishes to preserve in the Tumour / Cell Line case, de is the cancer type, while for the PBMC experiment, it refers to cell type - then we calculate
243
244
245
kBET 42,.
246
Coefficient
247
248
249
kBET and separately for every fixed-d, subpopulation and take the mean. We refer to these as the mean Silhouette Coefficient &y.c and the mean kBET metric m-kBET respectively. Full details of the datasets and metrics are given in Appendix D
251

---

Table 1: Tumour / Cell Line experiment results, with k 100.
Cell Line, and 0.01. Sk,c
and Sk,c are the two Silhouette Coefficient variants (see Section 6
The top scores are in bold.
Accuracy
8k,c
kBETk.
$R.0
m-kBET k.o
VAE
0.209
0.658
0.974
0.803
0.581
CVAE
0.328
0.554
0.931
0.684
0.571
VFAE
0.585
0.168
0.258
0.198
0.188
VAE
0.585
0.096
0.163
0.138
0.123
Celligner
0.578
0.082
0.525
0.568
0,226
CoMP (ours)
0.579
0.023
0.160
0.094
0.101
VAE
CVAE
VFAE
DnVAE
Celligner
CoMP
Figure 3: 2D UMAP projection of posterior means of z, from Tumour / Cell Line data. Tumours (blue) and cell lines (orange).
253
6.1
Alignment of tumour and cell-line samples
254
Despite their widespread use pre-clinical cancer studies, cancer cell-lines are known to have significantly different gene expression profiles compared to their corresponding tumour samples. Here we evaluate the ability CoMP to factorise out the tumour / cell line condition from its latent representations. This can be seen as both a dataset integration and batch effect correction task. In addition to the set of nearest neighbor- gaggboo-bbsedd mixing evaluations, we train Random Forest model on the representations of the tumour samples and their cancer-type labels and assess the prediction accuracy on held-out cell lines. To match the results from (11, the evaluations are performed on the 2D UMAP projections, The results are presented in Table 1
255
256
257
258
260
262
As expected, both the VAE and CVAE baselines fail the mixing task; the three explicitly penalised VAE models and, to a lesser extent, the Cellinger method have good mixing performances, with CoMP outperforming the benchmark models by significant margin on the silhouette coefficient and kBET metric, while successfully maintaining high accuracy in the cancer-type prediction task. We also see from Figure 3 that CoMP representations have the fewest instances of isolated tumour-oni only clusters. Finally, from our evaluation on the & and m-kBET metrics, we can deduce that the occurrence of cell lines one cancer type erroneously clustering around tumours of a different
263
264
285
266
267
266
type is less frequent for CoMP compared to the other models. In Appendix D we qualitatively validate this for several example clusters.
270
6.2 Interventions
Obtaining molecular measurements from biological tissues typically requires destructive sampling. For example, to obtain ScRNA-seq data, each cell is lysed so that the RNA molecules contained within it can be extracted and sequenced. This process destroys each cell, meaning that we are unable to study the gene expression profile of the same cell over time or under multiple experimental conditions. As we discussed in Section 2.2 counterfactual inference can be used to predict how the molecular status of destroyed biological sample would have differed it were measured under different experimental conditions, such as applications of different drugs.
273
275
277
279
To assess CoMP's s utility in counterfactual inference, we trained it on ScRNA- scRNA-seq data from PBMCs that were either stimulated with IFN-9 or left untreated (control) [41]. It is clear from Figure 4 that IFN-3 stimulation causes clear shifts in the latent space between stimulated and control cells from the
281
same cell type. Noticeably, the CD14 and CD16 monocyte and dendritic cell (DC) populations see greater shifts in their gene expression after stimulation. CVAE fails to align these particular cell types
283
in the latent space, while VAE, VFAE and CoMP perform better. However, stimulated and control

---

VAE
CVAE
trVAE
VFAE
CoMP
CD6 T
CD14 Mono
DC
stimuisted
Figure 4: 2D UMAP projections of posterior means of z; derived from stimulated and control PBMC scRNA-seq data. Top row: colours indicate immune cell types, bottom row: colours indicate condition (IFN-B stimulation or control).
CD14 Mono
DC
NK
Difference in expression
10
05
0.0
00
0050
up
-2
-00
-0.5 75
down
-100
-20
-125
VAE
OVAE
trVAE
VFAE
CoMP
VAE
CVAE
trVE
VFAE
CoMP
VAE
CVAE
trVAE
VFAE
CoMP
Figure 5: The difference in gene expression values for the top 50 differentially expressed genes
(up-regulated: red, down-regulated: blue) between IFN-8 stimulated cells and counterfactually stimulated control cells for CD14 monocytes, dendritic cells (DC) and natural killer (NK) cells. See Appendix D for further details.
285
cells are better mixed in the latent space derived from CoMP than those from the other models (see metrics presented in Appendix D)
287
Next we perform a counterfactual prediction task under a IFN-3 control-to-simmulation variable swap, i.c. the gene expression profiles for control cells were reconstructed through the decoder with the condition, c + stimulated. This means we utilise equation ① with our encoder go (z)x,, e) and decoder pe xx z e') in place of p(z)x, x, c) and p(x|z, c'). The degree to which the models respect the requirement z I1 will influence the quality of predictions. Figure 5 shows how the profiles of (actual) stimulated cells differ from the counterfactual predictions for selection of cell types (see Appendix D for the complete set of results). We see that baseline models tend to systematically underestimate the expression of genes up-regulated by stimulation and overestimate those down-
288
289
290
291
292
293
295
regulated. CoMP outperforms all other models by accurately predicting the expression alterations brought about by stimulation.
297
6.3
Fair Classification
The goal for this fair classification task is to learn representation on the Adult Income dataset that is not predictive of an individual's gender whilst still being predictive of their income. We compute a baseline by predicting gender and income labels directly from the input data and compare our method
299
301
to the published results for the VFAE [18] and the VAE. We also include results for a standard sampled from the posterior before classification, our experiments used the posterior means to avoid the noise from sampling acting to
VAE and CVAE. Unlike in [88, where the representations
303
304
mask the inclusion of predictive information about gender in the encodings.
305
CoMP achieves gender accuracy that close to random (67.5%), tying with the VFAE results from whilst also remaining competitive with the other methods on income accuracy (Table CoMP
18
8

---

Table 2: UCI Adult Income experiment results with k 1000, c Male for Sk.c. and & 100, a = 0.01 for kBET lower gender prediction accuracy is better; 0.675 is the lowest achievable.
Gender Acc.
Income Acc.
8k.c
kBETE.o
Original data
0.796
0.849
0.067
0.786
VAE CVAE VFAE (sampled) 18
0.764
0.812
0.054
0.748
0.778
0,819
0.054
0.724
0.680
0.815
VFAE (mean)
0.789
0,805
0,046
0.571
trVAE
0.698
0.808
0.066
0.731
CoMP (ours)
0.679
0.805
0.011
0.451
Original
trVaE
(mean)
CoMP
Figure 6: UMAP projections for the UCI Adult Income dataset, coloured by gender. Showing the original data and latents for tAV,E, VFAE and CoMP. Male (blue) and female (yellow).
also outperforms all methods on the nearest neighbour and silhouette metrics (Table Latent space mixing between males and females can be seen qualitatively in the 2D UMAP projection (Figure 66)
308
Conclusion
309
310
We presented Contrastive Mixture of Posteriors (CoMP) an effective means to perform batch correction, data integration, counterfactual inference and fair representation learning CVAE framework. Whilst CoMP covers the majority common use-cases for these tasks, there are several limitations that are avenues of future research. For example, in scRNA-seq analysis, there often the need to integrate more than two datasets together, or to adjust for continuous condition variables. Mathematically, CoMP is applicable to any number of discrete conditions, and it would be CoMP could tackle the case of continuous condition variable. Additionally, CoMP requires the condition variable to be fully observed: future work might attempt to generalise to the partially observed case.
Limitations
312
313
314
316
interesting to apply to setting with conditions. Extensions
318
319 Summary We identified marginal independence between the representation z and condition e as the mathematical thread linking data integration, counterfactual inference and faimess. We proposed CoMP, a novel method to enforce this independence requirement in practice. We saw that CoMP has several attractive theoretic properties. First. CoMP only uses the variational posteriors, requiring no additional discrepancy measures such as MMD. Second, we proved that the CoMP penalty can be interpreted as an upper-bound on weighted sum of KL divergences, connecting
321
322
323
324
well- founded divergence measure. Third, we demonstrated that, unlike MMD, CoMP gradients have self- normanormmliiing property, allowing one to obtain strong gradients for distant points in latent space with complex global structure. Empirically, we demonstrated CoMP's performance when applied to two biological and one fair representation learning dataset. These biological datasets are of critical importance in drug discovery, for example matching cell-lines to tumours for effective pre-clinical assay development of anti-cancer compounds. Overall, CoMP has the best in class performance on
325
326
327
328
330
all tasks across range of metrics that measure either latent space mixing or faimess.

---

References
332
333
[1] Valentine Svensson, Roser Vento-Tormo, and Sarah A Teichmann. Exponential scaling of single-cell ma-seq in the past decade. Nature protocols, 13(4):599-604, 2018.
334
335
[2] Aviv Regev, Sarah A Teichmann, Eric S Lander, Ido Amit, Christophe Benoist, Ewan Bimey, Bernd Bodenmiller, Peter Campbell, Piero Carninci, Menna Clatworthy, et al. Science forum: the human cell atlas. Elife, 6:e27041, 2017.
336
337
338
[3] Diederik Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv: 2013.
339
Danilo Jimenez Rezende, S. Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. ICML, 2014.
341
342
[5]
Romain Lopez. Jeffrey Regier, Michael B Cole. Michael I Jordan, and Nir Yosef. Deep generative modeling for single-cell transcriptomies. Nature methods, 15( (51)))10033 1058, 2018.
344
[6]
Gregory P Way and Casey S Greene. Extracting a biologically relevant latent space from cancer transcrip- tomes with variational autoencoders, In PACIFIC SYMPOSIUM ON BIOCOMPUTING 2018: Proceedings of the Pacific Symposium, pages 80-91, World Scientific, 2018.
345
347
[7]
Dongfang Wang and Jin Gu. Vase: dimension reduction and visualization of single-cell ma-seq data by deep variational autocncoder. Genomics, proteomics & bioinformatics, 16(5):320-331, 2018.
348
349
Christopher Heje Grønbech, Maximillian Fornitz Vording. Pascal Timshel, Casper Kaae Sønderby, Tune H Pers, and Ole Winther. sevae: Variational auto-encoders for single-cell gene expression data. Bioinformatics, 36(16) 4544442244 2020.
350
351
352
Mohammad Lotfollahi, F Alexander Wolf, and Fabian Theis. scGen predicts single-cell perturhation responses. Nature methods, 16(8):155, 2019.
353
354
[10] Mohammad Lotfollahi, Mohsen Naghipourfar, Fabian J Theis, and F Alexander Wolf. Conditional out-of-sample generation for unpaired data using tAE. arXiv preprint 19011:9100.1199,, 2019.
356
[1I] Matthew Amodio, D. V. Dijk, R. Montgomery, Guy Wolf, and Smita Krishnaswamy. Out-of-sample extrapolation with neuron editing. arXiv: Quantitative Methods, 2018.
357
358
[12 Johnson. Cheng Li, and Ariel Rabinovic. Adjusting batch effects in microarray expression data using empirical Bayes methods. Biostatistics, 8 1:118-27, 2007.
359
[13] J. Leek and John D. Storey. Capturing heterogeneity in gene expression studies by surrogate variable analysis. PLoS Genetics, 3, 2007.
381
382
[14] Laleh Haghverdi. A. Lun, Michael D. Morgan, and J. Marioni. Batch effects in single-cell ma- maquennnn data are corrected by matching mutual nearest neighbors. Nature Biotechnology. 36:421-427, 2018.
364
[15] Allison Warren, Andrew Jones, Tsukasa Shibue, William C Hahn, Jesse S Boehm. Francisca Vazquez. Aviad Tsherniak, and James M McFarland. Global computational alignment of tumor and cell line transcriptional profiles. Nature Communications, 12(22), 2021.
385
367
[16] Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. In Advances in neural information processing systems, pages 3483-3491, 2015.
368
369
370
[17] Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bemhard Schölkopf, and Alexander Smola. kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
371
372
[18] Christos Louizos, Kevin Swersky, Yujia Li, Max Welling, and Richard Zemel. The variational fair autoencoder. arXiv preprint arXiv: 1511::1111100030 100830, 2015.
373
374
[19] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint a7Xiv:1807.03748, 2018.
375
376
[20] Ting Chen, Simon Kombliihh Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. arXiv preprint arXiv:2002.05700 05709, 2020.
377
378
[21] Allison Warren, Yejia Chen, Andrew Jones, Tsukasa Shibue, William C Hahn, Jesse S Boehm, Francisca Vazquez, Aviad Tsherniak, and James M McFarland. Global computational alignment of tumor and cell line transcriptional profiles. Nature Communications, 12(1):1-12, 2021.
379
380
381
[22] Dheeru Dua and Casey Graff. UCI machine learning repository. 2017.
382
[23] Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematicul statistics, pages 400-407, 1951.
384
[24] Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. International Conference Learning Representations, 2017.
385
10

---

387
[25]
Emile Mathieu, Tom Rainforth, N. Siddharth, and Yee Whye Teh. Disentangling disentanglement in variational autoencoders. In In International Conference on Machine Leurning, pages 4402-4412. PMLR,
309
2019.
380
[66] Jakub Tomczak and Max Welling, Vae with a vampprior. In International Conference on Artificial Intelligence and Statistics, pages 1223- PMLR, 2018.
382
[27]
Judea Pearl. Causality. Cambridge university press, 2009.
[28] Kenneth A. Bollen. Structural equation models. Wiley, 2005.
394
[29] Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In International conference on machine learning, pages 3020-3029, 2016.
385
396
[30] Jonas Peters, Joris Mooij, Dominik Janzing, and Bemhard Schölkopf. Identifiability of causal graphs using functional models, arXiv preprint urXiv:1202.37577 2012.
397
398
[31] Rich Zemel, Yu Wu, Kevin Swersky, Toni Pitassi, and Cynthia Dwork. Learning fair representations. In International conference on machine learning, pages 325-333. PMLR, 2013.
399
400
[32] Matt J. Kusner, Joshua R. Loftus, Chris Russell, and Ricardo Silva. Counterfactual fairness. arXiv preprint
401
orXiv:170300683 2017.
402
[33]
Craig Glastonbury, Michael Ferlaino, Christoffer Nellaker, and Cécilia M Lindgren. Adjusting for confounding in unsupervised latent representations of images, arXiv preprint 1Xiv:1811.06999 2018.
403
[34]
Jean- Phanlppilip Vert, Koji Tsuda, and Bernhard Schölkopf. A primer on kernel methods. Kernel methods in computational biology, 47:35-70, 2004.
405
406
[35]
Ke Li and Jitendra Malik. Implicit maximum likelihood cstimation. arXiv preprint arXiv: arXiv:1809.090 1809.09077
407
2018.
408
[36]
Romain Lopez, Jeffrey Regier, Michael Jordan, and Nir Yosef. Information constraints on auto-encoding variational bayes. In Advances in Neural Information Processing Syxtems, pages 61 14-6125, 2018.
409
[37] Kaspar Märtens and Christopher Yau. Neural decomposition: Functional anova with variational autoen- coders. In International Conference on Artificial Intelligence and Statistics, pages 2917-2927. PMLR. 2020.
411
412
413
[38] Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling, Causal effect inference with deep latent-variable models. In Advances in Neural Information Processing Systems, pages 6446-6456, 2017.
414
415
416
[39] H. Kim, Seungjue Shin, Joonho Jang, Kyungwoo Song. Weonyoung Joo, Wanmo Kang, and II-Chul Moon. Counterfactual fairness with disentangled causal effect variational autoencoder. In AAAI, 2021.
418
[40] Stephea R. Pfohl, Tony Duan, Daisy Yi Ding, and Nigam H. Shah. Counterfactual reasoning for fair clinical risk prediction. In Finale Doshi- Jim Fackler. Ken Jung. David Kale, Rajesh Ranganath, Byron Wallace, and Jenna Wiens, editors, Proceedings of the 4th Machine Learning for Healthcare Conference, volume 106 Proceedings of Machine Learning Research, pages 325-358, Ann Arbor, Michigan, 09-10 10 Aug 2019. PMLR.
419
421
422
[41] Hyun Min Kang, Meena Subramaniam, Sasha Targ, Michelle Nguyen, Lenka Maliskova, Elizabeth McCarthy, Funice Wan, Simon Wong, Lauren Bymnes, Cristina M Lanata, et al. Multiplexed droplet single-cell ma-sequencing using natural genetic variation. Nature biotechnology, 36(1):89, 2018
425
426
[42] Maren Büttner. Zhichao Miao, F Alexander Wolf, Sarah A Teichmann, and Fabian J Theis. A test metric for assessing single-cell ma-seq batch correction. Nature methods, 16(1):43-49, 2019.
427
428
[43]
Peter J Rousseeuw. Silhouettes: a graphical aid to the interpretation and validation of cluster analysis. Journal of computational and applied mathematics, 20:53 65, 1987.
429
430
[44] Christian P. Robert and Judith Rousseau. How Principled and Practical Are Penalised Complexity Priors? Statistical Science, 32()3336-40, February 2017.
431
[45] Murray Aitkin. Posterior Bayes Factors. Journal of the Royal Statistical Society. Series B (Methodological), 53(1):111-14 1991.
433
434
[46] Paul Fearnhead and Dennis Prangle. Constructing summary statistics for approximate Bayesian computa-
435
tion: semi-automatic approximate Bayesian computation. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 74(3)411-477, 2012.
436
437
[47] Adam Foster, Martin Jankowiak, Mathew O' O'Meara, Yee Whye Teh, and Tom Rainforth. A unified stochastic gradient approach to designing bayesian-optimal experiments. In International Conference on Artificial Intelligence and Statistics, pages 2959-2969. PMLR, 2020.
438
440
[48] J. Weinstein, E. Collisson, G. Mills, K. Shaw, B. Ozenberger, Kyle Ellrott, Shmulevich, C. Sander, and Joshua M. Stuart. The Cancer Genome Atlas Pan- Cancer analysis project. Nature Genetics, 45:1113- 1120. 2013.
441
11

---

443
[49] DS Gerhard, S Hunger, C Lau, Maris, P Meltzer, S Meshinchi, E Perlman, J Zhang, J Guidry-Auvil, and M Smith. Therapeutically applicable research generate effective treatments (target) project: Half of pediatric cancers have their own" driver" genes. In PEDIATRIC BLOOD CANCER, volume 65, pages S45-S44 S45. WILEY 111 RIVER ST, HOBOKEN 07030-5774, NJ USA, 2018.
445
446
[50]
M. Ghandi, F. Huang. J. Jané Valbuena, G. Kryukov, Christopher Lo, E. McDonald, J. Barretina, E. Gelfand, C. Bielski, Haoxin Li, Kevin Hu, Alexander Y. Andreev-Drakhiin J. Kim, Hess. B. Haas, F. Aguet, Paolella, M. Lawrence, Rehan Akbani, Y. Lu, Hong L.. Tiv, P. Gokhale, Antoine de Week, Ali Amin Mansour, C. Oh, J. Shih, Kevin Hadi, Yanay Rosen, J. Bistline, K. Venkatesan, Anupama Reddy, Dmitriy Sonkin, Manway Liu, J. Lehár, J. Korn D. Porter, M. Jones, J. Golji, G. Capon- igro, Jordan E. Taylor, C. Dunning, Amanda L Creech, Allison Warren, James M. McFarland. Mahdi Zananighomi. A. Kauffmann, Nicolas Stransky. M. Imieliáski. Y. Maruvka, A. Cherniack, Aviad Tsher- niak, F. Vazquez, J. Jaffe, A. A. Lane, D. Weinstock, C. Johannessen, Michael P. Morrissey, F. Stegmeier, R. Schlegel, W. Hahn, G. Getz, G. Mills, J. Bochm, T. Golub, L. Garraway, and W. Sellers. Next-generation characterization of the Cancer Cell Line Encyclopedia. Nature, 569:503-5088 2019.
448
449
B. Weir, M. Rothberg,
450
453
454
455
Checklist
457
458
1. For all authors...
459
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
480
(b) Did you describe the limitations of your work? [Yes] See Section 7
462
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
463
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
464
465
2. If you are including theoretical results...
466
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
487
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix B
468
3. If you ran experiments..
469
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as URL)? [Yes] See Appendix
470
471
D
472
(b) Did you specify all the training details (c.g., data splits, hyperparameters, how they were chosen)? [Yes] See Appendix D
474
(c) Did you report error bars (e.g., with respect the random seed after running experi-
475
ments multiple times)? [Yes] See Appendix D
476
(d) Did you include the total amount of compute and the type of resources used (e.g., type
477
of GPUs, internal cluster, or cloud provider)? [Yes] See Appendix [
478
4. If you are using existing assets (c.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
480
(b) Did you mention the license of the assets? [Yes]
481
(c) Did you include any new assets either in the supplemental material or as URL? [Yes]
482
See Appendix [
(d) Did you discuss whether and how consent was obtained from people whose data you 're
using/curating? [N/A]
484
485
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
486
487
5. If you used crowdsourcing conducted research with human subjects...
488
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
489
490
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
482
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
483
12


Reviews:
[{"review_id": "_mLuoOqZaA0", "reviewer": "Reviewer_UbUq", "summary": "This paper considers the problem of data integration, counter factual inference and fairness concentrating on large data sets where each data point lives in high dimensions. The authors generate an embedding of the data into a latent space in such a way that the distribution of the embeddings is the same for data from all clusters. The authors propose to learn such an embedding through a variation autoencoder (VAE). They come up with a novel penalty called contrastive mixture of posteriors misalignment (CoMP) penalty which they use to with the ELBO loss to train the VAE. They show empirical results chiefly on single-cell RNA-seq data and also on one income data-sets. ", "questions": "", "limitations": "", "rating": 7, "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": "The authors consider the problems of data integration, counter factual inference and fairness. They primarily concentrate on single cell RNA-seq for the first two points - a new technology in which one can obtain gene expression on a cell by cell basis  giving us data-sets with millions of cells - and tens of thousand measurements of gene expression for each cell. However different experiments carried out by different scientists unfortunately have batch effects. Being able to combine data from experiments run different scientists to do analysis is necessary for scientific integration - making data integration and dealing with batch effects important problems to consider. The counter-factual inference in single cell datasets is less important as there usually are technical and biological replicates available for experiments. There of course is a plethora of literature on fairness. I did think that the connection between the three was nice, even though it is obvious. \n\nThe authors consider the method of embedding data from different batches into the same latent space as a means of getting rid of batch effects. However, this approach is explicitly not suitable for data integration in single cell RNA-seq because the data generated by different scientists is usually not from the exact same tissue from the same individual, and thus one would expect  the distribution of the latent space to be different in different batches even if there was no batch effect (without additioinal information the problem is not even well defined - but that's what makes working with single cell RNA-seq hard and interesting). \n\nThe authors consider a VAE as means of obtaining such an embedding. They do this by coming up with a novel penalty to the ELBO loss. The penalty penalises differences in the distribution of points in the latent space for different batches - and is different from existing VAE based penalties in that it uses the posterior of the model learnt in the penalty rather than a different kernel.\n\nThe authors then provide empirical results on Single Cell RNA seq and an income data. One concern I have is that the authors limit their comparisons to VAE based methods for Single Cell RNA-seq when they are not the most popular ways in literature to deal with batch effects and counter-factual inference not comparing with methods like those of [14], and many other papers exploring that problem such as [100], [101], and [102]. Similarly, there is no comparison for non-VAE based methods (for example [103]) for fairness claims. \n\nThe paper is well written and I liked the broad direction of the paper. But the formulation does not match the application of single cell RNA-seq that the paper purports to target. Furthermore, empirical results are found wanting comparing it with non VAE based approaches; which makes its contribution from the application perspective quite small. The only novelty in the paper seems to be using a different penalty in training the VAE - and I dont know the VAE literature enough to know who important or novel a contribution that is.\n\nTo me this is a borderline paper leaning towards rejection from an applications perspective. I like the idea and it is well written, but the formulation does not target any problem which analysts would be interested in and empirical results are not thorough.\n\n\n[100] Korsunsky I, Millard N, Fan J, Slowikowski K, Zhang F, Wei K, Baglaenko Y, Brenner M, Loh P-r, Raychaudhuri S. Fast, sensitive and accurate integration of single-cell data with Harmony. Nature Methods\n\n[101] Welch J, Kozareva V, Ferreira A, Vanderburg C, Martin C, Macosko E. Integrative inference of brain cell similarities and differences from single-cell genomics. bioRxiv. 2018:459891 Available from: http://biorxiv.org/content/early/2018/11/02/459891.abstract\n\n[102]Stuart T, Butler A, Hoffman P, Hafemeister C, Papalexi E, Mauck WM 3rd, et al. Comprehensive integration of single-cell data. Cell. 2019;177:1888–1902.e21.\n\n[103]Agarwal, Alekh, Alina Beygelzimer, Miroslav Dudík, John Langford, and Hanna Wallach. \"A reductions approach to fair classification.\" In International Conference on Machine Learning, pp. 60-69. PMLR, 2018.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "listed above", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4.5"}, "extra_text": {}}, {"review_id": "NnWTjXsaaDB", "reviewer": "Reviewer_SVjF", "summary": "The authors propose a penalty term called Contrastive Mixture of Posteriors (CoMP) for conditional variational autoencoders to learn latent representations invariant of a condition. Their penalty term doesn’t have extraneous requirements and is free from requiring a kernel (as is the case case with MMD) and enforces independence in latent space with respect to a condition (e.g. batch effect, sensitive information etc.). CoMP penalty itself is made up of two components which 1) encourage latent representations to be clustered together in both conditions - under a condition and under complement of said condition, 2) avoid tight clustering using a contrastive learning approach.\n \nThe work is useful in some computational biology tasks (batch correction, counterfactual inference) which require low-dimensional representations of high-dimensional sequencing data free from batch effect or other conditions, and also in fair classification to learn a classifier free from sensitive demographic information. For each of these tasks, the authors have presented one experiment and have compared their methods with existing methods based on variational and conditional variance autoencoders.", "questions": "", "limitations": "", "rating": 8, "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": "The paper is very well-written and easy to follow. The applications and tasks considered are well-motivated with Theorem 1. While relatively straightforward, I find the connections drawn between data integration, counterfactual inference and fair classification to be very useful as it brings together references and results from sub-fields that might not be very aware of each other. Finally, the proposed method is very relevant and well-established for both domains considered in the paper - computational biology and fair classification. The limitations mentioned are also well thought out. \n \nMy main remarks (roughly in decreasing order of importance) are:\n \n1) Assuming independence between z and c is absolutely crucial to the foundation of the studied setting, the correctness of the suggested method and to its theoretical analysis. However, the assumption rarely holds with biological data where different samples might, at the same time, present different phenotypes, technical biases (or batch effects), cell type compositions etc. I would be very curious to see results on the integration of datasets where we know the assumption is broken. For instance, if one cell type is present only in one batch, will their latent codes overlap with cells from a different cell type from the other batch? This issue, usually refered to as \"over-correction\" is absolutely crucial in practice. At least in the discussion, I think it could be nice to comment on it and give some insights on how this issue could be tackled both in theory and in practice.\n \n2) The batch correction experiment considered is for alignment of tumor RNAseq (TCGA) and cell-lines (Cellinger). For a proper evaluation of batch correction, further experiments especially on single-cell data, where batch effects are much more prominent and the task is generally harder, should have been considered.\n \n3) The only batch correction method evaluated outside of variational inference based approaches is Cellinger. It would be good to compare the model on single cell data with at least some popular (among single-cell practitioners) batch correction methods in use, e.g. Seurat and Harmony.\n \n4) While understandably, as argued in the paper, the goal is to learn meaningful latent representations free from batch effect, it would still be interesting to evaluate batch correction in the original gene space (Especially considering that MSEs of reconstructions for counterfactual inference experiment are already shown in Tables 9-12 in Appendix D). This would also provide more opportunities for downstream analysis and evaluation of corrected data.\n \n5) Only prostate cancer clusters have been visually shown in detail in Fig. 7 of Appendix D for trVAE and CoMP. It would be good to have detailed visualizations for more cancer clusters, or at least to include other methods for prostate cancer clusters. This is particularly important as in Section 6.1, it is mentioned that Appendix D includes qualitative results for several clusters.\n \n6) In both batch correction and counterfactual inference experiments, the UMAPs of original data are not shown (Only the latent representations in Figures 3 and 4.). While it is clear that original data would be clustered by condition (batch or original and stimulated), it would still be useful to see that.\n \n7) For fair classification, the authors mentioned that the paper shows “published results” from trVAE and VFAE, however in the original trVAE, this experiment was not performed. For completeness, it will be good to include the source where those results were published.\n\n-------------\nAfter checking the author's response, I am upgrading my rating from 7 to 8.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "While I think the most important limitation of this work (the need for independence between z and c) was not mentioned, some other limitations were briefly (but appropriately) discussed in the conclusion.\nThe societal impact were not really discussed but I am not sure how relevant it would be for this specific paper.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "Dg7VzK7hErK", "reviewer": "Reviewer_vhhX", "summary": "The paper proposes a penalty to learn representations that are\nconditionally independent of some observed variables. The idea is to\nuse the contrastive loss that encourage the representations from the\nsame conditional variable to be spread out and those from different\nconditional variable to be close. The paper demonstrates\napplication to fairness, counterfactual inference, and data\nintegration, and demonstrate the effectiveness of the penalty.\n", "questions": "", "limitations": "", "rating": 6, "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": "The paper proposes a penalty to learn representations that are\nconditionally independent of some observed variables. The idea is to\nuse the contrastive loss that encourage the representations from the\nsame conditional variable to be spread out and those from different\nconditional variable to be close. The paper demonstrates\napplication to fairness, counterfactual inference, and data\nintegration, and demonstrate the effectiveness of the penalty.\n\nBoth the theoretical and empirical results of the paper make sense. I\nhave a few major comments.\n\nOne comment is the choice of contrastive loss. It is understandable\nthat using the loss in Eq 3 would encourage tightly clustered\nrepresentations and produce suboptimally informative representations.\nOne thus may think contrastive loss may help avoid trivial solutions\nand preserve informativeness. However, it is unclear why the second\npenalty is sufficient to ``avoid tight clusters of points.'' Why is\nthe contrastive loss necessarily the right choice here? I think the\npaper could benefit from a discussion about it.\n\n\nA second comment is on the comparison between the misalignment penalty\nand the other MMD penalties in the literature. The paper argues that\nMMD suffers when there are many isolated cluster in the data. The\nproblem of MMD appears to be related to scaling; and the misalignment\npenalty has the self-normalized property. I wonder if it implies that\nMMD is equally effective if one manually rescale and recenter these\ndata points to achieve the normalization effect? What if the global\nstructure does not have many isolated clusters? Will the misalignment\npenalty need to pay a price? Are there other differences between the\nmisalignment penalty and the MMD approaches that made the misalignment\npenalty more preferable?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper discussed the limitations in sec 7.", "review_text": "", "extra_scores": {"ethical_concerns": "NA", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10"}, "extra_text": {}}]

Author rebuttal:

Meta-review: [This meta-review was written jointly by the SAC and AC after extensive discussion of the paper.]

This paper has interesting ideas and compelling results, and reviewers were overall supportive of the proposed work. Unfortunately, the paper has two entangled parts, where one of them has quite serious issues. The first part includes a Contrastive Mixture of Posteriors (CoMP) method that leverages a novel misalignment penalty to enforce certain independence relations; the second is related to counterfactual inference, data integration, and fairness.

Regarding the first part, CoMP enforces that a discrete auxiliary variable is independent of the latent variable inside a (conditional)-VAE without resorting to external discrepancies like an MMD. This is achieved by marginalizing out the inference network using empirical samples from the conditional distribution of the observation given settings of the discrete variable. This part of the paper is smooth, and the contributions are pretty substantial. It also connects quite nicely with the idea of variational inference and does not require an extraneous discrepancy measure or a separately tuned kernel, such as in previous approaches. The biggest benefit of this approach is that the scales of the terms in the loss are on the same scale as expectations of log probabilities or bounds of the same quantity. The approach could benefit many places where conditional-VAEs are used such as molecule generation (J. Lim, S. Ryu, J. W. Kim & W. Y. Kim, 2018) or dialog models (T. Zhao, R. Zhao & M. Eskenazi, 2017).

The second part of the paper regarding counterfactual inference, data integration, and fairness is quite concerning, and we feel there are fundamental misunderstandings regarding each one of them. The very claim that the paper introduces a unifying framework for solving these tasks is, from the beginning, completely off.

To understand why this is the case, consider the critical result used in the paper to compute Eq. 1, a quantity known as the "probability of causation," and written as $p(x_{c=c’} | x_i, c_i) = \int p(z|x_i, c_i) p(x | z, c’) dz$. By and large, the understanding of how to evaluate the l.h.s. of this expression is unsound. If one has the true, underlying structural causal model (SCM), say M*, the probability of causation could be evaluated in the manner described in Eq. 1, as introduced in Sec.7.1.1 of Causality (Pearl, 2000). On the other hand, whenever M* is unknown, this is not the case.  In particular, the probability distributions considered in the r.h.s. equation are arbitrary, representing a *learned* SCM M. The true SCM M* and the learned M are not necessarily the same.

The paper then proposes a method such that M will be trained to fit the data and guarantee the independence between Z and X. Still, it can be shown that it’s generically the case that there exists an M that will *almost always* disagree with the actual generating model M*, while agreeing with respect to the data. This means, therefore, that only with observational data, one cannot make statements about counterfactuals in generality. This is a well-known result in the field. For concreteness, please refer to the example discussed in Sec. 1.4.4 (Pearl, 2000) or examples 7-9 in (Bareinboim, Correa, Ibeling, Icard, 2020). This result is called the causal hierarchy theorem and can be found in Thm. 1 of the previous reference or Thm. 3 in https://arxiv.org/abs/2107.08558. Also, this has been shown in the context of autoencoders with continuous variables, for example, Appendix A.4 of https://arxiv.org/pdf/1907.03451.pdf,  has a construction where $c$ and $z$ are independent, but $x$ is constructed in a way that it is also marginally independent of $c$. This means the pair $(z, c)$ are independent and reconstruct $x$ and the pair $(x, c)$ are independent and reconstruct $x$. The pairs are observationally equivalent even with the extra independence constraint between the latent variable and $c$, but can be different under interventions.

Interestingly, the paper has one sentence (around line 124) acknowledging that it is not always possible to recover the correct model, while the rest of the paper seems to ignore this observation.  By and large, the challenge is that for almost any actual model M*, there are many models M that could fit the observed data equally well, but have very different implications in terms of their interventional and counterfactual distributions. This is certainly the most fundamental issue that disallows the paper to claim to be doing “counterfactual inference in generality.”

There are also serious issues regarding the tasks related to fairness and data fusion, and the purported “unification.” For example, many papers here at NeurIPS and other AI venues propose analyzing the fairness problem through causal lenses. The understanding from this work clarifies that enforcing independence among the protected attribute and the outcome variable is orthogonal to the issue at hand; in other, even when the independence is accomplished, biases may be amplified. Here are a few references that look at fairness from the viewpoint of causality (not an exhaustive list):
 
- Counterfactual Fairness, M. Kusner, J. Loftus, C. Russell, & R. Silva, NeurIPS’17.
- A Causal Framework for Discovering and Removing Direct and Indirect Discrimination. L. Zhang, Y. Wu, & X. Wu. IJCAI’ 2017
- Avoiding discrimination through causal reasoning. N. Kilbertus, M. Rojas-Carulla , G. Parascandolo, M. Hardt, & B. Scholkopf. NeurIPS’17.
- Equality of Opportunity in Classification: A Causal Approach. J. Zhang & E. Bareinboim, NeurIPS’18.
- Path-specific counterfactual fairness. S. Chiappa. AAAI’19.

Having said that, we note that it’s possible to stay oblivious to the causal inference literature and claim a purely statistical notion of fairness, which is anchored in conditional independence relations. In its current form, the paper claims to be doing counterfactual/causal inferences, so it’s expected the paper engages with the literature on fairness from the causal perspective, which goes beyond conditional independence relations and purely statistical measures. 

Finally, the data integration problem has also been studied through the causal framework, and I will suggest Ch. 10 of the Book of Why (Pearl & Mackenzie, 2018) or (Bareinboim & Pearl, PNAS’16), which survey the challenges involved in such endeavors. Even though data integration is a different task, most of the issues with the proposed approach come precisely from the misinterpretation of Eq. 1, following the aforementioned discussion.

All in all, we would like to make it clear that the first dimension of the paper, including CoMP, is appreciated, and we believe the contributions are quite significant. On the other hand, the second part suffers from fundamental flaws. Therefore, we don’t know how to accept the paper without re-thinking and re-writing the paper in a major way. One thought would be to demonstrate the value of CoMP on tasks that are more bread-and-butter for conditional-VAEs like generation (see the dialog and molecule generation references above).

Decision: Reject