Contrastive Mixture of Posteriors for Counterfactual Inference, Data Integration and Fairness

Learning meaningful representations of data that can address challenges such as1 batch effect correction , data integration and counterfactual inference is a central2 problem in many domains including computational biology . Adopting a Condi-3 tional VAE framework , we identify the mathematical principle that unites these4 challenges : learning a representation that is marginally independent of a condition5 variable . We therefore propose the Contrastive Mixture of Posteriors ( CoMP ) 6 method that uses a novel misalignment penalty to enforce this independence . This7 penalty is defined in terms of mixtures of the variational posteriors themselves,8 unlike prior work which uses external discrepancy measures such as MMD to9 ensure independence in latent space . We show that CoMP has attractive theoretical10 properties compared to previous approaches , especially when there is complex11 global structure in latent space . We further demonstrate state of the art performance12 on a number of real-world problems , including the challenging tasks of aligning13 human tumour samples with cancer cell-lines and performing counterfactual in-14 ference on single-cell RNA sequencing data . Incidentally , we find parallels with15 the fair representation learning literature , and demonstrate CoMP has competitive16 performance in learning fair yet expressive latent representations.17 1 Introduction18 Large scale datasets describing the molecular properties of cells , tissues and organs in a state of health19 and disease are commonplace in computational biology . Referred to collectively as ‘ omics data,20 thousands of features are measured per sample and , as single-cell methodologies have developed , it21 is now typical to measure such features across 105–106 samples [ 1 , 2 ] . Given these two properties of22 ‘ omics data , the need for scalable algorithms to learn meaningful low-dimensional representations23 that capture the variability of the data has grown . As such , Variational Autoencoders ( VAEs ) [ 3 , 4 ] 24 have become an important tool for solving a range of modelling problems in the biological sciences25 [ 5 , 6 , 7 , 8 , 9 , 10 ] . One such problem is utilising representations for counterfactual inference,26 e.g . predicting how a certain cell or cell-type , observed only in the control , would have behaved when27 exposed to a drug [ 9 , 10 , 11 ] . Another key problem is removing batch effects—spurious shifts in28 observations due to differing experimental conditions—from data in order to integrate or compare29 multiple datasets [ 5 , 12 , 13 , 14 , 15 ] .30 We present a formal account of these challenges and show that , to a great extent , they can be seen31 as different aspects of a the same underlying problem , namely , that of learning a representation32 that is marginally independent of a condition variable ( e.g . experimental batch , stimulated vs.33 control ) . Figure 1 [ CoMP ] illustrates what this looks like in practice : the complete overlap of the cell34 populations from different conditions in the latent space . This directly addresses batch correction , and35 in the case where we also have a generative model that maps from latent space back to the original36 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . data space , methods that solve this problem can also be applied to counterfactual inference [ 10 ] .37 This same mathematical requirement for independence also occurs in the fair representation learning38 literature , in which we seek a representation that removes a sensitive attribute , e.g . gender.39 Neither the VAE nor the conditional VAE ( CVAE ) [ 16 ] are typically successful at learning repre-40 sentations that achieve this desired independence , as shown in Figure 1 . Despite the CVAE being41 theoretically able to remove batch effects , there is no constraint that prevents it from from separating42 different conditions in latent space . Existing methods use a penalty to encourage the CVAE to learn43 representations that overlap correctly in latent space , with Maximum Mean Discrepancy ( MMD ) [ 17 ] 44 being the most common penalty , applied in the VFAE [ 18 ] and the more recent trVAE [ 10 ] . These45 methods , however , suffer from a number of drawbacks : conceptually , they introduce an extraneous46 discrepancy measure that is not a part of the variational inference framework ; practically , they re-47 quire the choice of , and hyperparameter tuning for , an MMD kernel ; empirically , whilst trVAE is a48 significant improvement over an unconstrained CVAE , Figure 1 [ trVAE ] shows that it can still fail to49 exactly align different conditions in latent space.50 To overcome these difficulties , we introduce Contrastive Mixture of Posteriors ( CoMP ) , a new51 method for learning aligned representations in a CVAE framework . Our method features the novel52 CoMP misalignment penalty that forces the CVAE to remove batch effects . Inspired by contrastive53 learning [ 19 , 20 ] , the penalty encourages representations from different conditions to be close , whilst54 representations from the same condition should be spread out . To achieve this , we approximate the55 requisite marginal distributions using mixtures of the variational posteriors themselves , leading to a56 penalty that does not require an extraneous discrepancy measure or a separately tuned kernel . We57 prove that the CoMP penalty is a stochastic upper bound on a weighted sum of KL divergences , so58 minimising the penalty minimises a well-established statistical divergence measure . We analyse the59 training gradients of the CoMP and MMD penalties , finding key differences that help explain why60 CoMP gradients are generally more stable and better suited to datasets with complex global structure.61 As shown in Figure 1 [ CoMP ] , our method can achieve visually perfect alignment on a number of62 real-world biological datasets . We apply CoMP to two challenging biological problems : 1 ) aligning63 gene expression profiles between tumours and their corresponding cell-lines , as tackled in [ 21 ] and64 2 ) estimating the gene expression profile of an unperturbed cell as if it had been treated with a65 chemical perturbation ( counterfactual inference ) [ 9 ] . We show that CoMP outperforms existing66 methods , achieving state-of-the-art performance on both tasks . Finally , given the connections to fair67 representation learning , we apply CoMP to the problem of learning a representation that is independent68 of gender in the UCI Adult Income dataset [ 22 ] , showing that we can learn a representation that is69 fully independent of the protected attribute whilst maintaining useful information for other prediction70 tasks . CoMP represents a conceptually simple and empirically powerful method for learning aligned71 representation , opening the door to answering high-value questions in biology and beyond.72 2 Background73 2.1 Variational Autoencoders and extensions74 We begin by assuming that we have n observations x1 , . . . , xn of an underlying data distribution.75 Variational autoencoders ( VAEs ) [ 3 , 4 ] explain the high-dimensional observations xi using low76 dimensional representations zi . The standard VAE places a standard normal prior z ∼ p ( z ) on77 the latent variable , and learns a generative model pθ ( x|z ) that reconstructs x using z , alongside78 an inference network qφ ( z|x ) that encodes x to z . Both θ and φ are trained jointly by maximising79 the ELBO , a lower bound on marginal likelihood given by log pθ ( x ) ≥ Eqφ ( z|x ) [ log pθ ( x|z ) ] −80 KL [ qφ ( z|x ) ‖p ( z ) ] . This can be maximised using stochastic optimisers [ 23 , 3 ] . Various extensions81 of the VAE have been proposed , such as the β-VAE [ 24 ] , which scales the KL term of the ELBO82 by a hyperparameter β . Because the isotropic normal prior may limit the expressivity of the model83 [ 25 ] , various authors have considered alternative priors . For example , [ 26 ] proposed the Variational84 Mixture of Posteriors ( VaMP ) prior , that replaces the isotropic Gaussian with a mixture of posteriors85 from the encoder network itself , evaluated at a number of learned pseudo-inputs.86 So far , we have assumed that the only data available are the observations x1 , . . . , xn , but in many87 practical applications we may have additional information such as a condition label for each ob-88 servation . For example , in gene knock-out studies , we have information about which gene was89 targeted for deletion in each cell ; in multi-batch experiments we have information about which exper-90 imental batch each samples was collected in . Thus , we augment our data by considering data pairs91 ( x1 , c1 ) , . . . , ( xn , cn ) where x is the same high-dimensional observation , and c is a label indicating92 the condition or experimental batch that x was collected under.93 Whilst VAEs are theoretically able to model the pairs ( xi , ci ) , it makes sense to build a model that94 explicitly distinguishes between the x and c. The simplest model is the Conditional VAE ( CVAE ) 95 [ 16 ] . In this model , a conditional generative model pθ ( x|z , c ) and a conditional inference network96 qφ ( z|x , c ) are trained using a modified ELBO . A key observation for our work is that the CVAE has97 many different ways to model the data . For example , it can completely ignore the condition c in pθ98 and qφ , reducing to the original VAE . Assuming that x is not independent of c , this failure mode of99 the CVAE would be apparent on a visualization of the representations . For example , different values100 of c might be visible as separate latent clusters , as shown in Figure 1 [ CVAE ] .101 2.2 Counterfactual inference102 If ( xi , ci ) represents an RNA transcript and the gene knock-out applied to the cell , a natural question103 to ask is “ How would the transcript have differed if a different knock-out c′ had been applied ? ” In104 general , counterfactual inference attempts to answer questions of the form “ How would the data have105 changed if ci had been replaced by c′ ? ” Answering counterfactual questions is a notoriously difficult106 task , because they naturally refer to unobservable data [ 27 ] . A principled approach to such questions107 is to adopt the framework of Structural Equation Models [ 28 , 27 ] . For example , we could assume that108 the data generating process is given as in Figure 2 . If this model is correct , counterfactual inference109 in the Pearl framework [ 27 ] can then be performed by : 1. abduction : inferring the latent z from x110 and c using p ( z|x , c ) , 2. action : swap c for c′ , 3. prediction : use p ( x|z , c′ ) to obtain a predictive111 distribution for the counterfactual . Thus , the counterfactual distribution of xi observed with condition112 ci but predicted for condition c′ is given by113 p ( xc=c′ |xi , ci ) = ∫ p ( z|xi , ci ) p ( x|z , c′ ) dz . ( 1 ) In order to make use of this relationship , we must fit a latent variable model [ 29 ] such as a CVAE that114 will estimate the encoding distribution p ( z|xi , ci ) and the generative distribution p ( x|z , c′ ) .115 3 Unifying counterfactual inference , data integration and fairness116 We have seen that batch effect correction , data integration and counterfactual inference are central117 problems of interest for the application of latent variable models in computational biology.118 For counterfactual inference , latent variable models such as the CVAE are increasingly popular119 choices [ 29 ] . The failure mode in which different values of c form separate latent clusters , however,120 can be catastrophic for this application . When this happens , simply switching ci to c′ is not correct,121 we have to account for the shift between clusters [ 9 ] . Mathematically , the latent space clustering122 phenomenon violates the assumption z ⊥⊥ c that is required by the model in Figure 2 . Thus , whilst it123 is not always possible to know when we have found the right causal model [ 30 ] , we can immediately124 say that a model in which z and c are dependent is not correct.125 Another key challenge for computational biology is data integration . Suppose our data126 ( x1 , c1 ) , . . . , ( xn , cn ) in which ci indicates the experimental batch , exhibits batch effects—these are127 changes in the observation xi due to the experimental conditions rather than true changes in the128 underlying biology . One approach to dataset integration is to create a representation z = z ( x , c ) 129 that ‘ subtracts ’ the batch effects . Downstream tasks can then work with z in place of x without130 learning signal based on misleading batch effects . To know when we have successfully subtracted131 batch effects , we might assume that there are no population-level differences between batches . In132 other words , the marginal distribution of z should be the same for each value of the condition c.133 Thirdly , this same notion of building a representation that can not be used to recover c has been studied134 widely in recent literature on fairness [ 31 , 18 , 32 , 33 ] . In particular , if we wish to make a predictive135 rule based on x that does not discriminate between individuals in different conditions c , we can use136 a fair representation z , one which can not be used to recover c , as an intermediate feature and train137 our model using z . Such a representation clearly needs to contain information from x , but without138 containing any information that could be used to recover c.139 To connect these three notions of ‘ alignment in representation space ’ we recall the key components140 of the CVAE—the encoder qφ ( z|x , c ) and decoder pθ ( x|z , c ) —and we now drop the θ , φ subscripts141 for conciseness . The marginal distribution of representations within condition c ∈ C is q ( z|c ) =142 Ep ( x|c ) [ q ( z|x , c ) ] , and the marginal distribution of over all conditions not equal to c is denoted143 q ( z|¬ c ) = ∑ c′∈C , c′ 6=c p ( c ′ ) q ( z|c′ ) ∑ c′∈C , c′ 6=c p ( c ′ ) . ( 2 ) The following Theorem brings together key notions in counterfactual inference , data integration and144 fair representation learning . See Appendix B for the proof.145 Theorem 1 . The following are equivalent : 1 ) z ⊥⊥ c under distribution q , 2 ) for every c , c′ ∈ C,146 q ( z|c ) = q ( z|c′ ) , 3 ) for every c ∈ C , q ( z|c ) = q ( z|¬ c ) , 4 ) the mutual information I ( z , c ) = 0 under147 distribution q , 5 ) z can not predict c better than random guessing.148 4 Contrastive Mixture of Posteriors149 We have seen that counterfactual inference , data integration and fair representation learning can be150 understood through the unified concept of learning a representation such that the latent variable z is151 independent of the condition c under the distribution q , so that the latent clusters with different values152 of c are perfectly aligned . Building off the CVAE , which rarely achieves this in practice , a number of153 authors have attempted to use a penalty term to reduce the dependence of z upon c during training.154 The most successful methods , such as trVAE [ 10 ] , are based on a Maximum Mean Discrepancy155 ( MMD ) [ 17 ] . We discuss this and other common methods in Section 5 . Whilst trVAE and related156 methods can work well , they require an MMD kernel , not a part of the original model , to be specified157 and its parameters to be carefully tuned . Experimentally , we observe that MMD-based methods158 can often struggle when there is complex global structure in the latent space . We also analyse the159 gradients of MMD penalties , showing that they have several undesirable properties.160 We propose a novel method to ensure the conditions of Theorem 1 do hold in a CVAE model . Our161 penalty is based on posterior distributions obtained from the model encoder itself . That is , we do not162 introduce any external discrepancy measure , rather we propose a penalty term that arises naturally163 from the model itself . Taking our inspiration from contrastive learning [ 19 , 20 ] and the VaMP prior164 [ 26 ] , we suggest a novel penalty to enforce equation condition 3 ) of Theorem 1 . This equation165 requires the equality of the marginal distribution q ( z|c ) and q ( z|¬ c ) for each c ∈ C. In practice,166 these marginal distributions can be approximated by finite mixtures . To encourage greater overlap167 between q ( z|c ) and q ( z|¬ c ) , we can encourage points with the condition c to be in areas of high168 density under the representation distribution for other conditions , i.e . areas in which q ( z|¬ c ) is also169 high . To encourage this , we can add the penalty term P0 ( zi , ci ) = − log q ( zi|¬ ci ) to the objective170 for the data pair ( xi , ci ) . When we minimise P0 , this brings the representations of samples under171 condition ci towards regions of high density under q ( z|¬ c ) . Since the density q ( z|¬ c ) is not known172 in closed form , we approximate q ( z|¬ c ) using other points in the same training batch as ( xi , ci ) .173 Indeed , suppose we have a batch ( x1 , c1 ) , ... , ( xB , cB ) . We let Ic denote the subset of indices for174 which cj = c and I¬c denote its complement . We use the approximation175 log q ( zi|¬ ci ) ≈ log ( 1 |I¬ci | ∑ j∈I¬ci q ( zi|xj , cj ) ) ( 3 ) and we will show in Theorem 2 , this approximation in fact leads to a valid stochastic bound.176 It may happen that the penalty P0 causes points to become too tightly clustered . Indeed , the177 penalty encourages latent variables to gravitate towards high density regions of q ( z|¬ci ) . Inspired178 by contrastive learning , we include a second term which promotes higher entropy of the marginal,179 thereby avoiding tight clusters of points . Combined with P0 , this leads us to a second penalty180 P1 ( zi , ci ) = log q ( zi|ci ) − log q ( zi|¬ ci ) . Again , the density q ( z|c ) is not known in closed form,181 but we can approximate it using points within the same training batch in a similar fashion to ( 3 ) .182 Combining both approximations and taking the mean over the batch gives our Contrastive Mixture of183 Posteriors ( CoMP ) misalignment penalty184 CoMP penalty = 1 B B∑ i=1 log ( 1 |Ici | ∑ j∈Ici q ( zi|xj , ci ) ) − log ( 1 |I¬ci | ∑ j∈I¬ci q ( zi|xj , cj ) ) . ( 4 ) where x1 : B , c1 : B , z1 : B ∼ ∏B i=1 p ( xi , ci ) q ( zi|xi , ci ) is a random training batch of size B , Ic denotes185 the subset of { 1 , . . . , B } with condition c and I¬c = { 1 , . . . , B } \ Ic . Our method therefore utilises186 a training penalty for CVAE-type models that encourages the conditions of Theorem 1 to hold by187 using mixtures of the variational posteriors themselves to approximate q ( z|c ) and q ( z|¬ c ) . We do188 not introduce an additional kernel or hyperparameter-heavy discrepancy measures.189 As hinted at by the definition of P1 , CoMP can be seen as approximating a symmetrised KL-190 divergence between the distributions q ( z|c ) and q ( z|¬ c ) . In fact , the following theorem shows that191 the CoMP misalignment penalty is a stochastic upper bound on a weighted sum of KL-divergences.192 Theorem 2 . The CoMP misalignment penalty satisfies193 E∏B i=1 p ( xi , ci ) q ( zi|xi , ci ) [ 1 B B∑ i=1 log ( 1 |Ici | ∑ j∈Ici q ( zi|xj , ci ) ) − log ( 1 |I¬ci | ∑ j∈I¬ci q ( zi|xj , cj ) ) ] ≥ ∑ c∈C p ( c ) KL [ q ( z|c ) ||q ( z|¬ c ) ] and the bound becomes tight as B →∞.194 The proof is presented in Appendix B . Our result reveals that our new penalty directly enforces195 condition 3 ) of Theorem 1 by reducing the KL divergence between each pair q ( z|c ) , q ( z|¬ c ) weighted196 by p ( c ) . As with standard contrastive learning , our method benefits from larger batch sizes . We197 add the CoMP misalignment penalty to the familiar β-VAE objective to give our complete training198 objective for a batch of size B as199 LCoMPB = 1 B B∑ i=1 [ log p ( xi|zi , ci ) + β log p ( zi ) q ( zi|xi , ci ) − γ log ( 1 |Ic| ∑ j∈Ic q ( zi|xj , ci ) 1 |I¬c| ∑ j∈I¬c q ( zi|xj , cj ) ) ] ( 5 ) with one new hyperparameter γ that controls the strength of the regularisation we apply to enforce200 the requirements z ⊥⊥ c. Theorem 2 shows that , if LβB is the standard β-VAE objective , then we are201 maximising E [ LCoMPB ] ≤ E [ LβB ] − γ ∑ c∈C p ( c ) KL [ q ( z|c ) ||q ( z|¬ c ) ] .202 4.1 Analysing CoMP gradients203 Before presenting empirical results on the performance of CoMP , we attempt to understand how it204 differs from existing penalties in the literature . Specifically , we compare CoMP with a Gaussian205 posterior family with MMD using a Radial Basis Kernel [ 34 ] . In Appendix C , we show that both206 methods can be interpreted as applying a penalty to each element zi , ci of the training batch . We207 show further that , under certain conditions , the gradient of the MMD penalty for zi , ci takes the form208 ∇ziPMMD ( zi , ci ) = 2 |Ici | 2 ∑ j∈Ici e−‖zi−zj‖ 2 ( zj−zi ) − 4 |I¬ci | |Ici | ∑ j∈I¬ci e−‖zi−zj‖ 2 ( zj−zi ) , ( 6 ) whilst the CoMP penalty gradient takes the form209 ∇ziPCoMP ( zi , ci ) = 2 ∑ j∈Ici e−‖zi−µzj | 2 ( µzj − zi ) B ∑ j∈Ici e−‖zi−µzj | 2 − 2 ∑ j∈I¬ci e−‖zi−µzj ‖ 2 ( µzj − zi ) B ∑ j∈I¬ci e−‖zi−µzj ‖ 2 ( 7 ) where µzj is the variational mean for zj . One important feature of the MMD gradients is that , if210 ‖zi − zj‖2 is large for all j 6= i , for instance when the point zi is part of an isolated cluster , then211 the gradient to update the representation zi will be small . So if zi is already very isolated from the212 distribution q ( z|¬ ci ) , then the gradients bringing it closer to points with condition ¬ ci will be small.213 In comparison to the MMD gradient , it can be seen that gradients for CoMP are self-normalised . This214 means that the gradient through zi will be large , even when zi is very far away from any points with215 condition ¬ ci . This , in turn , suggests that that CoMP is likely to be preferable to MMD when we have216 a number of isolated clusters or interesting global structure in latent space , something which often217 occurs with biological data . The CoMP approach also bears a relationship with nearest-neighbour218 approaches [ 35 ] . Indeed , for a Gaussian posterior as σ → 0 , the ¬ ci term of the gradient places all219 its weight on the nearest element of the batch under condition ¬ ci.220 5 Related Work221 The problem of batch correction in data integration has been addressed using linear [ 12 , 13 ] and222 nonlinear methods [ 14 , 15 ] that perform transformations of the original feature space . In both cases,223 the goal is to transform the feature space so that information related to the scientific question of224 interest is retained while dependence on the batch ( or nuisance covariate ) is reduced . Methods based225 on representation learning learn a low-dimensional representation , z = q ( x ) , which is independent226 of nuisance factors while also being a faithful representation of the original data [ 18 , 5 , 36 , 10 , 37 ] .227 Of these , the work that is most similar to ours are the VFAE [ 18 ] , in which the authors introduce an228 MMD [ 17 ] penalty to encourage the marginal distributions of z under different values of c to be close,229 and the trVAE [ 10 ] , where the MMD penalty is applied to the output of the first layer of the decoder,230 rather than to z directly . Representation learning algorithms for counterfactual inference have been231 shown to benefit from a penalty enforcing distributional similarity between the representations of the232 treated and untreated samples [ 12 ] . Elsewhere , authors have applied the variational autoencoder to233 inference on causal graphs [ 38 , 39 , 40 ] .234 6 Experiments235 We perform experiments on three datasets ; 1 ) Tumour / Cell Line : bulk expression profiles of236 tumours and cancer cell-lines across 39 different cancer types ; 2 ) Single-cell PBMCs : single-cell237 gene expression ( scRNA-seq ) profiles of interferon ( IFN ) -β stimulated and untreated peripheral blood238 mononuclear cells ( PBMCs ) [ 41 ] ; 3 ) UCI Adult Income : personal information relating to education,239 marriage status , ethnicity , self-reported gender of census participants and a binary high / low income240 label ( $ 50,000 threshold ) [ 22 ] . All experiments used a 90/10 training/validation split.241 The two broad objectives across our experiments are 1 ) to demonstrate the extent to which the two242 random variables zi and ci are independent , and 2 ) to quantify useful information retained in zi . To243 benchmark CoMP on the first objective , we use the following pair of k nearest-neighbor metrics:244 kBETk , α [ 42 ] , the metric used to evaluate batch correction methods in biology , and a local Silhouette245 Coefficient [ 43 ] sk , c . In both cases a low value close to zero would indicate good local mixing of246 sample representations . As for the second objective , if we assume the existence of an additional247 discrete label di that represents information one wishes to preserve – in the Tumour / Cell Line248 case , di is the cancer type , while for the PBMC experiment , it refers to cell type – then we calculate249 kBET and s separately for every fixed-di subpopulation and take the mean . We refer to these as the250 mean Silhouette Coefficient s̃k , c and the mean kBET metric m-kBET respectively . Full details of the251 datasets and metrics are given in Appendix D.252 6.1 Alignment of tumour and cell-line samples253 Despite their widespread use in pre-clinical cancer studies , cancer cell-lines are known to have254 significantly different gene expression profiles compared to their corresponding tumour samples.255 Here we evaluate the ability of CoMP to factorise out the tumour / cell line condition from its latent256 representations . This can be seen as both a dataset integration and batch effect correction task . In257 addition to the set of k nearest neighbor-based mixing evaluations , we train a Random Forest model258 on the representations of the tumour samples and their cancer-type labels and assess the prediction259 accuracy on held-out cell lines . To match the results from [ 21 ] , the evaluations are performed on the260 2D UMAP projections , The results are presented in Table 1.261 As expected , both the VAE and CVAE baselines fail at the mixing task ; the three explicitly penalised262 VAE models and , to a lesser extent , the Cellinger method have good mixing performances , with263 CoMP outperforming the benchmark models by a significant margin on the silhouette coefficient264 and kBET metric , while successfully maintaining a high accuracy in the cancer-type prediction265 task . We also see from Figure 3 that CoMP representations have the fewest instances of isolated266 tumour-only clusters . Finally , from our evaluation on the s̃ and m-kBET metrics , we can deduce that267 the occurrence of cell lines of one cancer type erroneously clustering around tumours of a different268 type is less frequent for CoMP compared to the other models . In Appendix D we qualitatively validate269 this for several example clusters.270 6.2 Interventions271 Obtaining molecular measurements from biological tissues typically requires destructive sampling.272 For example , to obtain scRNA-seq data , each cell is lysed so that the RNA molecules contained273 within it can be extracted and sequenced . This process destroys each cell , meaning that we are274 unable to study the gene expression profile of the same cell over time or under multiple experimental275 conditions . As we discussed in Section 2.2 , counterfactual inference can be used to predict how the276 molecular status of a destroyed biological sample would have differed if it were measured under277 different experimental conditions , such as applications of different drugs.278 To assess CoMP ’ s utility in counterfactual inference , we trained it on scRNA-seq data from PBMCs279 that were either stimulated with IFN-β or left untreated ( control ) [ 41 ] . It is clear from Figure 4 that280 IFN-β stimulation causes clear shifts in the latent space between stimulated and control cells from the281 same cell type . Noticeably , the CD14 and CD16 monocyte and dendritic cell ( DC ) populations see282 greater shifts in their gene expression after stimulation . CVAE fails to align these particular cell types283 in the latent space , while trVAE , VFAE and CoMP perform better . However , stimulated and control284 cells are better mixed in the latent space derived from CoMP than those from the other models ( see285 metrics presented in Appendix D ) .286 Next we perform a counterfactual prediction task under a IFN-β control-to-stimulation variable287 swap , i.e . the gene expression profiles for control cells were reconstructed through the decoder with288 the condition , c→ stimulated . This means we utilise equation ( 1 ) with our encoder qφ ( z|x , c ) and289 decoder pθ ( x|z , c′ ) in place of p ( z|x , c ) and p ( x|z , c′ ) . The degree to which the models respect the290 requirement z ⊥⊥ c will influence the quality of predictions . Figure 5 shows how the profiles of291 ( actual ) stimulated cells differ from the counterfactual predictions for a selection of cell types ( see292 Appendix D for the complete set of results ) . We see that baseline models tend to systematically293 underestimate the expression of genes up-regulated by stimulation and overestimate those down-294 regulated . CoMP outperforms all other models by accurately predicting the expression alterations295 brought about by stimulation.296 6.3 Fair Classification297 The goal for this fair classification task is to learn a representation on the Adult Income dataset that is298 not predictive of an individual ’ s gender whilst still being predictive of their income . We compute a299 baseline by predicting gender and income labels directly from the input data and compare our method300 to the published results for the VFAE [ 18 ] and the trVAE . We also include results for a standard301 VAE and CVAE . Unlike in [ 18 ] , where the representations z are sampled from the posterior before302 classification , our experiments used the posterior means to avoid the noise from sampling acting to303 mask the inclusion of predictive information about gender in the encodings.304 CoMP achieves a gender accuracy that is close to random ( 67.5 % ) , tying with the VFAE results from305 [ 18 ] whilst also remaining competitive with the other methods on income accuracy ( Table 2 ) . CoMP306 also outperforms all methods on the nearest neighbour and silhouette metrics ( Table 2 ) . Latent space307 mixing between males and females can be seen qualitatively in the 2D UMAP projection ( Figure 6 ) .308 7 Conclusion309 Limitations We presented Contrastive Mixture of Posteriors ( CoMP ) as an effective means to310 perform batch correction , data integration , counterfactual inference and fair representation learning311 in a CVAE framework . Whilst CoMP covers the majority of common use-cases for these tasks , there312 are several limitations that are avenues of future research . For example , in scRNA-seq analysis , there313 is often the need to integrate more than two datasets together , or to adjust for continuous condition314 variables . Mathematically , CoMP is applicable to any number of discrete conditions , and it would be315 interesting to apply it to a setting with > 2 conditions . Extensions of CoMP could tackle the case316 of a continuous condition variable . Additionally , CoMP requires the condition variable c to be fully317 observed : future work might attempt to generalise to the partially observed case.318 Summary We identified marginal independence between the representation z and condition c as319 the mathematical thread linking data integration , counterfactual inference and fairness . We proposed320 CoMP , a novel method to enforce this independence requirement in practice . We saw that CoMP321 has several attractive theoretic properties . First , CoMP only uses the variational posteriors , requiring322 no additional discrepancy measures such as MMD . Second , we proved that the CoMP penalty can323 be interpreted as an upper-bound on a weighted sum of KL divergences , connecting it to a well-324 founded divergence measure . Third , we demonstrated that , unlike MMD , CoMP gradients have a325 self-normalising property , allowing one to obtain strong gradients for distant points in a latent space326 with complex global structure . Empirically , we demonstrated CoMP ’ s performance when applied to327 two biological and one fair representation learning dataset . These biological datasets are of critical328 importance in drug discovery , for example matching cell-lines to tumours for effective pre-clinical329 assay development of anti-cancer compounds . Overall , CoMP has the best in class performance on330 all tasks across a range of metrics that measure either latent space mixing or fairness.331 References332 [ 1 ] Valentine Svensson , Roser Vento-Tormo , and Sarah A Teichmann . Exponential scaling of single-cell333 rna-seq in the past decade . Nature protocols , 13 ( 4 ) :599–604 , 2018.334 [ 2 ] Aviv Regev , Sarah A Teichmann , Eric S Lander , Ido Amit , Christophe Benoist , Ewan Birney , Bernd335 Bodenmiller , Peter Campbell , Piero Carninci , Menna Clatworthy , et al . Science forum : the human cell336 atlas . Elife , 6 : e27041 , 2017.337 [ 3 ] Diederik P Kingma and Max Welling . Auto-encoding variational bayes . arXiv preprint arXiv:1312.6114,338 2013.339 [ 4 ] Danilo Jimenez Rezende , S. Mohamed , and Daan Wierstra . Stochastic backpropagation and approximate340 inference in deep generative models . In ICML , 2014.341 [ 5 ] Romain Lopez , Jeffrey Regier , Michael B Cole , Michael I Jordan , and Nir Yosef . Deep generative modeling342 for single-cell transcriptomics . Nature methods , 15 ( 12 ) :1053–1058 , 2018.343 [ 6 ] Gregory P Way and Casey S Greene . Extracting a biologically relevant latent space from cancer transcrip-344 tomes with variational autoencoders . In PACIFIC SYMPOSIUM ON BIOCOMPUTING 2018 : Proceedings345 of the Pacific Symposium , pages 80–91 . World Scientific , 2018.346 [ 7 ] Dongfang Wang and Jin Gu . Vasc : dimension reduction and visualization of single-cell rna-seq data by347 deep variational autoencoder . Genomics , proteomics & bioinformatics , 16 ( 5 ) :320–331 , 2018.348 [ 8 ] Christopher Heje Grønbech , Maximillian Fornitz Vording , Pascal N Timshel , Casper Kaae Sønderby,349 Tune H Pers , and Ole Winther . scvae : Variational auto-encoders for single-cell gene expression data.350 Bioinformatics , 36 ( 16 ) :4415–4422 , 2020.351 [ 9 ] Mohammad Lotfollahi , F Alexander Wolf , and Fabian J Theis . scGen predicts single-cell perturbation352 responses . Nature methods , 16 ( 8 ) :715 , 2019.353 [ 10 ] Mohammad Lotfollahi , Mohsen Naghipourfar , Fabian J Theis , and F Alexander Wolf . Conditional354 out-of-sample generation for unpaired data using trVAE . arXiv preprint arXiv:1910.01791 , 2019.355 [ 11 ] Matthew Amodio , D. V. Dijk , R. Montgomery , Guy Wolf , and Smita Krishnaswamy . Out-of-sample356 extrapolation with neuron editing . arXiv : Quantitative Methods , 2018.357 [ 12 ] W. Johnson , Cheng Li , and Ariel Rabinovic . Adjusting batch effects in microarray expression data using358 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 variable360 analysis . PLoS Genetics , 3 , 2007.361 [ 14 ] Laleh Haghverdi , A. Lun , Michael D. Morgan , and J. Marioni . Batch effects in single-cell rna-sequencing362 data are corrected by matching mutual nearest neighbors . Nature Biotechnology , 36:421–427 , 2018.363 [ 15 ] Allison Warren , Andrew Jones , Tsukasa Shibue , William C Hahn , Jesse S Boehm , Francisca Vazquez,364 Aviad Tsherniak , and James M McFarland . Global computational alignment of tumor and cell line365 transcriptional profiles . Nature Communications , 12 ( 22 ) , 2021.366 [ 16 ] Kihyuk Sohn , Honglak Lee , and Xinchen Yan . Learning structured output representation using deep367 conditional generative models . In Advances in neural information processing systems , pages 3483–3491,368 2015.369 [ 17 ] Arthur Gretton , Karsten M Borgwardt , Malte J Rasch , Bernhard Schölkopf , and Alexander Smola . A370 kernel two-sample test . The Journal of Machine Learning Research , 13 ( 1 ) :723–773 , 2012.371 [ 18 ] Christos Louizos , Kevin Swersky , Yujia Li , Max Welling , and Richard Zemel . The variational fair372 autoencoder . arXiv preprint arXiv:1511.00830 , 2015.373 [ 19 ] Aaron van den Oord , Yazhe Li , and Oriol Vinyals . Representation learning with contrastive predictive374 coding . arXiv preprint arXiv:1807.03748 , 2018.375 [ 20 ] Ting Chen , Simon Kornblith , Mohammad Norouzi , and Geoffrey Hinton . A simple framework for376 contrastive learning of visual representations . arXiv preprint arXiv:2002.05709 , 2020.377 [ 21 ] Allison Warren , Yejia Chen , Andrew Jones , Tsukasa Shibue , William C Hahn , Jesse S Boehm , Francisca378 Vazquez , Aviad Tsherniak , and James M McFarland . Global computational alignment of tumor and cell379 line transcriptional profiles . Nature Communications , 12 ( 1 ) :1–12 , 2021.380 [ 22 ] Dheeru Dua and Casey Graff . UCI machine learning repository , 2017.381 [ 23 ] Herbert Robbins and Sutton Monro . A stochastic approximation method . The annals of mathematical382 statistics , pages 400–407 , 1951.383 [ 24 ] Irina Higgins , Loic Matthey , Arka Pal , Christopher Burgess , Xavier Glorot , Matthew Botvinick , Shakir384 Mohamed , and Alexander Lerchner . beta-vae : Learning basic visual concepts with a constrained variational385 framework . International Conference on Learning Representations , 2017.386 [ 25 ] Emile Mathieu , Tom Rainforth , N. Siddharth , and Yee Whye Teh . Disentangling disentanglement in387 variational autoencoders . In In International Conference on Machine Learning , pages 4402–4412 . PMLR,388 2019.389 [ 26 ] Jakub Tomczak and Max Welling . Vae with a vampprior . In International Conference on Artificial390 Intelligence and Statistics , pages 1214–1223 . PMLR , 2018.391 [ 27 ] Judea Pearl . Causality . Cambridge university press , 2009.392 [ 28 ] Kenneth A. Bollen . Structural equation models . Wiley , 2005.393 [ 29 ] Fredrik Johansson , Uri Shalit , and David Sontag . Learning representations for counterfactual inference . In394 International conference on machine learning , pages 3020–3029 , 2016.395 [ 30 ] Jonas Peters , Joris Mooij , Dominik Janzing , and Bernhard Schölkopf . Identifiability of causal graphs using396 functional models . arXiv preprint arXiv:1202.3757 , 2012.397 [ 31 ] Rich Zemel , Yu Wu , Kevin Swersky , Toni Pitassi , and Cynthia Dwork . Learning fair representations . In398 International conference on machine learning , pages 325–333 . PMLR , 2013.399 [ 32 ] Matt J. Kusner , Joshua R. Loftus , Chris Russell , and Ricardo Silva . Counterfactual fairness . arXiv preprint400 arXiv:1703.06856 , 2017.401 [ 33 ] Craig A Glastonbury , Michael Ferlaino , Christoffer Nellåker , and Cecilia M Lindgren . Adjusting for402 confounding in unsupervised latent representations of images . arXiv preprint arXiv:1811.06498 , 2018.403 [ 34 ] Jean-Philippe Vert , Koji Tsuda , and Bernhard Schölkopf . A primer on kernel methods . Kernel methods in404 computational biology , 47:35–70 , 2004.405 [ 35 ] Ke Li and Jitendra Malik . Implicit maximum likelihood estimation . arXiv preprint arXiv:1809.09087,406 2018.407 [ 36 ] Romain Lopez , Jeffrey Regier , Michael I Jordan , and Nir Yosef . Information constraints on auto-encoding408 variational bayes . In Advances in Neural Information Processing Systems , pages 6114–6125 , 2018.409 [ 37 ] Kaspar Märtens and Christopher Yau . Neural decomposition : Functional anova with variational autoen-410 coders . In International Conference on Artificial Intelligence and Statistics , pages 2917–2927 . PMLR,411 2020.412 [ 38 ] Christos Louizos , Uri Shalit , Joris M Mooij , David Sontag , Richard Zemel , and Max Welling . Causal413 effect inference with deep latent-variable models . In Advances in Neural Information Processing Systems,414 pages 6446–6456 , 2017.415 [ 39 ] H. Kim , Seungjae Shin , Joonho Jang , Kyungwoo Song , Weonyoung Joo , Wanmo Kang , and Il-Chul Moon.416 Counterfactual fairness with disentangled causal effect variational autoencoder . In AAAI , 2021.417 [ 40 ] Stephen R. Pfohl , Tony Duan , Daisy Yi Ding , and Nigam H. Shah . Counterfactual reasoning for fair clinical418 risk prediction . In Finale Doshi-Velez , Jim Fackler , Ken Jung , David Kale , Rajesh Ranganath , Byron419 Wallace , and Jenna Wiens , editors , Proceedings of the 4th Machine Learning for Healthcare Conference,420 volume 106 of Proceedings of Machine Learning Research , pages 325–358 , Ann Arbor , Michigan , 09–10421 Aug 2019 . PMLR.422 [ 41 ] Hyun Min Kang , Meena Subramaniam , Sasha Targ , Michelle Nguyen , Lenka Maliskova , Elizabeth423 McCarthy , Eunice Wan , Simon Wong , Lauren Byrnes , Cristina M Lanata , et al . Multiplexed droplet424 single-cell rna-sequencing using natural genetic variation . Nature biotechnology , 36 ( 1 ) :89 , 2018.425 [ 42 ] Maren Büttner , Zhichao Miao , F Alexander Wolf , Sarah A Teichmann , and Fabian J Theis . A test metric426 for assessing single-cell rna-seq batch correction . Nature methods , 16 ( 1 ) :43–49 , 2019.427 [ 43 ] Peter J Rousseeuw . Silhouettes : a graphical aid to the interpretation and validation of cluster analysis.428 Journal of computational and applied mathematics , 20:53–65 , 1987.429 [ 44 ] Christian P. Robert and Judith Rousseau . How Principled and Practical Are Penalised Complexity Priors ? 430 Statistical Science , 32 ( 1 ) :36–40 , February 2017.431 [ 45 ] Murray Aitkin . Posterior Bayes Factors . Journal of the Royal Statistical Society . Series B ( Methodological ) ,432 53 ( 1 ) :111–142 , 1991.433 [ 46 ] Paul Fearnhead and Dennis Prangle . Constructing summary statistics for approximate Bayesian computa-434 tion : semi-automatic approximate Bayesian computation . Journal of the Royal Statistical Society : Series B435 ( Statistical Methodology ) , 74 ( 3 ) :419–474 , 2012.436 [ 47 ] Adam Foster , Martin Jankowiak , Matthew O ’ Meara , Yee Whye Teh , and Tom Rainforth . A unified437 stochastic gradient approach to designing bayesian-optimal experiments . In International Conference on438 Artificial Intelligence and Statistics , pages 2959–2969 . PMLR , 2020.439 [ 48 ] J. Weinstein , E. Collisson , G. Mills , K. Shaw , B. Ozenberger , Kyle Ellrott , I. Shmulevich , C. Sander , and440 Joshua M. Stuart . The Cancer Genome Atlas Pan-Cancer analysis project . Nature Genetics , 45:1113–1120,441 2013.442 [ 49 ] DS Gerhard , S Hunger , C Lau , J Maris , P Meltzer , S Meshinchi , E Perlman , J Zhang , J Guidry-Auvil , and443 M Smith . Therapeutically applicable research to generate effective treatments ( target ) project : Half of444 pediatric cancers have their own '' driver '' genes . In PEDIATRIC BLOOD & CANCER , volume 65 , pages445 S45–S45 . WILEY 111 RIVER ST , HOBOKEN 07030-5774 , NJ USA , 2018.446 [ 50 ] M. Ghandi , F. Huang , J. Jané-Valbuena , G. Kryukov , Christopher Lo , E. McDonald , J. Barretina , E. Gelfand,447 C. Bielski , Haoxin Li , Kevin Hu , Alexander Y. Andreev-Drakhlin , J. Kim , J. Hess , B. Haas , F. Aguet,448 B. Weir , M. Rothberg , B. Paolella , M. Lawrence , Rehan Akbani , Y. Lu , Hong L. Tiv , P. Gokhale , Antoine449 de Weck , Ali Amin Mansour , C. Oh , J. Shih , Kevin Hadi , Yanay Rosen , J. Bistline , K. Venkatesan,450 Anupama Reddy , Dmitriy Sonkin , Manway Liu , J. Lehár , J. Korn , D. Porter , M. Jones , J. Golji , G. Capon-451 igro , Jordan E. Taylor , C. Dunning , Amanda L Creech , Allison Warren , James M. McFarland , Mahdi452 Zamanighomi , A. Kauffmann , Nicolas Stransky , M. Imieliński , Y. Maruvka , A. Cherniack , Aviad Tsher-453 niak , F. Vazquez , J. Jaffe , A . A . Lane , D. Weinstock , C. Johannessen , Michael P. Morrissey , F. Stegmeier,454 R. Schlegel , W. Hahn , G. Getz , G. Mills , J. Boehm , T. Golub , L. Garraway , and W. Sellers . Next-generation455 characterization of the Cancer Cell Line Encyclopedia . Nature , 569:503–508 , 2019.456 Checklist457 1 . For all authors ... 458 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s459 contributions and scope ? [ Yes ] 460 ( b ) Did you describe the limitations of your work ? [ Yes ] See Section 7461 ( c ) Did you discuss any potential negative societal impacts of your work ? [ N/A ] 462 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to463 them ? [ Yes ] 464 2 . If you are including theoretical results ... 465 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] 466 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] See Appendix B467 3 . If you ran experiments ... 468 ( a ) Did you include the code , data , and instructions needed to reproduce the main experi-469 mental results ( either in the supplemental material or as a URL ) ? [ Yes ] See Appendix470 D471 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they472 were chosen ) ? [ Yes ] See Appendix D473 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experi-474 ments multiple times ) ? [ Yes ] See Appendix D475 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type476 of GPUs , internal cluster , or cloud provider ) ? [ Yes ] See Appendix D477 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 478 ( a ) If your work uses existing assets , did you cite the creators ? [ Yes ] 479 ( b ) Did you mention the license of the assets ? [ Yes ] 480 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ Yes ] 481 See Appendix D482 ( d ) Did you discuss whether and how consent was obtained from people whose data you ’ re483 using/curating ? [ N/A ] 484 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable485 information or offensive content ? [ N/A ] 486 5 . If you used crowdsourcing or conducted research with human subjects ... 487 ( a ) Did you include the full text of instructions given to participants and screenshots , if488 applicable ? [ N/A ] 489 ( b ) Did you describe any potential participant risks , with links to Institutional Review490 Board ( IRB ) approvals , if applicable ? [ N/A ] 491 ( c ) Did you include the estimated hourly wage paid to participants and the total amount492 spent on participant compensation ? [ N/A ] 493

Summary: The paper proposes a penalty to learn representations that are conditionally independent of some observed variables. The idea is to use the contrastive loss that encourage the representations from the same conditional variable to be spread out and those from different conditional variable to be close. The paper demonstrates application to fairness, counterfactual inference, and data integration, and demonstrate the effectiveness of the penalty.