Counterfactual Generative Modeling with Variational Causal Inference
Authors: Yulun Wu
Louis McConnell
Claudia Iriondo
Venue: ICLR 2025 Poster (2025)
Keywords: deep probabilistic modeling
bayesian networks
variational inference
causal inference
counterfactual generative modeling
causal machine learning
deep unsupervised learning
graphical models
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Abstract: Estimating an individual's potential outcomes under counterfactual treatments is a challenging task for traditional causal inference and supervised learning approaches when the outcome is high-dimensional (e.g. gene expressions, facial images) and covariates are relatively limited. In this case, to predict one's outcomes under counterfactual treatments, it is crucial to leverage individual information contained in the observed outcome in addition to the covariates. Prior works using variational inference in counterfactual generative modeling have been focusing on neural adaptations and model variants within the conditional variational autoencoder formulation, which we argue is fundamentally ill-suited to the notion of counterfactual in causal inference. In this work, we present a novel variational Bayesian causal inference framework and its theoretical backings to properly handle counterfactual generative modeling tasks, through which we are able to conduct counterfactual supervision end-to-end during training without any counterfactual samples, and encourage disentangled exogenous noise abduction that aids the correct identification of causal effect in counterfactual generations. In experiments, we demonstrate the advantage of our framework compared to state-of-the-art models in counterfactual generative modeling on multiple benchmarks.

Published as a conference paper at ICLR 2025
COUNTERFACTUAL GENERATIVE MODELING WITH VARIATIONAL CAUSAL INFERENCE
Yulun Wu
Louie McConnell Genentech mcconn 3êgene
Claudia Iriondo Genentech ir www.ooooocgeeec.ceoooo
University of California, Berkeley edu
ABSTRACT
Estimating an individual's counterfactual outcomes under interventions is a chal- lenging task for traditional causal inference and supervised learning approaches when the outcome is high-dimensional (e.g. gene expressions, facial images) and covariates are relatively limited. In this case, to predict one's outcomes under counterfactual treatments, it is crucial to leverage individual information contained in the observed outcome in addition to the covariates. Prior works using variational inference in counterfactual generative modeling have been focusing neural adaptations and model variants within the conditional variational autoencoder for- mulation, which we argue is fundamentally ill-suited to the notion of counterfactual in causal inference. In this work, we present a novel variational Bayesian causal inference framework and its theoretical backings to properly handle counterfactual generative modeling tasks, through which we are able to conduct counterfactual supervision end-to-end during training without any counterfactual samples, and en- courage disentangled exogenous noise abduction that aids the correct identification of causal effect counterfactual generations. experiments, we demonstrate the advantage of our framework compared to state-of-the-an models in counterfactual generative modeling on multiple benchmarks.
INTRODUCTION
In traditional causal inference, heterogeneous treatment effect is typically formulated as the interven- tional model p(Y [X, do()") (or E[Y [X, do(T))) with outcome Y, covariates X, treatment T, and estimated by the observed covariate-specifi efficacy pX (X,T) T) under the ignorability assumption (Rosenbaum & Rubin 1983 Pearl 1995). However, in cases such as the single-cell perturbation datasets (Dixit et al. 2016;] Norman et al. 2019, Schmidt et al., (022) where Y has thousands of dimensions while X has only two or three categorical features, such model could hardly be relied on to produce useful individualized results. To construct outcome Y' under a counterfactual treatment T" in such high-ding -igh--imensional outcome scenario, it is important and necessary to learn the individual- specific efficacy X,, conditioning on the factual outcome itself, which leverages the rich information embedded in the factual outcome that cannot be recovered by the handful of covariates. For example, given a cell with type A,549 (X) that received SAHA drug treatment (T), we may want know what its gene expression profile (Y') would have looked like if it had received Dex drug treatment (I") instead. In this case, we would want to take the profiles of other cells with type A549 that indeed received Dex as reference, such that the counterfactual construction Y' exhibits the treatment characteristics of Dex on A549, but would also want to extract as much individual information as possible from its own expression profile (Y) such that the counterfactual construction could preserve this cell's individuality such as cell state.
Yet the lack of observability of counterfactual outcome Y' makes this objective intractable and hence presents major difficulty for supervised learning when Y is taken as an input. In previous works
that involve self-supervised counterfactual generation (Kocaoglu et al. 2017 Louizos et al. 2017
Yoon et al.] 2018, Pawlowski al. 2020 Yang et al. [202] Kim et al. 2021 Lotfollahi et al. 2021
Sauer & Geiger 2021 Feng et al. 2022 Shen et al. (2022) based on Variational Autoencoders (VAE)
(Kingma& & Welling 2013) and Generative Adversarial Networks (GAN) (Goodfellow et al. 2014),
there is no counterfactual supervision during training, and hence no explicit regulation on the trade-off
between treatment characteristics and individuality. When such regulation does not exist, there is

---

Published as a conference paper at ICLR 2025
disconnection between training and counterfactual inference. To see this, consider a conditional generative model ge (y. t) that is only supervised on the reconstruction loss with respect to outcome y during training. In this case, the model is free to ignore the extra treatment and converge to state such that go (y, 1) go go 0) (Chen et al. 2016), hence preserving maximum individuality yet minimum treatment characteristics in counterfactual construction go (y: E') during inference. This is particularly an issue for high-fidelity counterfactual generation with state-of-the-ar Hierarchical Variational Autoencoders (HVAE) (Vahdat & Kautz 2020) Child 2020), as the bottleneck latent representations possess much larger number of dimensions compared to the observed outcome y. In more recent works on counterfactual generative modeling with HVAEs and diffusion models
(Sanchez & Tsattaris 2022 Monteiro et al. 2023; Ribeiro et al.) 2023), only Ribeiro et al. (2023) touched on the issue of counterfactual supervision, yet its proposed method does not have a theoretical backing in variational inference, nor could it be conducted end-to-end during training. As we will see in the proceeding sections of this work, there is a fundamental incompatibility of the conditional VAE formulation with counterfactual generative modeling, which is the root cause of such issue. In fact, this incompatibility also exhibits itself in another important aspect of counterfactual generative modeling which is latent disentanglement latent representation) in conditional HVAEs and diffusion models is able to encode individuality better than traditional VAEs, but at the same time is still heavily entangled with treatment variables and often result in lingering characteristics of the factual treatment in counterfactual construction (Monteiro) al. 2023 Sanchez & Tsaftaris 2022). More introduction regarding latent disentanglement can be
the partially abducted exogenous noise (i.e. shallow-level
found in Appendix A
Our contributions in this work are summarized as follow: 1) we propose a formulation and stochastic optimization scheme for counterfactual generative modeling by specifically formulating counter- factual variables and using variational inference to derive the evidence lower bound (ELBO) for the individual-level likelihood p(Y'IY, X, T, T") a well-motivated and fitting objective for coun- terfactual generative modeling instead of the marginal-level (interventional) conditional likelihood p(Y|X, T) or even joint likelihood X. used by VAE-based prior works. To emphasize, the contribution here is not to propose another model variant under the conditional VAE formu- lation such as conditional HVAEs diffusion models (Sanchez & Tsaftaris 2022; Monteiro et al. 2023: Ribeiro et al. 2023), but to fundamentally change the VAE formulation. We call this framework Variational Causal Inference (VCI). An explanation to why the conditional VAE formulation (including diffusion models) is ill-suited to counterfactual generative modeling and a straightforward comparison between VAE's formulation and ours can be found in Appendix C 2) Our proposed optimization scheme allows us to conduct counterfactual supervision end-to-end during training alongside self supervuperi,,,n even without any counterfactual observations, which has not been done in prior works to the best of our knowledge. 3) This workflow of constructing and supervising counterfactual outcomes during training presents us the unique opportunity to naturally enforce disentangled exogenous noise abduction through distribution alignment on matching pairs (Shu et al. 2019 Locatello et al. 2020 Brehmer et al. 2022) without having paired data, which has not been done in counterfactual generative modeling, to the best of our knowledge. 4) We further propose a robust estimation scheme for high-dimensional marginal causal parameters leveraging this individual- counterfactual construction framework. No prior work in counterfactual generative modeling conducts such asymptotically efficient marginal estimation upon acquiring individual predictions, to the best of our knowledge. 5) In experiments, our proposed method is evaluated on datasets with vector outcomes- single cell perturbation datasets, as well as datasets with image outcomes facial imaging and handwritten digits datasets, and compared to state-of-the-art models in the two domains. The results show that ours outperformed state-of-the-ae notable margins. Related work and comparative analysis can be found in Appendix
in both domains with
PROPOSED METHOD
2.1 FORMULATION AND INTUITION
Let Y : Ω + (J) Ey) be the outcome, X Ω- + (x,Ex) be the covariates, T Ω + (T, E)) be the treatments and Z Ω + (Z, Σz) be a latent feature vector on probability space (Ω, E, P). Suppose the causal relations between random variables (and random vectors) follow structural causal model (SCM) (Pearl, 1995) defined in Figure where we adopt twin networks (Balke & Pearl 1994) to formulate counterfactuals Y' and 7" as separate variables apart from Y and 7 having
2

---

Published as a conference paper at ICLR 2025
Y
Z
Z
Y
Y
Y
Ux
T"
T"
T
X
X
T
(a) Bayesian network
(b) Counterfactual view of Bayesian network
Figure 1: The causal diagram where variables are generated from the following SCM: X = [x(Ux): T= = 1=ff((X fr UT) T fr Ui7) where Ur. U1. are i..d..; Z = (x, Uy); Y = fy (Z, Y' = = (y) where Ey, Ey are i..d.,, independent of all Us, and constants w.p. 1 if the consistency assumption is assumed. Endogenous variables follow the Bayesian network on the left under the ignorability assumption, with permissibly one additional edge from either Z T to X if Uy or UT is dependent of Ux. We call O = (X, T, Y) observed variables, B (,T Y) sample variables (T" sampled for model training and evaluation), D (X, variables and W = (Z, X, T, T", Y, y") all variables. White nodes are observed, light grey nodes are assigned during training and inference, dark grey nodes are unobserved.
T, T, T", Y, Y") full data
conditional distribution p(Y", T Z, X) identical to that of its factual counterpart p(Y. 7 X). Different from Vlontzos et al. (2023). ) we introduce latent Z high-dimensional outcome settings as feature vector of Y such that random factor cy does not capture any more uncertainty that can be retrieved from its counterfactual counterpart, i.e, ey Y' and similarly ey 1 Y. The consistency and ignorability assumptions corresponding to this causal formulation can be found in Appendix Under the consistency assumption (Assumption @ and Remark 2 Ey = 0 and there is no more uncertainty in Y beyond Z, i.c. the shared random factor Uy fully captures the exogenous noise injected to Y and Y'. Hence, in this setting can be seen as an unobserved summary random vector of the covariates X and exogenous noise Uy- Note that the consistency assumption is not generally required for our subsequent variational inference, but is required for YáorTo to be formally defined as the counterfactual of Y under the three-layer causal hierarchy Pearl (2009). The ignorability assumption (Assumption and Remark [ is generally required, under which do(T" t') and T" = 1' result in the same outcome distribution on Y' conditioned on X or Z. See Appendix D for connection between our formulation and the traditional SCM formulation.
Semi-autoencodinng Under this formulation, Z has posterior distribution p(Z Y,T, X) given the observed variables; Y and Y' can be constructed by p(Y p(Y| T) and p(Y' T'T') respectively with latent features Z. Similar prior works in autoencoder (Vincent et al. 2008 Bengio al. 2014), we can estimate the latent recognition model (i.e. exogenous noise abduction model) and outcome construction model with deep neural network encoder qs and decoder pg. Given a sample b (z,t, t', y). while the reconstruction ye.s of is self-supervised, we can only assess the counterfactual construction Vh.6 under E' by looking its resemblance to similar individuals that indeed received treatment I'. Hence, naively, we may conduct counterfactual supervision during training with semi-autoencoding (SAE) loss function as follow
Le.s = L2 o. y) w.
(I)
where w is scaling coefficient and is the traditional covariate outcome model fit on the observed variables X.T.Y) (notice that p' = t') = p(Y| =t)). The intuition is that, if vo. is indeed one's outcome under I', then the likelihood of 1/.5 coming from the outcome distribution of individuals with the same attributes æ that factually received treatment t' should be high. In practice, we can fit ( T') in an end-to-end fashion using discriminator D(X, T", K, Y) on factual triplets (x, t, y) and counterfactual triplets (x, f...v.. vo.5) with the adversarial approach (Goodfellow et al. 2014). Specifically, use discriminator loss T.D (b. v6,0)) = (Dx, v) Log[I=D[z,,) Vo.g) [)]] to train model D, and use generator loss (p(Y")x,e') (v6.))))v6.))))) ) = log D(x, v6,6) in (b) (notice that p(X,T) = T")). This end-to-end approach prevents (po. qa) from exploiting pre-trained model. In cases where covariates are limited and discrete such as the
3

---

Published as a conference paper at ICLR 2025
Reconstractin
Counter Sccuaal supervisin
No.6
Ue
Outcome Construction Model
Outcome Construction Model
Latent Recognitoo
Recognition
1
Model
divergence
Figure 2: Model workflow of the Variational Causal Inference (VCI) framework. In a forward pass, the encoding model takes outcome as well as its treatments and covariates (if any) as inputs and attains latent feature 2; (2, and (z, where the counterfactual treatments are separately passed into the decoding model to attain reconstruction ye, and counterfactual construction vé for which reconstruction loss and counterfactual supervision loss evaluated; vl.. is then passed back into the encoding model along with t' and to attain again, which encouraged to match the latent feature x previously acquired from encoding the factual outcome y and treatment through the latent divergence term. The neural architectures of the latent recognition and outcome construction models for image generation tasks can be found in Appendix
single-cell perturbation datasets, p(Y| X, 1XX77 T) can simply be smoothed empirical outcome distribution under treatment 7 stratified covariates X. The SAE approach motivates the derivation of main theorem presented in the next section, and serves baseline comparison to our main algorithm in the ablation study of the experiments section.
2.2
VARIATIONAL CAUSAL INFERENCE
Now we present our main result by rigorously formulating the objective, providing probabilistic theoretical backing, and specifying an optimization scheme that reflects the intuition described in the previous section. Suppose we want optimize p(Y|Y, X, T, T") instead of the traditional outcome likelihood pY 7). The following theorem induces the evidence lower bound (ELBO) and thus provides a roadmap for stochastic optimization:
Theorem 1. Suppose a collection of random variables W follows a causal structure defined by the Bayesian network in Figure Then J(D) = loglp(Y X, T, T')] has the following variational lower bound:
J(D) 2 E Y. log (Y (p(([X, PY" (,T")
DKL (p( T. p(Z)'(
where Dlp |I
log log
Proof of Theorem D and proofs of theoretical results the following sections can be found in Appendix M Theorem suggests that, in order to maximize the counterfactual outcome likelihood P(Y'),, Y I,, I"),, we can maximize the following estimated ELBO parameterized with estimating models po and dø given each sample b = (x,,, 1', v):
Jo. (b) Egs(slu.t.a) log (ul2,, t)] + log (x,, t']] DRL (1o.,) t, x) q (0(2 t',
where Vé. = E ( (0 Y = ( ........ 2,,1 Y and can be estimated with the approaches described at the end of Section 2.. Note that p(y|2, t) does not impose gradient on (pe, qo) and is
thus omitted in Jo (b). The stochastic optimization of ELBO is hence conducted the objective Jpase (po, 9s) = Epa.. (b) where ~ Pdata x)) A figure demonstrating the workflow of this stochastic optimization scheme be found in Figure 2
As can be seen from Equation 3 the ELBO consists of the aadividual-specifie factual outcome likelihood (elp.ta) log(po(ye (vl 2, and the covariate-specific counterfactual outcome likelihood
4

---

Published as a conference paper at ICLR 2025
loglp(u6..o l2,, t")], echoing the intuition highlighted by Equation with additional divergence term -Dx.lgo(=ll, x, q0(2) v. ,',))) that regularizes the similarity across latent distributions under different potential scenarios. Note that this framework is fundamentally different from VAE formulations such as CVAE (Sohn al. 2015) or CEVAE (Louizos al. 2017), since the variational lower bound is derived directly from the causal objective and the KL-divergence term serves causal purposes that are entirely different from the prior latent bounding purpose of VAEs. An interpretation of this divergence term is given in the paragraph below. A comparison between VCI and conditional
VAE using standard VAE notations given in Appendix C detaching patterns for training stability is given in Appendix K An approach to conduct counterfactual supervision implicitly using traditional variational inference is given Appendix
summary of different optional gradient
F
Divergence Interpretation The divergence term D DKL (90 (2ly,,,,))) || 90(2ly6 t', x)] is the key of the framework and it serves two purposes. Firstly, it adds robustness to latent encoding by encouraging the encoding model to recognize individual features contained in both factual and counterfactual outcomes. In fact, these common features well disentangled from the treatment variables, which we will discuss in-depth in the next section. Secondly, it encourages the preservation of individuality in counterfactual outcome constructions. To see this, notice that if the counterfactual outcome construction was only penalized by the covariate-specifi likelihood loss logl Vo lzr t')], the counterfactual decoding model po(''=za, 1') could learn to completely discard the identity of individual subjects represented in latent features 24 once it detects a counterfactual treatment t' * in the inputs. Such behavior is regulated by the divergence term, since otherwise we would not be able to recover a latent distribution (z(2)) 116 x) close to qa (Zly, (,y,,,,)) 2).
Robust Marginal Estimation Similar to the augmented inverse propensity weighted (AIPW) estimator in traditional causal inference, the individual-level model predictions acquired by our framework can also aid the robust estimation of average treatment effect (p) = EppY do( for treatment level a of interest if such estimation is desired. Under the formulation in Figure [ we have the following robust estimator that is asymptotically efficient under some regularity conditions (Van
Der Laan & Rubin 2006):
I(tx a)
ē(tx|Ir)) yk (1 (1- I(ta- = =a) EEpo (Y I2k T' = =a..
(o)
(4)
ε(tx|2z))
where O% (Ik, Wx) p(z, t, y) is the vector of observed triplet of the k-th individual, 2 (N tk. 2x ), and è is an estimation of the propensity score that satisfies the positivity assumption (Robins 1986). See Appendix G for the derivation of this estimator. Notice that the regression adjustment term Ep, I'' ="=6 a] in ours has the ability to vary across different individuals within the same covariate group, which is not the case for that of the AIPW estimator in traditional causal formulation. Hence, it is meaningful to conduct the robust estimation of covariate-ppecific marginal treatment effect E(p) Eyp Ey, (x) for given covariate of interest, Appendix for details.
2.3 LATENT IDENTIHIABILITY AND EXOGENOUS NOISE DISENTANGLEMENT
In any deep probabilistic model, there is a certain degree of freedom and arbitrarity to the learnt latent distribution, even with latent restrictions. Naturally, one might question the latent identification of our variational causal framework: when identifying counterfactual Y' under T", does the learnt latent Z inevitably earry over some characteristics of the factual treatment in Y and T"? Ideally. we would like to aid counterfactual construction by attaining a clean disentanglement between the learnt latent and the observed treatment, such that the counterfactual generation does not preserve influence from the factual treatment. Such disentangled exogenous noise abduction fits the semantic of our generating process (Figure D. where Z and T' are not descendant of each other and Z 1 T|X, as well as the semantic of causal identification in the traditional formulation, where the identifiability of Yáo(r) relies on the ignorability assumption that exogenous noise does not contain unobserved confounders. In this section, we show that optimizing the VCI framework in fact grants such disentanglemert between Z and T under mild assumptions, leading to more in-depth understanding of the purpose and behavior of the latent divergence term in VCT's 's optimization scheme.
5

---

Published as a conference paper at ICLR 2025
Similar to Shu et al. (2019), we formally define and measure disentanglement in terms of oracle consistency and oracle restrictiveness, except that our definitions are generalized to the case of stochastic models and models with auxiliary attributes. For notation simplicity of the oracle operation on functions with auxiliary inputs, we use <(a; b) to denote >(c, b) where c = ((a) if Ç is a deterministic function and ((a) if Ç is probability measure. Denote the true data generating distribution as Po- The definitions of oracle consistency, oracle restrictiveness, and disentanglement are then given Appendix F
Consider our model class { X + [0, 1ر q. : y x T x X + [0, 1| } and :
let S == (.. V == TY. We regard T as part of the model outputs such that the oracle operations above are well-defined, i.e. let H. = {(P..0 q.) (s(s) = (Y x 6, (T))), (. x) = q. (, 1, x) x (T)) where d is the Dirac measure (7) = if else 0). Then we have the
following lemma:
Lemma 1. If q minimizes the latent divergence Equation 2 over Po(D), i.e.
Eo DEL (g(2/y, x, || <<//// , x)]]===
then q disentangles Z and T.
Lemma buulds direct connection between the latent divergence term and disentangled abduction of exogenous noise. Intuitively, it states that if the encoder encodes the same latent representation for any pair of potential outcomes of the same individual, then the latent must have been disentangled from the treatment. However, does optimizing the ELBO necessarily imply that the latent divergence term reaches minimum? We proceed to show that this can be implied under mild assumptions.
Proposition 1. With the true distribution po satisfying Assumption 2 in Appendix H and Assumption 3 in Appendix 1.2.2 2.2.2 if (p+:qe) maximizes the evidence lower bound over po (D) and any i.e.
(p+, q.) E argg maxlE {1Eg (slu.t,a) log Ip. (y) z, t)] + log lp(y' [x, t']
P.q.
-DKI. (=lv,,,))) || (.g.g.g.. (xle',
then q. disentangles Z and T.
Intuitively, Proposition states that as long as the treatments do not have any unknown side effect and do not completely wipe out different features in different individuals, the solution to maximizing the ELBO disentangles Z and T. In practice, we can formulate model class as deterministic functions (i.e. degenerate distributions) plus noise similar to prior work, and encoder noise standard normal:
{p. g. ( t) == òg (, = (Y - (2,,
{q. e, (Zly, L, .r) = (is.a.)) (. Ey v.t.z) 2,, (,,) = 10)
and learn model e,, 9. and c, during training. The following proposition states that disentanglement under the estimating model's own oracle can be achieved under this setting by optimizing the VCI objective in real-world optimizations, where neither the true generating mechanisms nor the true factual-counterfaatual pairs are available.
Proposition 2. Given dataser {Ir, tk, H let Paata be its empirical distribution and the model class in Equation and Equation 8 satisfy the following uniqueness conditions:
1) Va E T, we have (Y + (Y +
2) Vx,t,y,t Pdata (I, t, v) (t' (t'(z)
> 0 and v. = Eg
we have
Pdata (x. t,, v/) = 0 if (z, (z,t,u v) # (x,tt, V.Y.))
If (pe,q)) maximizes the VCI objective in Equation 3 Jover Pdata and any p. i.e. (p., (-) E arg max Jpans (p., 9.)
P.
then (g., q.) disentangles Z and T under J, C.Pdata-
We note that uniqueness condition 2) in Proposition 2 reasonable in high-dimensional outcome settings: taking the facial imaging dataset as an example, wouldn't make sense for a model's constructed image for certain individual to be exact match of the factual outcome of another individual in the dataset, as long as the dataset does not contain multiple images of the same individual in the same environment.

---

Published as a conference paper at ICLR 2025
Corollary 1. Under the same settings and conditions of Proposition 2 (ă., êx) disentangles Z and under Jy exPaata
Corollary| 1 states that the optimal model is guaranteed to disentangle Z and 7 under the constructed latent distribution in inference mode without noise injection. In summary, the disentanglement results in our framework rely on the simple fact that the latent divergence term is established on individual level, and does not necessarily compete with outcome supervisions under mild assumptions. Prior works in VAEs have proposed bounding latent on conditional marginal prior (Sohn et al., 2015 Khemakhem et al. 2020), and VCI can be seen as a further relaxed form of latent restriction in the sense that the latent for any given subject need not be restricted to any prior distribution we only require the latent distributions of a subject's s different potential outcomes to match each other. More insights on how the latent divergence term helps the training of VCI in practice can be found in the ablation study in the experiments section.
3 EXPERIMENTS
We present experiment results of our framework on two datasets with vector outcomes (sci-Plex dataset from Srivatsan al. (2020) (Sciplex) and the CRISPRa dataset from Schmidt et al. (2022) (Marson)) and two datasets with image outcomes: Morpho-MNIST (Castro et al. (211)) and CelebA- HQ (Karaas et al., 2017). Results from the former are compared to state-of of-the--rt models in single-cell perturbation prediction and results from the latter are compared to tfte-off-hhear models counterfactual image generation. In both cases, our model exhibited superior performance against benchmarks. Details of model and dataset settings for each experiment can be found in Appendix M
3.1 SINGLE-CELL PERTURBATION DATASETS
Same as Lotfollahi et al. (021), we evaluate our model and benchmarks on a widely accepted and biologically meaningful metric the R² (coefficient of determination) of the average prediction against the true average from the out-if-distributtin (OOD) set on all genes and differentially- expressed (DE) genes. Definitions of OOD set and DE genes can be found in Appendix M. Results over five independent runs are shown in Table !
Table 1: R of OOD predictions on single-cell perturbation datasets. AE is naive baseline adapting Autoencoder to counterfactual generation, see Appendix L.I GANITE is GANITE's counterfactual block adapting high-dimensional outcome plus multi-level treatment (see Appendix
Sciplex
Marson
all genes
DE genes
all genes
DE genes
AE
0.740 + 0.043
0.421 + 0.021
0.804 + 0.020
0.448 + 0.009
CEVAE Louizos et al.
2017
0.760 = 0.019
0.436 0.014
0.795 + 0.014
0.424 + 0.015
GANITE Yoon et 2018)
0.751 0.013
0.417
0.014
0.795 + 0.017
0.443 + 0.025
CPA Lotfollahi et al.
2021
0.836 = 0.002
0.474 0.014
0.876
+ 0.005
0.549 = 0.019
VCI
0.832 + 0.008
0.496 0.011
0.891 + 0.007
0.658 + 0.040
For each perturbation of each covariate level (e.g. each cell type of each donor) in the OOD set, the R² (coefficient of determination) is computed with the average outcome predictions for all genes and DE genes using samples from the validation set against the true empirical average over samples from the OOD set. The average Ri2 over all perturbations of all covariate levels is then calculated as the evaluation metrie R². In these experiments, our variational Bayesian causal inference framework excelled state-of -o---o--trtttt models in both experiments, with the largest fractional improvement on DE genes which are most causally affected by the perturbations. Results of marginal estimation using the robust estimator can be found in Appendix N.I.
3.2 MORPHO-MNIST
In an ideal world, we would like to evaluate model performance against benchmarks simply by measuring the error between the counterfactual prediction and the counterfactual truth for each
7

---

Published as a conference paper at ICLR 2025
Table 2:
The errors between counterfactual predictions and counterfactual truths on Morpho-MNIST. Images are scaled to [0. 1] in evaluations. On top of randomly sampling counterfactual treatment, we also randomly sample the ratio of modification, which could be any value such that the thickness and intensity of the counterfactual truth is in range. The standard error of runs and the violin plot of errors are reported in Appendix N.5
Image MSF (-10 ?)
Thickness (th) MAR
Intensity (in) MAR (-10-)
do(th)
do(in)
mix
do(th)
do(in)
mix
do(th)
dofin)
mix
DEAR
Shen el al.
2022
6,93
3.21
5.08
0.82
0.74
0.85
6.23
3.39
4.61
Dif--SCM
(Suncee & satturris 2022]
7.72
4,90
6,30
0.70
0.76
0.73
3.48
2.17
2,80
CHVAE
Monteiro cl 2023
4.92
4.81
4.81
0.61
0.66
0.62
3.23
1888
2.81
CHVAE
Monteiro
2023
6.93
6.19
6,40
0.71
1.34
0.97
2.37
1.31
1.89
MED
Ribeuro el al
2023
2.76
0.54
1.65
0.54
0.35
0.46
1.17
0.43
0.78
MED
Kiheiro el
2023
2.20
0.61
1.39
0.31
0.29
0.33
0.44
0.34
0.40
SAE
(Section
0.67
'0220
00.47
0.32
0.29
0.52
0.41
VCI (Section
0.42
0.13
0.36
0.20
0.14
0.22
0.42
0.22
0.33
individual. When the counterfactual truth is attainable, there is no need to resort to approximation metries such as Monteiro et al.. (2023): Melistas et al. (2024). For this reason, we specifically chose the Morpho-MNIST (Castro et 2019)) dataset to present our main evaluation results because the counterfactual truth can be directly computed on this dataset based on the intervention (modifying thickness or modifying intensity of hand-written digits) even with the existence of exogenous noise. Similar to Ribeiro et al. (2023), we evaluate counterfactual constructions under single modifications as well mix of modifications. and compare them to statc-of-the-arr models high-fidelity image counterfactual generation. Contrary to Ribeiro et al. (203)), the magnitude of modification is randomly sampled, which makes the task significantly harder. The results shown in Table (standard error in Table(8 which demonstrate that ours beat sate-of-thhe-rte by a wide margin. To understand why conditional diffusion model such as Diff-SCM (Sanchez & Tsaftaris 2022) does not perform well in counterfactual generative modeling, see Appendix C. Note that Ribeiro et al. (2023) only evaluates mean absolute error (MAE) on the thickness and intensity of counterfactual constructions, which took treatment characteristics into account but completely left out how much individuality of the original images has been preserved, whereas mean squared error (MSE) on the image is comprehensive metric that takes both factors into account and should be the primary evaluation metric when the counterfactual truth is available. A model inspection using axiomatic soundness metries (Monteiro et al., 2023) can be found in Appendix N.4
To investigate the impact of counterfactual supervision and latent divergence terms optimizations, we present an ablation study below to compare VCI to SAE and hierarchical autoencoder (HAE) with the same neural architecture.
do(thickkesss)
do(intensity)
9.0
1.8
5.4
HAE
1.6
4.8
SAF
7.5
1.4
4.2
VCI
MSE (10
6.0
1.2
3.6
1.0
3.0
4.5
0.8
2.4
0.6
1.8
0.4
0.2
0.0
50
100
150
200
50
100
150
200
50
100
150
200
Figure 3: Ablation Study: the error of counterfactual prediction across epochs during the training of HAE (VCI without counterfactual supervision and latent divergence). SAE (VCI without latent divergence) and VCI over five independent runs. Note that VCI with latent divergence but without counterfactual supervision does not make logical sense, but for the completeness of the ablation study, present the results for such setting in Appendix N.2
incorporating counterfactual supervision alone could greatly accelerate optimization in the initial epochs of training, and achieves an ideal model state in the early stage. This is the same phenomenon observed in the optimization single-cell perturbation datasets with
As can be seen in Figure 3
8

---

Published as a conference paper at ICLR 2025
respect R However, although both SAE and VCI beat HAE convincingly in terms of best result, SAE is largely unstable in the long-run without the proposed latent divergence restriction, and the full VCI optimization scheme consistently achieves better performance in terms of best result and final result. To further help the readers understand how and why the latent divergence term stabilizes VCI training, we include discussion below along with some illustrations do(intensity) in Figure
9
9
9
9 9
9
9
9
9
9
9
9
S
S
5
5
5
5
5
5
5
5
5
7
7
7
7
/
7
7
7
7
7
7
7
3
3
3
3
3
3
3
3
3
3
3
3
6
6
6
6
6
6
6
6
6
6
6
6
SAE
(h) VCI
Figure 4: Illustration of the latent divergence term's long-term impact. Samples are drawn from the last epoch of SAE and VCI on do(intensity). The left most image of each set the original image. More sampled results from VCI on do(thickness) and (ointessity) can be found in Figure 12 Results from intervening digit i.e. do(digit) can be found in Figure 13
In a GAN training scheme, the discriminator a neural model that could often struggle with out-of- distribution evaluations. As be seen from the generated samples by SAE in Figure when the intervention labels are close to the original label, the counterfactual generations are reasonable. For intervention labels that are significantly different from the original, the images become nonsensical. This is because the combination of an intervention label that is far from the original and an image similar to the original highly -utto--distribution and rarely present in the observed dataset. In other words, the discriminator has never encountered an image with the same digit and style as the original having such extreme intensity variations, so whatever the generator produces, the discriminator cannot tell if it is real or fake due to the lack of reference. This results in the generator exploiting the neural discriminator in these -iitti--dittitttionn scenarios the long run. However, the latent divergence term prevents the generator from producing these nonsensical images because, otherwise, the encoder would not be able to recover a latent distribution close to that the original image. An evaluation of latent disentanglement under the latent divergence term can be found in Appendix N.3
3.3 CELEBA-HQ
For any real-world dataset where the counterfactual truth not available, there is not definitive metric to evaluate how good the counterfactual constructions really are. For that reason, facial imaging datasets are prevailing benchmarks for examining counterfactual goodness because even in the lack of quantitative metric, human can judge the quality of counterfactuals without the need for any domain knowledge. Same as Monteiro et al. (2023), we use the CelebA-HQ (Karras et al.] (017) dataset 64 resolution to evaluate the model's capability of counterfactual constructions on two factors - smiling and glasses. Some sampled results from our model are shown in Figure 5
While MSE of the counterfactual construction cannot be measured, readers can empirically compare ours to prior work such as CHVAE (Monteiro et al. (023) which serves as the backbone for the state-of-the-ard model in high fidelity counterfactual image generation MED (Ribeiro et al.) 2023). Aside from compiling counterfactual credibility across all generation tasks, ours show strong capability of disentangling factual treatment from observed images in counterfactual construction, which can be particularly observed in the task of "removing glasses" as shown in Figure [dd] Note that contrary to prior works, we construct and supervise counterfactual outcomes during training, and the reconstruction and counterfactual construction share the same decoding mechanism. Hence, "reconstruction" is really counterfactual construction with the factual label. This can be most clearly observed in the third set of comparison in line of Figure 5al where the original image has label of "not smiling" yet the person is really in between smiling and not smiling. In this case, the "reconstruction" really counterfactual construction with treatment abiding to the original label.

---

Published as a conference paper at ICLR 2025
(a) Add smiling. Left: original, middle: not smiling (reconstruction) right: smiling (counterfactual)
Remove smiling. Left: original, middle: smiling (counterfactual), right: smiling (reconstruction).
(c) Add glasses. Left: original, middle: glasses (reconstruction), right: with glasses (counterfactual)
(d) Remove glasses. Left: original, middle: glasses (counterfactual) right: with glasses (reconstruction).
Figure 5: Results on the test set of CelebA-HQ.
Lastly, we include discussion here regarding counterfactual fairness. Note that adding glasses is relatively easy engineering task, but relatively hard deep learning task on CelebA-HQ due to the small sample size of images with glasses. As can be seen from Figure 5c our model learned not only attach glasses at the right position, but also attach glasses according to a person's 's style and context. However, readers may naturally wonder if this could induce any fairness issue, such as attaching certain type of glasses more frequently on certain demographic. We note that our framework can reduce spurious unintended relations between given factors, as long as practitioners measure these factors during data collection and include them as observed labels for the composed datasets, since the learnt latent representations of our framework are encouraged to be disentangled from the observed factors. In the case of CelebA, style of glasses is unmeasured, and the model cannot reasonably tell if attaching different types of glasses on different demographics unintended or not.
CONCLUSION
In this work, we introduced a variational Bayesian causal inference framework for estimating high- dimensional counterfactual outcomes, as well as consequent robust marginal estimators. With this framework, treatment characteristics and individuality in the predicted outcomes can be explicitly balanced and optimized, and learnt latent representations are disentangled from the treatments. As for limitations, the theoretical results of this work are established upon the common causal assumption that there exists no unobserved confounders, through which causal effects are identifiable. Besides, when multiple treatments are concerned, this work does not explicitly address the situation where a causal relation exists between treatments. Thus, a straightforward extension of this work could be incorporate explicit modeling of the causal relation among treatment variables and update the descendants of given counterfactual treatment, similar to Pawlowski et al.) (2020), before feeding them into counterfactual outcome construction framework.
10

---

Published as a conference paper at ICLR 2025
AUTHOR CONTRIBUTIONS
Yulun Wu developed the proposed framework, conducted experiments on single-cell perturbation datasets, Norpho-MNIST and CelebA-HQ, and wrote the manuscript; Louie McConnell conducted experiments Morpho-MNIST; Claudia Iriondo reviewed related work and edited the manuscript.
ACKNOWLEDGMENTS
The authors thank Cian Eastwood, Toru Shirakawa, Picter Abbeel, Jason Hartford, Dominik Janzing, Ivana Malenica, Ahmed M. Alaa, and David E. Heckerman for the insightful discussions. The authors thank Layne C. Price, Zichen Wang. Robert A. Barton, Vassilis N. loannidis, and George Karypis for their feedback in the early development.
REFERENCES
Alexander Balke and Judea Pearl. Probabilistic evaluation of counterfactual queries. Proceedings of the AAAI conference artificial intelligence, volume 12, 1994.
Yoshua Bengio, Eric Laufer, Guillaume Alain, and Jason Yosinski. Deep generative stochastic networks trainable by backprop. In International Conference Machine Learning, pp. 216-2262-44 234.. PMLR, 2014.
Johann Brehmer, Pim De Haan, Phillip Lippe, and Taco S Cohen. Weakly supervised causal representation learning. Advances in Neural Information Processing Systems, 35:38319-38331. 2022.
Daniel C Castro, Jeremy Tan, Bernhard Kainz, Ender Konukoglu, and Ben Glocker. Morpho-mnist: quantitative assessment and diagnostics for representation learning. Journal of Machine Learning Research, 20(178):1-22 2019.
Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Info- gan: Interpretable representation learning by information maximizing generative adversarial nets. Advances in neural information processing systems, 29, 2016.
Rewon Child. Very deep vaes generalize autoregressive models and can outperform them on images. arXiv preprint arXiv:2011. 2020.
Saloni Dash, Vineeth N Balasubramanian, and Amit Sharma. Evaluating and mitigating bias in image classifiers: A causal perspective using counterfactuals. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 914-924, 2022.
Atray Dixit, Oren Parnas, Biyu Li, Jenny Chen, Charles P Fulco, Livnat Jerby-Arnon, Nemanja D Marjanovic, Danielle Dionne, Tyler Burks, Raktima Raychowdhury, et al. Perturb-seq: dissecting molecular circuits with scalable single-cell raa profiling of pooled genetic screens. cell, 167(7): 1853-1866, 2016.
Ruili Feng, Jie Xiao, Kecheng Zheng, Deli Zhao, Jingren Zhou, Qibin Sun, and Zheng-Jun Zha. Principled knowledge extrapolation with gans. In International Conference on Machine Learning, pp. 6447-6464. PMLR, 2022.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances neural information processing systems, 27, 2014.
Luigi Gresele, Julius Von Küigelgen, Vincent Stimper, Bernhard Schölkopf, and Michel Besserve. Independent mechanism analysis, new concept? Advances in neural information processing systems, 34:28233-290 28248, 2021.
Hermanni Hälvä, Sylvain Le Corff, Lue Lehéricy, Jonathan So. Yongjie Zhu, Elisabeth Gassiat, and Aapo Hyvarinen. Disentangling identifiable features from noisy data with structured nonlinear ica. Advances in Newral Information Processing Systems, 14:1124-1633, 2021.
11

---

Published as a conference paper at ICLR 2025
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6833 2020.
Aapo Hyvarinen, Hiroaki Sasaki, and Richard Turner. Nonlinear ica using auxiliary variables and generalized contrastive learning. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 859-868. PMLR, 2019.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv:1710.10196. 2017.
Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoencoders
and nonlinear ica: A unifying framework. In International Conference on Artificial Intelligence and Statistics, pp. 2207-2217. PMLR, 2020.
Hyemi Kim, Seungjae Shin, JoonHo Jang. Kyungwoo Song, Weonyoung Joo, Wanmo Kang, and
II-Chul Moon. Counterfactual fairness with disentangled causal effect variational autoencoder. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 8128-8136, 2021.
Diederik Kingma and Max Welling.
Auto-encoding variational bayes.
arXiv preprint
arXiv:1312.61 14, 2013.
Murat Kocaoglu, Christopher Snyder, Alexandros G Dimakis, and Sriram Vishwanath. Causal- gan: Learning causal implicit generative models with adversarial training. arXiv preprint 1oXXv:1709,02033,, 2017.
Aneesh Komanduri, Yongkai Wu, Feng Chen, and Xintao Wu. Learning causally disentangled repre- sentations via the principle of independent causal mechanisms. arXiv preprint arXiv:2306.02/13 2023.
Aneesh Komanduri, Chen Zhao, Feng Chen, and Xintao Wu. Causal diffusion autoencoders: Toward counterfactual generation via diffusion probabilistic models. arXiv preprint arXiv:2404.17735, 2024.
Sébastien Lachapelle, Pau Rodriguez. Yash Sharma, Katie Everett, Rémi Le Priol, Alexandre Lacoste, and Simon Lacoste-Julien. Disentanglement via mechanism sparsity regularization: A new principle for nonlinear ica. In Conference on Causal Learning and Reasoning, pp. 428-484- PMLR, 2022.
Jonathan Levy. Tutorial: Deriving the efficient influence curve for large models. arXiv preprint arXiv:1903.01706 2019.
Timothy Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint (50.:::50900977, 2015.
Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly-supervised disentanglement without compromises. In International Conference on Machine Learning, pp. 6348-6359, PMLR, 2020.
Mohammad Lotfollahi, Anna Klimovskaia Susmelj. Carlo De Donno, Yuge Ji, Ignacio L Ibarra, Alexander Wolf, Nafissa Yakubova, Fabian Theis, and David Lopez-Paz. Learning interpretable cellular responses to complex perturbations in throuuhut screens. bioRxiv, 2021.
Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep latent-variable models. Advances in neural information processing systems, 30, 2017.
Thomas Melistas, Nikos Spyrou, Nefeli Gkouti, Pedro Sanchez, Athanasios Vlontzos, Yannis Pana- gakis, Giorgos Papanastasiou, and Sotirios A Tsaftaris. Benchmarking counterfactual image generation. arXiv preprint arXiv:2403.20223 2024.
Miguel Monteiro, Fabio De Sousa Ribeiro. Nick Pawlowski. Daniel C Castro, and Ben Glocker. Measuring axiomatic soundness of counterfactual image models. arXiv preprint arX1v:2303.01274, 2023.
12

---

Published as a conference paper at ICLR 2025
Thomas M Norman, Max A Horlbeck, Joseph M Replogle, Alex Y Ge, Albert Xu, Marco Jost, Luke A Gilbert, and Jonathan S Weissman. Exploring genetic interaction manifolds constructed from rich single-cell phenotypes. Science, 365(6455):786-795 2019.
Nick Pawlowski, Daniel Coelho Castro, and Ben Glocker. Deep structural causal models for tractable counterfactual inference. Advances in Neural Information Processing Systems, 33: 857-869, 2020.
Judea Pearl. Probabilistic reasoning in intelligent systems: networks of plausible inference. Morgan kaufmann, 1988.
Judea Pearl. Causal diagrams for empirical research. Biometrika, (2(4):669-688 1995.
Judea Pearl. Causality. Cambridge university press, 2009.
Fabio De Sousa Ribeiro, Tian Xia, Miguel Monteiro, Nick Pawlowski, and Ben Glocker, High fidelity image counterfactuals with probabilistic causal models. arXiv preprint 2xXiv23306.15764 2023.
James Robins. A new approach to causal inference mortality studies with a sustained exposure period-applicen application to control of the healthy worker survivor effect. Mathematical modelling, 7 (9-12):1393-151) 1512, 1986.
Paul R Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55, 1983.
Pedro Sanchez and Sotirios A Tsaftaris. Diffusion causal models for counterfactual estimation. arXiv preprint arXiv:2202. 10166, 2022.
Axel Sauer and Andreas Geiger. arXiv:2101.06046, 2021.
Counterfactual generative networks.
arXiv preprint
Ralf Schmidt, Zachary Steinhart, Madeline Layeghi, Jacob W Freimer, Raymund Bueno, Vinh Q Nguyen, Franziska Blaeschke, Chun Jimmie Ye, and Alexander Marson. Crispr activation and interference screens decode stimulation responses in primary human t cells. Science, 375(6) eabj4008, 2022.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06344 2017.
Xinwei Shen, Furui Liu, Hanze Dong, Qing Lian, Zhitang Chen, and Tong Zhang. Weakly supervised disentangled generative causal representation learning. The Journal of Machine Learning Research, 23(1):1099411048; 2022.
Rui Shu, Yining Chen, Abhishek Kumar, Stefano Ermon, and Ben Poole. Weakly supervised disentanglement with guarantees. arXiv preprint arXiv:/910.09772, 2019.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in newral information processing systems, 28, 2015.
Peter Spirtes, Clark N Glymour, and Richard Scheines. Causation, prediction, and search. MIT press, 2000.
Sanjay R Srivatsan, José I. McFalline-Figuerra, Vijay Ramani, Lauren Saunders, Junyue Cao, Jonathan Packer, Hannah Pliner, Dana Jackson, Riza M Daza, Lena Christiansen, et al. Massively multiplex chemical transcriptomics at single-cell resolution. Science, 367(6473) 2020.
Arash Vahdat and Jan Kautz. Nvac: A deep hierarchical variational autoencoder. Advances in neural information processing systems, 33:1
2020.
Mark J Van Der Laan and Daniel Rubin. Targeted maximum likelihood learning. The international journal of biostatistics, 2(1), 2006.
Aad W Van der Vaart. Asymptotic statistics, volume Cambridge university press, 2000.
13

---

Published as a conference paper at ICLR 2025
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q- learning. Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016.
Pascal Vincent, Hugo Larochelle, Yoshua Bengio, and Pierre-Antoinn Manzagol. Extracting and composing robust features with denoising autoencoders. In Proceedings of the 25th international conference on Machine learning, pp. 1996-1103, 2008.
Athanasios Vlontzos, Bernhard Kainz, and Ciarán M Gilligan-Lee. Estimating categorical counter- factuals via deep twin networks. Nature Machine Intelligence, 5(2):159- 11011199968,,,, 2023.
Julius von Kügelgen, Michel Besserve, Wendong Liang, Luigi Gresele, Armin Kekić, Elias Barein- boim, David M Blei, and Bernhard Schölkopf. Nonparametric identifiability of causal representa- tions from unknown interventions. arXiv preprint arXiv:2236..005: 2023.
Mengyue Yang, Furui Liu, Zhitang Chen, Xinwei Shen, Jianye Hao, and Jun Wang. Causalvae: Disentangled representation learning via neural structural causal models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9593-9602, 2021.
Jinsung Yoon, James Jordon, and Mihaela Van Der Schaar. Ganite: Estimation of individualized treatment effects using generative adversarial nets. In International Conference on Learning Representations, 2018.
14

---

Published as a conference paper at ICLR 2025
A INTRODUCTION ON LATENT DISENTANGLEMENT
The study of latent disentanglement in deep causal modeling has mainly focused on two areas. Most prior works lie in the first that is causal representation learning, in which researchers seek to learn latent variables corresponding to true causal factors and identify the structure of their causal relations YYng et al. 2021] Hälvá al. 2021 Gresele et al. 2021 Lachapelle al. 2022, Komanduri et al.) 2023). predominantly leveraging recent advances in non- independent component analysis (ICA) (Hyvarinen al., 2019; Khemakhem et al., 2020) and identify causal structure up to Markoy equivalence (Spirtes et al. 2000) or graph isomorphism under rather heavy assumptions. The second area and the area this work concerned
counterfactual generative modeling, in which the focus to conduct disentangled exogenous noise abduction that aids counterfactual outcome generation (Lotfollahi et al. 2021 Shen et al., (022). In this context, the goal is to acquire latent representations disentangled from the observed treatment to aid the correct identification of the causal effect of counterfactual treatments.
B RELATED WORK
Key prior works deep causal generative modeling fall primarily into three categories: variational-
based methods CEVAE (Louizos et al. (017). DeepSCM (Pawlowski et al. 2020), Diff-SCM
(Sanchez & Tsaftaris 2022), CHVAE (Monteiro et al. 2023). MED (Ribeiro et al. 2023); adversarial-
based methods CausalGAN (Kocaoglu et al. (017). GANITE (Yoon et al.) 2018); hybrid method
DEAR (Shen et al. 2022) as well as our framework VCI. comparative analysis of ours against these related works can be found in Table 3
Table 3: A comparative analysis of related work. "Variational Objective Distribution" describes the objective that the ELBOs are derived for in variational-baat method: "Joint" indicates joint distributions such as p(y, xr or p(y, 2) (note that different works could have different notations, here we use y, x, t, to represent outcome, covariates, treatment, latent respectively). and "Inter- ventional" indicates conditional outcome distribution p(ylær, or p(ylt). "Marginal Distribution Alignment" describes whether the stochastic optimization objective involves mechanisms that match distribution of outcome constructions with some learnt marginal outcome distribution.
Variational Objective Distribution
Hierarchical Model Structure
Marginul Distribution
Causal Discovery
End-to-Bind Counterfactual Supervision
Latent Disentanglement
Type
Alignment
CEVAE
Variational
Joint
X
X
X
CausalGAN
Adversarial
X
X
X
X
GANITE
Adversarial
X
X
X
X
CPA
Advensarial
X
X
X
V
DEAR
Variational
Adversarial
Joint
X
X
X
DeepSCM
Variational
Interventional
X
X
X
X
X
Diff-SCM
Variational
Interventional
X
X
X
X
CHVAE
Variational
Interventional
X
X
X
X
MED
Variational
Interventional
X
X
Variational
VCI
Counterfactual
Adversarial
*DAR learns restrictive representation for latent (2 is disentangled from treatment f when encoding f) but does not learn consistent representation for (t is disentangled from when encoding 2).
CEVAE uses variational autoencoders (VAEs) infer latent confounders and estimate individual treatment effects (TTE). DeepSCM integrates deep learning into SCMs using normalizing flows and variational inference for counterfactual inference. DifT-SCM leverages generative diffusion models, using forward and reverse diffusion processes guided by anti-causal predictors to generate counterfactuals. CHVAE extends VDVAE to a conditional model, performing abduction by encoding images and parent attributes, and includes penalty for counterfactual conditioning. MED uses hierarchical VAEs with deep causal mechanisms, inferring exogenous noise by conditioning on observed data and parent variables. GANITE uses counterfactual generator create proxies for
15

---

Published as a conference paper at ICLR 2025
x'
VAE
VCI
(c) VCI (unfolded)
Figure 6: The data generating formulation of conditional generative modeling under VAE formula- tion and counterfactual generative modeling under VCI formulation using standard VAE notations. p(z|2, c) = p(21|z,, c') for d' White nodes are observed, light grey nodes are assigned during training and inference, dark grey nodes are unobserved.
Table 4: The objective and ELBO for conditional generative modeling under the VAE formulation and counterfactual generative modeling under the VCI formulation, using standard VAE notations. can be seen, the difference in motivations between the two objectives manifest themselves in the derived ELBOs some meaningful and intuitive ways: VAE does not have the motivation to maximize counterfactual outcome likelihood, hence no such term log (2(z',, to conduct supervision of counterfactuals; VCI does not have the motivation to denoise or to generate samples from pure noise duing inference time, hence such term DEL [xL[g, (2|22, c) po (=|c)] bound latent distribution some marginal prior.
VAE
VCI
Objective
p(a|c)
p(a' (ar c,e")
NAE (z)
c)= Esttalnn log o))
1 (a,0,e)
= Balltane) log (pu(orlz. +
log (i
ELBO
DeL.
co( c)
Pa (2(o)
DxL
[o(=]a, e)
unobserved outcomes and an ITE generator to estimate potential outcomes. DEAR combines VAE-like generative loss with a supervised loss in bidirectional generative model, trained with GAN algorithm requiring knowledge of the causal graph structure and extra labels for latent factors.
As illustrated in Section and Appendix the primary contribution of our framework VCI is the revolution of VAE formulation which leads brand new stochastic optimization scheme that optimizes the ELBO of an actual counterfactual objective. This brings about end-to-end counterfactual supervision and exogenous noise disentanglement as presented Table
C COMPARISON OF CONDITIONAL VAE AND VCI
For readers that are more familiar with deep generative modeling than causal inference, we provide a straightforward comparison between conditional VAE and VCI (fundamentally a comparison between conditional generative modeling and counterfactual generative modeling) in this section using VAE's notations. For data x and condition c, the differences in formulations and objectives Figure and Table 4 In causal inference, counterfactual outcome is an individual-level concept for a given individual that we observed outcome/data æ under treatment/condition c, what would their outcome have been if they had received treatment/condition c' instead? This is a "would have" question, meaning that we seek to find out what the alternative outcome z' would be in a parallel world where everything in the state of the universe remained the same, except that the treatment c' (and its impact) was different. With this mind, modeling formulation is unorthodox for counterfactual generative modeling the ELBO serves to optimize the marginal- (rung 3) (Pearl 1995). and the learnt model generates samples towards the marginal-level distribution p(x(c) during inference time, and the question being asked here is "what will an outcome r be under condition c?". This is "will" question it is generating any outcome under condition not given a specific individual, not given specific state of the universe. Therefore, contrary to prior works in HVAEs and diffusion models (Sanchez & Tsattaris
shown in
hard to see why the conditional generative
likelihood p(x|c), which is interventional (rung 2) and not counterfactual
2022 Monteiro al. 2023 Ribeiro al.
16

---

Published as a conference paper at ICLR 2025
2023) which focus on model design adaptations to make the conditional VAE formulation work in counterfactual generative modeling, we first derive the orthodox formulation and objective for counterfactual generative modeling, then make the necessary model designs afterwards based on the accordingly derived ELBO.
C.I DIFFUSION MODELS
Due to the popularity of the diffusion models (Ho al. 2020), we feel the necessity to include discussion here specifically about its compatibility with counterfactual generative modeling. It is very important to note that, although diffusion models are the state-of-the-en art in generative and conditional generative modeling, it has not been shown that it has better capability in counterfactual generative modeling than ordinary HVAEs with learnable encoder (Monteiro et al., 2023 Ribeiro et al., 2023). and the reason is straightforward once we truly understand the fundamental incompatibility between the goal of counterfactual inference and the diffusion mechanism. Counterfactual inference entails abducting and preserving the exogenous noise in the original outcome; diffusion models, the other hand, entail diffusing and denoising the original image it is a mechanism that inherently does not respect the consistency assumption (Assumption (2) and fundamentally contradicts exogenous noise abduction. Even if certain prior work such as (anchez & Tsaftaris (2022) uses the learnt noise model to replace random noise in inference time, it is still performing the illogical operation of attempting to construct counterfactual from a noisy version of the original outcome, in which the exogenous noise has already been discarded. And this is on top of the fact that diffusion models, as models under the VAE formulation, are working towards the wrong interventional objective as discussed above. Prior works that attempted at utilizing diffusion models in counterfactual generative modeling (Dash et al.] [2022] Sanchez & Tsattari 2022, Komanduri et al.] 2024) clearly exhibited this flaw: as can seen from the CelebA results in Dash et al. (2022) and the MNIST results in Komanduri et al. (2024), exogenous noise/individuality has been largely discarded in counterfactual generation because the diffusion model is doing what is intended to do: diffusing exogenous and
generating samples that fit sceminglessly into the conditional/margina likelihood p(z]cc;; (anchez (anchez & Tsaataris (2022) on the other hand, hardly showed any evidence that the model is capable of preserving individuality- -MIITT was the only benchmark in their experiments and only intervening digit was performed, which does not have measurable counterfactual truth as intervening thickness and intensity in Morpho-MNIST. In general, the "interventions" performed in Sanchez & Tsaftaris (2022) all drastically modified the objects in the original images and makes it quite impossible to tell exogenous noise are preserved either quantitatively or qualitatively. When one drastically modifies the objects, the problem moves very far away from the consistency assumption, and it becomes really questionable whether the results under such "intervention" can even be defined as counterfactuals or not. We want to clarify that we think Dash et al. (2022), [anchez & Tsaftaris (2022) and Komanduri] ctt al. (2024) are meaningful and well-written works that conducted the important advancement of experimenting diffusion models in counterfactual generative modeling, however, diffusion models have not shown the capability of exogenous noise abduction on the level of sate-of-thee-ar HVAEs
(Monteiro et al. 2023 Ribeiro et al. (023) as it stands.
To give the readers a sense of how much difference makes to use the VCI framework instead of conditional diffusion models for counterfactual generative modeling, we presented the visual results on CelebA in Figure 5 and the results on "intervening" digits of MNIST in Figure 13 for readers to compare aforementioned prior arts.
D COMPARISON OF TRADITIONAL CAUSAL FORMULATION AND VCI FORMULATION
We note that we formulate counterfactuals Y and T' as separate variables for the cleanliness of graphical model and notation simplicity in variational inference, with the goal to bridge readers with variational inference background and causal inference background. In essence, it is no different from the traditional causal formulation where the do() operation is directly imposed on factual variable T, as T is just a dummy variable with no observations. Denote Ydo(T=ar) as Ye, see Figure 7 for a demonstration of our formulation in traditional notations.
Then under the ignorability assumption (Assumption [ note that there is no Ey and ey in the traditional SCM formulation as there is only one copy of exogenous noise Uy injected factual and
17

---

Published as a conference paper at ICLR 2025
Vy
Z
Uy
Y
Ye
Yr
Ux
Ur
Ux
Ur
do(T = =
do(T 1)
X
T
X
(a) SCM, traditional
SCM, with fy fx fv(f2(X,Ux),))
Y
Uy
Ye
UT
Ux
t
T
X
(c) SCM, with fr (X, (x Uv), unfolded
Figure 7: The evolution of our formulation from traditional SCM with observed triplets (X, T. Y). Figure 7b is simply extending the production function towards Y with an intermediate variable 2 representing two of its inputs (when fz(X, Xx, = (X, Uy), fy is just fy). Then, Figure 7c is just the unfolded graphical model of Figure 7b the do() operation replaces T in the production function of Y with /", i.e. YU = (x, Uy).
counterfactual outcome, but in general if there exists ey and by independent of Uy the following result holds long ey, ,y 1 UT), have the following version of Theorem |
Corollary 2. Suppose a collection of random variables follows a causal structure defined by the Bayesian network in Figure 7b Then loglp(Yr X, (,), has the following variational lower bound:
log (Y,
BIIZYT, log - ( (V)()
DEL (P(Z],, (()
(10)
where Dip || q)
log log
Now define dummy variable
T' -T U7),
(11)
with U/ satisfying the ignorability assumption (Remark D and
y'= (X, Uv) = fy (f2( Then, Y! dT(TT= by) (XX Yp. Hence, have
(12)
= ", do(T" = ()) XX,,"" =
(13)
for any t.. Therefore, Figure Ia and Theorem is essence just Figure 7c and Corollary with different notations and variable definitions.
E IMPLICIT COUNTERFACTUAL SUPERVISION
We demonstrated in Appendix that the marginal-level observed likelihood is unorthodox for counterfactual generative modeling and the individual-level likelihood log[p(Y"|YyX> X, T') is better motivated for this task. However, for researchers who insist on working towards the traditional likelihood objective log[p(Y\X,, (XXT)) T'] composed only of observed variables, there is approach to conduct counterfactual supervision implicitly - to view Y' as a latent variable restricted on the prior p(Y' X, T') and view counterfactual construction task latent recognition task:
18

---

Published as a conference paper at ICLR 2025
Proposition 3. J(O) loglp(Y has the following variational lower bound:
J(O)= 2 Rp(22/,T..x) log | 7)) Pp(zIY,..))DR DKI. lx(Y' (("(XX)
DRI. (p(Z T, X)))(") p(|]"" T ())(()2,
(14)
This induces the following ELBO in stochastic optimization:
(o) t, log (yl t)] Bao(ala DEL (po 100(/12,, l2,t') p(y'l2, ( t')l DRL. [= 2, po (v')=,,') Il co(2/// (lll/,/.))pp) x))p(y/ [z,
(15)
Equation 15 avoids estimating Y' with single sample as in the objective (Equation 3) derived from the explicit ELBO (Equation ( However, it could be lot more sampling inefficient as the latent divergence term requires Monte-Carlo samples from both and Y' to be estimated.
F DEFINITION OF ORACLE CONSISTENCY, ORACLE RESTRICTIVENESS, AND DISENTANGLEMENT
Let = 1100(0000 p: - [0,1 V: [0, 1|=5} be hypothesis class of stochastic models (p. The input space of is measurable space S where a feature vector s E is drawn; the input space of q is the product space of V and X' where vector of interest E V and some auxiliary attributes æ E X' are drawn. The outputs of p and q are probability measures on V and S respectively. Let the true model (p".9") be the true conditional distributions (p()), (S of random variables V and S. The following definitions provide criteria under which two parts Si and S of feature variable S separated by an index set are disentangled by a model.
Definition (Oracle Consistency). Consider
(s1,8,) and s' (s1, s(1) where
x ((X)
(16)
((S((z)
(17)
811 14 li p(S,,111,2))
(18)
We say that q is consistent with S1 if is the true distribution po and
E E l/ pp" (8;a) qı pp" (8': ;1)] = 0,
(19)
where )r(srlu,a) = Say ā(sr ous,, x).. More generally, we say that (ø, ÿ) is consistent with S1 under if
Epællår p(s: x) q1 . p(s';)))) = 0.
(20)
As discussed in Shu al. (2019), oracle consistency (or encoder consistency) states that for any fixed
choice of S₁, resampling S.z should not affect the oracle's measurement of S₁. A more general notion of oracle consistency is defined in addition, as it is valuable in real-world optimization to evaluate the estimating model's consistency under its own oracle q p.
In contrast, restrictiveness states that for any fixed choice of S.,, resampling S should not affect the oracle's 's measurement of S.y:
Definition 2 (Oracle Restrictiveness). Consider s = 8V,) and s' = (55, where
x ~ p(X)
(21)
S p( 5()))
(22)
ST, 8T,5 s üd ((Sr|s(1:3)
(23)
We say that q is restricted to Si if is the true distribution po and
Epo.r |lā pp" (x; ã p (s': 0,
(24)
where á (5 = -s, q(sl, 8 a)) More generally, we say that (ø. ②) is restricted to S₁ under
p if
E (())2)) 7. ( 0.
(25)
19

---

Published as a conference paper at ICLR 2025
Notice that the oracle consistency and restrictiveness of S₁ is equivalent to the oracle restrictiveness and consistency of Sr,, .. we have the following definition of disentanglement:
Definition 3 (Disentanglement), We say that q disentangles S₁ and Su if is consistent with and restricted to S₁. More generally, we say that (p, i) disentangles S1 and Su under p if (P, i) is consistent with and restricted S₁ under p.
G ROBUST MARGINAL EFFECT ESTIMATION
By Van der Vaart 22000, a sequence of estimators L is asymptotically efficient if 1 - 1/vn Ei=1 1 Wx) + (1) where Vi is the efficient influence function of (p) and W₁ p(W). The theorem below gives this efficient influence function for the average treatment effect of the treated (ATT) (p)) E, D =o) and thus provides a construction of an asymptotically efficient
regular estimator for V::
Theorem 2. Suppose W : Ω + Rw follows a causal structure defined by the Bayesian network in Figure 7 where the counterfactual conditional distribution p(Y', Z,, X) is identical to that of its factual counterpart p(Y, T7,T Z, X). Then (p) has the following efficient influence function:
(W) 1(T = ()) (Y Ep (VIZ, + Rpp 2 2, T o) V.
(26)
X
This leads to the construction of 1 (0) in Equation which is asymptotically efficient under some regularity conditions (Van Der Laan & Rubin 2006). with the asymptotic distribution given by
Vn(ū, (p)) 4 N Ep (W)))
(27)
where the variance can be estimated empirically in practice and confidence bounds can be constructed accordingly. The robust estimator within the VCI framework can also be extended to estimating covariate-specific marginal treatment effect E(p) Ep ==] for given covariate of interest:
Corollary 3. Under the same settings as Theorem function:
E(p) has the following efficient influence
(WW 1(X c)
1/ (W).
(28)
(
Note that when p(X) is estimated empirically, the asymptotically efficient estimator for E(p):
1 (xx =))
1(th 1()==a) y/ + (1 1ltx-
e(1x =()==a))) IEp [Y' T = al
=a))
(o
(29)
1=1 p(xx)
1 nhe Th
I(th. =Q)
yk. + (1 I(te, = aa)
Epo [X' =a
(30)
e(tu, (
E(ta,
simply an application of Equation 4 on the set of observations with covariates e, where (ki, ) are indices such that Tk.
H FORMAL DEFINITION OF COMMON CAUSAL ASSUMPTIONS UNDER THE VCI FORMULATION
Assumption 1 (Ignorability). There is no unobserved confounders, i.e. Uy 1 Ur, where Uy (Uy,e,,ey).
Note that this assumption can be extended to the counterfactual treatment variable T' w.o.l.g.:
Remark 1 (Ignorability). Note that T' in our formulation is a dummy variable with no observations, and we only care about the outcome distribution of Y' under assignment do(T" = t') for any given t'. Hence, the ignorability Uy L U/. and in addition, Ux 1 Uf, can be assumed w.o.Igg just so that the outcome distributions under do(T' = t') reduces to that under T' = t' conditioned on X or 7. for notation simplicity, without violating any actual mechanism the observed data.
20

---

Published as a conference paper at ICLR 2025
The ignorability assumption is universally assumed and entailed in our formulation in Figure
Assumption 2 (Consistency). An individual's 's observed outcome under treatment is the same as its potential outcome under that treatment, i.e. Y Y do(T'=Tr where Y do(T"==) (z, Ty yielded by replacing the production equation for T" with T =T
The consistency assumption, commonly stated causal inference literature, is indeed an assumption and not just direct consequence of the SCM under our formulation, since ey and éy are two copies of independent variables. In essence, it states that ey = and there is only one copy of exogenous noise Uy injected Y and Y':
Remark 2 (Consistency). Under the consistency assumption, (x) must be deterministic for any and I, otherwise we would have + YooTr= since (y(2,t,ty) and (v L, kv) are independent for any and Hence there exists deterministic futnntin fy such that (Z, (2, T, ). So w.o.l.g., we let ey 0 under the consistency assumption.
Note that this assumption does not result in a collapse of counterfactual inference to interventional inference in the traditional sense, latent Z is unobserved and entails further uncertainty Uy beyond the observed covariates X.
1 PROOF OF THEORETICAL RESULTS
1.1 EVIDENCE LOwER BOUND
1.1.1 PROOF OF THEOREM
Proof. By the d-separation (Pear 1988) of paths on the causal graph defined in Figure la we have
log [p(Y'|y, X, T.). log (ZIT.X) [p(Y' Z, Y, X)T)T") 2 EXIZYTX)) log [p (Y' (Z Y, X,Y,X,TT))
(31)
(Jensen's inequality)
(32)
P(Y', XX,T),,
EAIZIT., log
(33)
p(Y', T,Y)X,T,TT)
EPIZITT,X) log
(34)
P XX)xXX,,X)((
P Z,T))(Z)YY, T', ())((Y'
EPIZIKTT. log
(35)
p(Z T, (Y
EX(ZIV,,)) log [Y(Y DEL xX p(ZYY,,
-D (x( p(Y'|x,T))
(36)
Note that the above steps still hold with an additional edge from Z or T to X (but not both) if Uy or Ur dependent of Ux.
I.1.2 PROOF OF PROPOSITION 3
Proof. By the d-separation (Pear 1988 of paths on the causal graph defined in Figure Ta we have
P(Y, Z, T',, p(Z, "',
log [p(Y|X, log |y
log (Y..
(37)
PYY Z, X,,,7
E,(..IIYX log
(Jensen's inequality)
(38)
p(Z, Y X,T,T")
p(Y' "xTT")((,,,_""
Ex log
(39)
p(Z,Y Y Y,
P (') T", X)P(Y
-ENZFIXXTTY log
(40)
PIZIY.T.X'p(('II.TT))
p(
-EIZYIXXXXTT log
(41)
21

---

Published as a conference paper at ICLR 2025
--B DKL (p ||
DKL (')(Y// "'
+ -,(22)1;,,)) log [2,
(42)
Note that the above steps still hold with an additional edge from Z or to (but not both) if Uy or Ur dependent of Ux-
1.1.3 PROOF OF COROLLARY 2
Proof. Note that Theorem J only relies on the conditional dependency structure in Figure T and the only difference between the dependency structure of (X, Z. T. T", Y,Y") in Figure [ and that of (X, Z, t,, Y,) in Figure is the dependency between (X, T") and (X,t') However, such difference does not change the conditional independence of latent posterior:
X,, p Yr l', X). the conditional independence of out- TY) and DlY 1', T) p(Y, ("(), and the condi- p(Y T) and p(Yr|X,,''s t', T)) p(y X, t'). Hence, the same proof for Theorem | applies to Corollary 2 by replacing 7" with and y' with Ye.
[ and p( (r , L', come distribution 2,
X,T)
tional independence of covariate-specifit outcome distribution XY,
1.2
DISENTANGLED EXOGENOUS NOISE ABDUCTION
1.2.1 PROOF OF LEMMA
Proof. Let S₁ Z and S T. On one hand,
pp llă p ((s1, s11)=) - ãi op x))==)]
hww.e//tnn.n.....ccc.....eoeeeee (v, (1 x))l
(43)
(44)
Epplastt llgl x)),
(45)
S Epo(d) V DEL la(-ly, t. (0 [1,))] (Pinsker's inequality)
(46)
0,
(47)
since Equation 5 jimpiies that DKL [g(-ly,
i, x) DKL(a(-lv,
g(-ll', d', (',=)] = 0 a.e. with respect to po- Hence q is
consistent with Z. On the other hand, is restricted to by design:
E
Pho.p-llg 1l o ((8), s)));)) 2. ((s), (();=))
(v(((==))
(48)
Pente:: (0 - (v(55,4,)))
(49)
( (+)))
(50)
O.
(51)
Therefore q disentangles Z and T.
1.2.2 PROOF OF PROPOSITION
Assumption 3 (Uniqueness). Two different individuals do not have the exact same outcome (a.s.) under the same treatment a, i.e. for y ~ p(Y|z,())))) a.), y' p(Y|='(, a., we have y # of a.s. if +
The uniqueness assumption weak and automatically satisfied under common assumptions in prior works such injective decoder, continuous outcome or diffeomorphism (Khemakhem et al. 2020 von Kügelgen et al., 2023). Note that this assumption could be unrealistic in traditional causal inference but very reasonable in high -igh--iimnsional outcome settings take facial imaging dataset for example, no two individuals with different facial features should look exactly identical pixel-by by-pixel under the same treatment.
22

---

Published as a conference paper at ICLR 2025
Proof. Denote
LB 9.) == Bpo(d) ( (slust))) log (ylz, f)] log (1,
DKL 1q. (-l,, x, q. x1)
(52)
and suppose that (p-, q.) ) E arg p.aa LB( but
DKL (..(,,,,,,)) a) g g-(-lv// [',, >0,
(53)
which we can factorize as where p²(2,t) (T)) Then we consider model (per. que) such that Vz, t, we have
Epo(x,,,t.''... DKL le. op op' (2,1;z) op (z, (2,1';a))
(54)
q. (,, (z)=)
(55)
P-. (Y,(())())) -0oo221,00) (Y)
(56)
where fo(2,t, 1,.,t,y cy) Po t) is the true SCM and y is any injective transformation of 2 (e.g. (()) =). Equation 55 is well-defined by Assumption 2 and which guarantees that fo(2, + fo(=', 0) for + =;; Equation 56 is well-defined due to the injectiveness of 7, Qa = Qo. P.. elsewhere. It follows that
Epo DKL (-1,,,, t, x) ll (..(,,b
Bpolz,,, " DKL la. (2, q. op (2,(x();
(57)
Epo( DKL (f O), (l (1,,,,
(58)
E DKL (x)) (+)
(59)
< (y DKL la. ly. f, x)
(60)
by Assumption in the meantime
E Poo( (el (alutt)))
log (vl2,
log (p-- (yl 2.., t)]
(61)
(po(a...)) R (x.. 1f0xx..))).). log [p.. (fo(2, [)]
(62)
Ppo( log (p.. (F0 ()))
(63)
Epo(k....) log (o( (fl
(64)
Epo(a....) By (x. (o(...)))..) log [ (Foe,,
(65)
2 Epo(a.s.t) B fo(a. (ot. log (yl (.))
(66)
Epo(s,,,t) Ez. (xttr)) log (p- (p((2, t)]
(67)
Ppp By m.) log (y]=,t)
(68)
by Assumption Hence we have
LB(p.. (ox) > LB(p
(69)
which contradicts (p..q.) q.) E arg max .g. LB(p..q.). Therefore (p., 9.) must satisfy
Epo(ay DKL. [g- (-1v, 6, x. |l g-1//1/ f, v,,]] 0
(70)
and it follows that q. disentangles 2 and T by Lemmal
1.2.3 PROOF OF PROPOSITION
Proof. The strategy is similar to the proof of Proposition We first prove statement similar to Lemma that states a sufficient condition for empirical disentanglement (sufficiency), and then argue that this condition is attainable by optimizing Jpan. (p., (attainability).
23

---

Published as a conference paper at ICLR 2025
Sufficiency
We prove that if
I pa (0. A. v) Peasa (t' 1 Paata (1"))) DEL. ((.ll..,,,2)) || (0""",,,, """""""""" [,,,,) =
(71)
where y. = Ep, ( (PPPat--1 2o0ry and y. Ep 2 Zern) Y, then (g., 9) disentangles and T under Pe =fy e. -Pdaaa-
Let S₁ Z and Sy1 = T. On one hand, notice that for any x, z s.t. Pe 2) O, have only if pe(, > 0 only if Pdata
(t(2,, > 0 only if (e(x, O for any t. It follows that
(2; (x. > O only if pdata (t)=)) >
Ep llĀ 0 (2;)) ब - 0z xz)
11 9 (z, [':=)) - a (2,1"=,)]
(72)
Ep
(/, qr(y"" 1" ,))]
(73)
Epane (2. 1.0) (P
(.118" , ,t" .>)l
(74)
Pe DKL la( U' 2) (-1," 1" x)] (Pinsker's inequality)
(75)
0,
(76)
(4((1//l//// (/ (()))) || q(-lv", 1", ,x)] a.e. with respect to Pdata- Hence (g., 9) is consistent with Z under Pe. On the other hand, (g., 9) is restricted to Z under pe by design:
since Equation 71 implies that DRI,(4(
Ep... 1ly o (e', - t;2)]
Ep.c..thi.. o प. " t;:)]]
(77)
Pe (1 Y,t,)) (1 Y,t,,, (Epl)(())) 6.(-
(78)
(79)
(80)
Therefore (g. disentangles and T under J, e Pdata-
Attainability Let A = (K Wa, Ham be the set of all unique entries in {( La, ( Suppose that (p.. ) E arg max p... Spanc (v. but
(1 DKL (-1(%,,%,, f',x) g- (- ,,,,)] > 0, (81)
where 1/2 = Ep- (11E (Inia 2,0Y and y/ = ((b)) i.e. E1 : 1 S S n and
1) {()
Pdata (Z1)) Pdaca (taca(((()) (z1) > such that
DKL (0(- (20),, mn) ll (((1)!?2) (12) .mm] >0
(82)
where ( Ep. (-I,,, im.
2,00 2,00)) and y!?
2, Then consider
model (p.., (=a) such that V(x, v, E A, we have
T.. 2 E p. B v.d Z,0 Z.) a, g. (Z1y,,,,,))
(83)
for any a s.t. Pdaaa (T = al x) > 0. Equation 83 is well-defined since
Ep. Fy we. 2.,) Y Ep. ( m
2a) 2a
+ 9- R. (-1,,1.0) 7, g- F (1 (HW 7, (-9. 7 = EgG- (84)
+ e. (7 x) e. (7/y', ,1) = (v. = (v', (=))== = g. (7l, 1,x) ) = (Z'y'''',n) (85)
by uniqueness condition 1) and the construction of {q.}. 9.a 9. elsewhere and p.. p. everywhere.
Firstly, show that (pe.. q.x) is still member of the model class of interest. Since q.u is still normally distributed with unit variance by definition and g.. still satisfies uniqueness condition by definition, only need to show that (s.) still satisfies uniqueness condition 2). This
24

---

Published as a conference paper at ICLR 2025
is because Vx,, t, V, Pdaca (2,, t, y) Pdata (t'(z) > O. we have (x,t,g) E A and thus (æ,t,y)
(a,, at p.p (-IPo =ny) 7,00 Y) for some (Ik,, tk. 3kk)) E A only if
2 V) (2r, a, (-Ix.")) (slun, only if (xr, ax Ep (-1E, tl me, 2,00 Y) = (Ikk, the Vk,) ) by the fact that (g,, satisfies unique- ness condition 2). this case, we have (z,, t, v, = (IX lxi Vk.) and
(I, O, (. (la, 2 2,xx
E
(86)
q:o (2,,, = 9-o (... (110,, Aa. *x, 2 2,0)) Y, ax
(87)
= q- (Z tR Tk.) = qo (2) t, x).
(88)
If (2,,, o (20, ot, Bp. (-122, -2, 2.a)
Y) for any (IT,, , 1k; yh.) E A, then g.. L, x,
q. (Z t) by definition. Therefore we have 9as (Z g. t, x) = 9. Zly t, x.). ) follows that y/., =
Egg (-12,t))e, (=lu.t,,,)) Y y. and (gex; C.x) satisfies uniqueness condition 2) by the fact that (g., ew) satisfies uniqueness condition 2).
Secondly, we show that (p-, q.) is no longer optimal, thus raise a contradiction. On one hand, we must have
Bpaa DEL le- ( ly,, 1, |l g. (-a(: 1',, xx)]
Paata (21, 4,, ()) Paata (1)) (q. (-lun, lt, 222) (. ,10, (0, ,11),)) mn)
+Patt (1(2)11)) DKL q- (1gg, ll q. (- ((2), 2000
(89)
0,
(90)
otherwise we would have q- (Z ti, mi) = g- (Z ((1) m) = (z(y)) ((2),, x)) a.e. which
contradicts Equation 82 Hence
E Epant KL -9-- (- (1, x,
Pdata (x. t,
Σ Pdata (t'lz) DEL la (-ly,, x.) |l q-.., 2),
(91)
(2,1,))E
=
Σ
Pdata (x, t, v,
Σ (tattat'|zz) DKL (-1,,, t,x) (...1,,,,)))
(92)
(2,1,))EA
0 < Epeces DKL la. (-ly, t, x) ll (-b6,, f.. x]]
(93)
by Equation 83 and the fact that qoo Ely,, x) = qu (2 y,, t, x) for any (x, 4,) E A as shown above.
On the other hand, we have
Epasca ( log pp.. (y) t) + log [ply'n |x, t')}}
Eppaat {E log [p. (ulz, 1)] + log (/ [r. (')]]
(94)
simply by the construction of p.s and the fact that q.. (Z 1,, q- (Z 2/,, 2,), v.. = v/ for any (x. f, E as shown above. Hence we have
Spans (pr., 9o-) > Spanm (p..anpp.. (p., which contradicts (p., q.) E arg max, p... Jpsaia (p., Therefore (p-, 9.) must satisfy
(95)
E Epaan (t..) (flkaaa(""))))paaaa peaas ( (x) DKL. Ig- (-11/e l', x)) ll q. (-19", 1", ["] = 0 and it follows that (g., q-) disentangles Z and T under I, e. PaPata by sufficiency.
(96)
1.2.4 PROOF OF COROLLARY II
Proof. By the proof of Proposition 2 (attainability). we have
Ppaauu(sstta)))aa PPaaa Paa)) DKL. [g. (-1V6, x.)x) ll (-1v", t" [=],, =0 where 1/4 (-1g (11s Z and y/ = (Ppt(ii(a))))))))))))))). hence Bg. =
(97)
Eyg ( v", Z a.e. on Pdata and it follows that
E Epanca(r (t). w)an DKL [e- (-1u/, xxx) || es (-1'",
=Epasta (x. Panan ( (f( paa (1*)z) DKL. (8,,,ci0 ( || 8g (158 in (n)).,,,))))) =0. 0.
(98)
Notice that y¹. = Ep.( (l Z,y") Y and y/" Ep. (.EE, (is..)) Z, Y, we have (ū., ē.) disentan- gles 7 and 7 under J, edPaata by the proof of Proposition 2 (sufficiency)
25

---

Published as a conference paper at ICLR 2025
1.3 MARGINAL ESTIMATOR
1.3.1 PROOF OF THEOREM
= Figure [ Following Van der Vaart 2000). we define a path Pe(A) +[
Proof. (p) has the identification Ppp)=
ai] under
on density
p of random vector A as submodel that passes through at é in the direction of the score S(A) de log(pe(A) l==0
Following the key identity presented in Levy 20199
de Pe (1:(AlppaA)) = =\-11
-AA,) = (()))A = >, (A,,) = X (((A)))aa) = Ar-1))
(99)
where pa( denotes the parent nodes of variable A,, we have
d de c=0
d
(pe))
Bp,
Ep,
12'" all
(100)
de e=
d
D /ell'/1,, a),x(
(101)
de l==0
I Ju
d
de ( a))p(2222)
(dominated convergence)
(102)
[ V'((,,) 1) d WAR
de c= Pel''lz, T' a)
(103)
2.2. V/ = de pr(2,)) p
+
(104)
1(t' = a) p(t'lz)
p(1(xx) 1/ p(y, 1)2,,,)(,(())) x) de pe
y'P(U'12, a) d de (2)22)
+
(105)
1/,2
I(U' =a) ,
=1 x) (S(w) - (sW,),,,,,,
p(('(z)
V'P(U' =0) )' (S(S)))) [S(W)))
(
(106)
Vi
I(t 11=0 a)
up(u p/(u.t/2.200 2)/,,))p(/// =),()) ((()) - E (0,)),,,,,,t''
p(t|x)
1
y v/ply |z, T' =) (B |S(W))
(107)
T(t T(t=a) up(w)
S(w)
p([[z))
E [[S(W))'',,,,,4 2,,44, 1(t a)
p(t(z)) up(y|2,0 1)
+
[l 2 =a)
[[S(W)) -" a
(108)
I[t (
(w:)
T'
T' = a] v p(w)
(109)
p(t|x)
by the assumption of Theorem and factorization according to Figure Hence
d
S(W), I(T a) (Y R, (VIZ, + 2 = al 1T/
((p.)
(110)
de
1=0
x(TIX)
and we have (T) (V-
= a V.
26

---

Published as a conference paper at ICLR 2025
1.3.2 PROOF OF COROLLARY 3
Proof. The proof simple extension on the proof of Theorem
We have
d c=0 de
d
(pp))
Ep, [..
Y Z, =
|X cl
(111)
dc
dele=0
( 3 ppe (v 2,, =a)p c)]. de dele=0/ v/
d
(112)
Hence, following the same derivation, Pe x) in Equation 102 and onwards becomes pe (2 = c). while p(z, x) in Equation 105 and onwards becomes T(x: c)/p(x) p(2, Therefore, have (W) I(X = c)/p(X) (W).
J MODEL ARCHITECTURE FOR IMAGE GENERATION TASKS
The specific model architecture used for image generation can be found in Figure g Note that the image generation tasks in our experiments, condition belongs rather than X. this is not the case, the non-intervenable conditions would compose covariates X and would only be inputted to the encoding model but not the decoding model.
treated every condition as intervenable, thus every
01:64
72612
32H32
(4)
Figure 8: Model architecture of the latent recognition model (grey) and outcome construction model (pink) for image generation tasks.
To incorporate (t') and æ into the convolutional layers, we first embed them into vector representa- tions (tabular embedding for categorical variables; sinusoidal embedding with non-linear mapping for continuous variable), then repeat and expand each value of the vectors to match the input resolutions of the corresponding blocks and concatenate them to image representations as extra channels, similar Monteioo al. (2023). Note that there are other ways to conduct image-vector incorporation such as the scale-shift approach commonly used in diffusion models to incorporate time step variable (Sanchez & Tsaftaris 2022). Each convolutional block (each rectangle in Figure & contains 1 embedding layer, a 3 3x3 (or 1x if resolution is lower than 3x3) botteleneck layer, and 3x3 (or 1x1 1 if resolution is lower than 3x3) output layer.
K DETACHING PATTERNS
In practice, there are few optional gradient detaching options when evaluating the divergence term to enhance the stability of stochastic optimization. Our default behavior is to use copy of de for the evaluation of the divergence term which is updated after every epoch, while preserving the gradient path and gradient computation for V6.0- This technique is analogous to the handling of surrogate objectives commonly applied reinforcement learning (Lillicrap et al. 2015) Van Hasselt et al. 2016; Schulman al. (2017), and the copy of 9s is seen the target network. In addition, provide the option of detaching o from Vé (focusing on the second purpose described in divergence interpretation: preservation of individuality in decoding) and the option of not detaching the gradient
27

---

Published as a conference paper at ICLR 2025
computation for 46 (reflecting the first purpose in divergence interpretation: recognition of mutual features in encoding) during the evaluation of the latent divergence term.
L BENCHMARK ADAPTATIONS
L.1 AUTOENCODER
The adapted autoencoder reconstructs the outcome during training similar generic autoencoder, but takes treatment and covariates as additional imputs. During test time, we simply plug in the counterfactual treatments along with factual outcomes and covariates generate the counterfactual outcome predictions.
L.2 GANITE
GANITE Yoon et al. 2018); 's counterfactual generator does not scale with a combination of high- dimensional outcome and multi-level treatment, thus here only input randomly sampled counterfactual treatment to the generator and correspondingly construct one counterfactual outcome for each sample. See Figure 9 for the original and adapted structure of the model.
Vf-
yf.
Yef-
Yef.
x
Vf
t
yi
(a) original
(b) adapted
Figure 9: GANITE's counterfactual generator. tef is random sample of t, passed into the generator as part of the input (z, tef Ms), separately from input (x, uy) of the factual generation.
The discriminator predicts the logits L,, les of yy, Bef. separately. The cross entropy loss of (l,, kef) against (1, 0) is then calculated.
M EXPERIMENT SETTINGS
In this section, we describe the settings of our experiments in detail. For complete informa- tion on hyperparameter settings, see our codebase at hthttps::/giinh thub. com/ rulun irriational-cauaa wawpall--casaal-leeeeeeeee iinference
M.1 SINGLE-CELL PERTURBATION DATASETS
For both datasets in experiments, two thousand most variable genes were selected for training and testing. Marson has 2,000 dimensional outcomes, 3 categorical covariates (cell type, donor indicator, stimulation indicator) and 1 categorical treatment (target gene), Sciplex has 2,000 dimensional outcomes, 2 categorical covariates (cell type, replicate indicator) and categorical treatment (used drug). Data with certain treatment-covaraat combinations are held as the out-of-distribttion (OOD) set and the rest are split into training and validation set with four-to-one ratio.
Out-of-Diste Distribution Selections We randomly select a covariate category (e.g. cell type) and hold out all cells in this category that received one of the twenty perturbations whose effects the hardest to predict. We use these held-out data to compose the out-of-distribution (OOD) set. We computed the Fuclidean distance between the pseudobulked gene expression of each perturbation
28

---

Published as a conference paper at ICLR 2025
against the rest of the dataset, and selected the twenty most distant ones the hardest-to tardest-to-prediet perturbations. This is the same procedure carried out by Lotfollahi et al. (2021).
Siiferentially-Expeeseeee Genes In order to evaluate the quality of the predictions on the genes that were substantially affected by the perturbations, select sets of 50 uufferentially-expessseee (DE) genes associated with each perturbation and separately evaluate model performance on these genes. This is the same procedure carried out by Lotfollahi al. (202T).
Hyperparameter Settings All common hyperparameters of all models are set the same as the defaults of CPA (Lotfollahi et al. [021]: an universal number of hidden dimensions 128; number of layers 6 (encoder 3, decoder 3); an universal learning rate 3 3-,, weight decay rate 4- Contrary to CPA, we use step-based learning rate decay instead of epoch-based learning rate decay, and decay step size is set to 400, 000 while decay rate remains the same at 0.1. Batch size is 64 for Marson and 128 for Sciplex.
Other Details We used the empirical outcome distribution (with Gaussian kernel smoother) strati- fied by X and T to estimate the covariate-specifii outcome model pX The fully-attached detaching pattern is applied where both Ve.s and qo not detached when evaluating the divergence term. Models are trained on Amazon web services' accelerated computing EC2 instance G4dn which contains 2nd Generation Intel Xeon Scalable Processors (Cascade Lake P-8259CL) and up 8 NVIDIA T4 Tensor Core GPUs
M.2 MORPHO-MNIST
We used the original Morpho-MNIST training set for model training, and the original Morpho- MNIST testing set as the observed samples for model testing. The training and testing set have 28x28 28 dimensional outcomes, covariates and 2 continuous treatments (thickness, intensity). Three different testing sets are counstructed based on the counterfactual treatment settings: only modifying thickness for do(th); only modifying intensity for (ii;:); randomly modifying thickness intensity for mix. For each observed sample in the testing set, we randomly sample the amount modification based on the range of the selected modification.
Data Augmentation For each image in the training (28x28), we pad by margin of 4 pixels on each side 36)36), then randomly crop 32x32 region. This is the same procedure carried out by Ribeiro et (2023)
Hyperparameter Settings The model width, depth and resolutions of our framework are set to the same as Ribeiro et al. (202)): number of channels for the hierarchical encoding blocks are set to (16, 32, 64, 128, 2 with resolutions {32, 16, 8, 4, 1} and number of blocks {4, 4, 4, 4, 4); number of channels for the hierarchical decoding blocks are set to {256, 128, 64, 32, 16) with resolutions (1,4, 4, 8, 16, 32) and number of blocks {4, 4, 4, 4, 4). Each block contains a bottleneck layer with one-forth the number of channels the input/output number of channels. Training is conducted with a batch size of 32, an universal learning rate 1-4, and weight decay rate 4 Learning rate decays at epoch 100 and linearly decays 0 at epoch 200.
Other Details We used the adversarial training approach estimate the covariate-specific outcome model pY||x, 7). The default detaching pattern (see Appendix K is applied where a copy of go is used critic evaluate the divergence term, and updated after every epoch. Models are trained on Amazon web services' accelerated computing EC2 instance G4dn which contains 2nd Generation Intel Xeon Scalable Processors (Cascade Lake P-8259CL) and up to 8 NVIDIA T4 Tensor Core GPUs.
M.3 CELEBA-HE HQQ
The two factors of interest smiling and glasses, regarded intervenable. Therefore, the training and testing set have 64x3 dimensional outcomes, no covariates and categorical treatments. The train-test split of the CelebA-HQ dataset is inherited from the original CelebA dataset.
29

---

Published as a conference paper at ICLR 2025
Data Augmentation For each image, crop the x28x 128 center region, then randomly crop x20x 120 region within the center region. Afterwards, we resize this region to 64x64, and apply horizontal flip with probability 0.5. Note that data augmentation is applied to both training and testing set in this experiment, since we are sampling and empirically examining the counterfactual results and not evaluating any metric.
Hyperparameter Settings The model width and resolutions of our framework are set to the same as Monteiro al. (2023): number of channels for the hierarchical encoding blocks are set to {32, 64, 128, 256, 512, 1 [024] with resolutions {64, 32, 16, 8, 4, 1}: number of channels for the hierar- chical decoding blocks are set to {1024, 512, 256, 128, 64, 32} with resolutions {1, 4, 8, 16, 32, 64). The number of blocks is set to {3, 12, 12, 6, 3,3 3} for the encoder and {3.3 6, 12, 12, 3} for the de- coder. This chosen to be slightly more shallow than Monteiro et al. (2023) (f4, 12, 12, 4, 4] and (4,4,8, 12, 12, 4)) Each block contains bottleneck layer with one- the number of channels as the input/output number of channels. Training conducted with a batch size of 32. An universal learning rate 3 3-4, and weight decay rate 4-7 are inherited from the experiments with single-cell perturbation datasets. Step size for learning rate decay is 400, 000 with decay rate 0.1.
Other Details We used the adversarial training approach to estimate the covariate-specifer outcome model p T). The default detaching pattern (see Appendix K) is applied where a copy of qo is used critic to evaluate the divergence term, and updated after every epoch. Models are trained on Amazon web services' accelerated computing EC2 instance P3 which contains high frequency Intel Xeon Scalable Processor (Broadwell E5-2686 v4) and up to NVIDIA Tesla V100 GPUs, each pairing 5, 12 CUDA Cores and 640 Tensor Cores.
N ADDITIONAL EXPERIMENTS AND DETAILS
N. MARGINAL ESTIMATIONS ON SINGLE-CELL PERTURBATION DATASETS
In this section, we use the same evaluation metric as Section 3.1 but compute each R with the robust marginal estimator and compare the results to that of the empirical mean estimator. Note that on OOD set, the robust estimator reduces to empirical mean since no perturbation-covariatss combination exist in validation set. Therefore, we compute the R of the marginal estimators using samples from the training set against the true empirical average on the validation set in these experiments. In each run, we train a VCI model for individualized outcome predictions, and calculate the evaluation metric for each marginal estimator periodically during the course of training. Table shows the results on Marson. Note that the goal of robust estimation is to produce less biased estimators with tighter confidence bounds, hence we report the mean and standard error over five independent runs reflect its effectiveness regarding this goal.
Table 5: Comparison of marginal estimators on Marson Schmidt et al. 2022
All Genes
DE Genes
Episode
mean
robust
mean
robust
40
0,9141 1 0.0159
0.9343 1 0.0080
0.7108 + 0.0735
0.9146 + 0.0305
80
0.9171 + 0.0104
0.9349 + 0.0068
0.7274 + 0.0462
0.9163 1 0.0254
120
0,9204 = 0.0097
0.9352 + 0.0063
0.7526 + 0.0447
0.9182 + 0.0229
160
0.9157 + 0.0043
0.9355 + 0.0053
0.7383 + 0.0402
0.9203 + 0.0199
As is shown the table, the robust estimator provides a crucial adjustment to the empirical mean of model predictions especially the sard-to-predice elements of diggeniiensional vectors. Such estimation could be valuable in many contexts involving high-dimensional predictions where deep learning models might plateau rather low ceilings.
N.2 MORE COMPLETE ABLATION STUDY
Note that the latent divergence term is essentially a regularization term to counterfactual supervision, as discussed in Section 2.2 hence it would not be meaningful to train VCI with the latent divergence
30

---

Published as a conference paper at ICLR 2025
term and without the counterfactual supervision term it would just further worsen the training-
inference disconnection problem described in Section [] without any real benefit. However, we acknowledge that presenting this case would make the ablation study more complete, and therefore have made this inclusion (marked as HAE-A) in Figure T0 As expected, HAE-A is even worse than
HAE.
do(thickness)
do(intensity)
mixed
9.0
5.4
HAE
1.6
4.8
HAE A
MSE (10-))
7.5
1.4
4.2
SAE
6.0
1.2
3.6
VCI
3.0
4.5
3.0
1.8
1.5
1.2
0.6
0.0
50
100
150
200
50
100
150
200
50
100
150
200
Figure 10: Ablation Study: Added HAE-A representing training VCI with the latent divergence term and without the counterfactual supervision term.
N.3 EVALUATION ON LATENT DISENTANGLEMENT
The definition of disentanglement (Definition 3 is extended from a well-established definition of disentanglement in (Shu et al. (209)), based the oracle consistency (Definition [T) of the latent representation. It a different notion from latent disentanglement in the non-linear ICA literatures (Appendix A) although sometimes sharing non-distinguisaable terminology. this context, the best way to evaluate disentanglement is by definition, using oracle consistency, and we present the ablation results in Table Note that the oracle consistency is measured by KL--divergence for computational efficiency due to the lack of analytic form for the total variational distance between Gaussians. As stated by the Pinsker's Inequality, the total variational distance bounded once KL-divergence controlled.
Table 6: Oracle consistency under different strengths (denoted coefficient W2) on the exogenous disentanglement term.
W2
oracle consistency
3979.86 + 202.42
0.01
0.84 + 0.13
0.02
0.75 + 0.11
0.03
0.53 + 0.10
0.04
0.51 + 0.07
0.05
0.51 + 0.08
N.4 COMPOSITION, EFFECTIVENESS, AND REVERSIBILITY
We examine our model through the composition, effectiveness, and reversibility metrics (Monteiro al. 2023) and use the best benchmark model in Table as a reference. The results are shown in Table Note that these results only serve model inspection purposes and should not be used as strict performance evaluations in favor of Table 2 These metrics are merely surrogate measurements of counterfactual models' axiomatic characteristics, and they are inconclusive in deciding the best- performing counterfactual model. For example, naive model that simply returns the original image as counterfactual outcome at all times is completely useless for counterfactual inference, but will achieve the best score on 2 of these metrics a perfect composition score of 0 and a perfect reversibility score of O. There is trade-off between composition and effectiveness effectiveness and reversibility, and it is not conclusive which counterfactual model performs the best unless there is clear domination of one model over the others all three metrics.
well as
31

---

Published as a conference paper at ICLR 2025
Table Composition, effectiveness, and reversibility of ours against the best benchmark model in 2 Images are scaled to [0. 1] in evaluations. All losses are on the L-2 scale (note that the factors measured by effectiveness thickness and intensity are continuous values).
8
composition
effectiveness
reversibility 4
MED
283,99 67.76
0.01024 + 0.00767
MED 3
0
151.13 = 86.46
0.00816 + 0.00306
SÃE "0.00008 :E 0.000084
T89.03 22.92
0.00102110000000
VCI 0.00026 0.000055
17.64 6.23
0.00076 = 0.00015
it is expected that our model does not reach composition of O because the consistency assumption does not necessarily hold strictly and the model does not necessarily perform exact reconstruction in regions directly affected by the intervention, which is not issue for the purpose of counterfactual generative modeling. Note that the composition metric does not interact with counterfactual interventions all, while the other two metrics interact with counterfactual interventions at least in some capacity.
As discussed in Section 3.3
N.5
MORE DETAILS ON MORPHO-MNIST EXPERIMENTS
Table 8 shows the standard error of metrics reported in Table 2 across five independent runs. Figure I1 shows the violin plot of errors reported in Table 2 and Table across the corresponding runs.
Figure T2 and Figure 13 show some sampled results from VCI.
Table 8: Standard error of metrics.
Std. of Image MSE (-10 3)
Std. of Thickness (th) MAE
Intensity (in) MAE (-10 !)
do(th)
do(in)
mix
do(th)
do(in)
mix
do(th)
do(in)
mix
DEAR
0.20
0.06
0.11
0.18
0.10
0.21
0.32
0.49
0.25
Diff-SCM
0.02
0.02
0.02
0.01
0.01
0.01
0.01
0.02
0.01
CHVAE
0.53
0.86
0.62
0.03
0.15
0.10
0.33
0.32
0.20
CHVAE
0.90
0.83
0.94
0.04
0.21
0.08
0.19
0.10
0.17
MED
0.92
0.26
0.59
0.11
0.08
0.12
0.31
0.19
0.23
MED
1.34
0.40
0.86
0.05
0.06
0.07
0.25
0.09
0.15
SAE
0.12
0.05
0.05
0.10
0.04
0.05
O..1
0.10
0.08
VCI
0.07
0.01
0.07
0.03
0.01
0.05
0.11
0.05
0.06
32

---

Published as a conference paper at ICLR 2025
do(thickness)
do(intensity)
mixed
0.10
0.09
0.08
Image MSE
0.07
0.06
0.05
0.04
0.03
0.01
0.00
1.6
Thickness MAE
1.4
0.6
0.64
0,56
0.48
Intensity
0.40
0.32
0.24
0.16
0.08
0.00
DEAR
CHVAE(B 1)
CHVAE(R 3)
MEDIO = 1)
MEDI = 3)
SAE
VCI
DEAR
CHVAE(SB 1)
MEDIP 1]
SAE
VCI
DEAR
CHVAEIB 1]
CHVAE(BB 3)
SAE
VCI
MEDIB=
MEDIB =
CHVAE(B
MEDIP
Model
Figure I1: Violin plot of Image MSE, Thickness MAE and Intensity MAE on benchmark models and ours across five independent runs.
33

---

Published as a conference paper at ICLR 2025
1
]
1
/
(
2
2
2
2
2
2
2
2
2
2
2
2 2
2
2
2
2
2
2
2
d
2
2
a
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
5
5
5
5
5
5
5
5
5
-
-
5
5
5
5
5
5
5
5
5
5
5
5
5
5
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
(a) do(thickness).
(b) do(intensity).
Figure 12: Sampled results the test of Morpho-MNIST. The left most image of each set is the original image. These results are con-cherry-pickee in terms of model performance (picking was conducted make sure there is variety of digits and styles, but based on performance).
34

---

Published as a conference paper at ICLR 2025
0
0
2
4
5
6
7
8
9
/
0
3
4
5
6
7
8
2
D
/
2
a
4
5
6
7
8
9
3
0
/
2
3
4
5
6
7
8
9
4
O
/
2
4
5
6
7
&
4
5
0
2
3
4
5
6
7
9
6
0
3
4
6
6
3
9
7
0
2
3
4
5
6
7
8
9
8
0
2
3
2
6
7
8
9
9
0
2
3
9
5
6
7
8
9
Figure 13: Sampled results the test of Morpho-MNIST digit intervention. Note that these results do not have quantifiable counterfactual truths which we can evaluate on the capacity of intervening thickness and intensity in Table 2
35


Reviews:
[{"review_id": "KB58aKclqS", "reviewer": "Reviewer_m2tW", "summary": "This paper studies counterfactual generative modeling in terms of neural variational causal inference. The authors claimed that they proposed a new VAE formulation for this purpose with fundamental changes and achieved latent disentanglement through distribution alignment. Several experiments are conducted to demonstrate the author's claim.", "questions": "None", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "The work is theoretically grounded.", "weaknesses": "The manuscript should properly discuss the previous works, such as CEVAE, in the main paper.\n\nFigure 1 could be more specified with additional links that indicate the inference network in the VAEs.\n\nIf the proposed method has novelty or contribution in terms of latent identifiability and disentanglement, appropriate experiments should be conducted with fair comparison against the works within such sub-fields.\n\nWhat I am concerned about is the following. I found that this paper is an extension of a workshop paper where some theoretical results and experimental results are almost duplicated from the previous version. While the recent AI community is rapidly growing and being generous in such situations, I don't think this is the right direction for the development of academia.", "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": "", "extra_scores": {}}, {"review_id": "dkGjENbOXC", "reviewer": "Reviewer_QZf3", "summary": "This work studies counterfactual generation through a newly proposed variational inference paradigm specific to causal models. The proposed Variational Causal Inference (VCI) framework considers the generation of counterfactual outcomes relying on the consistency assumption, where potential outcomes and the true outcomes are the same under a counterfactual treatment. The authors propose a data-generating process from a counterfactual perspective where the latent variables derived from covariates, treatment, and outcome are aligned with latent variables derived from counterfactual outcomes. Utilizing the new counterfactual formulation, an appropriate variational lower bound is derived. The authors show theoretical identifiability guarantees of the proposed formulation. Empirical results show the effectiveness of the VCI framework in generating accurate counterfactuals.", "questions": "- There does not seem to be any explicit modeling of causal mechanisms in the proposed framework. For example, the causal graph thickness → intensity can be imposed on the MorphoMNIST dataset. Previous work, such as DeepSCM [3], CausalHVAE [4], and CausalDiffAE [5], explicitly model the causal mechanism between the two variables. For systems with complex causal relations, can the proposed framework model the causal mechanisms?\n- The notion of disentanglement in this work is quite different than in disentangled representation learning from non-linear ICA based approaches. If T is observed, what is the intuitive meaning behind disentangling the latent factors Z from T? It is not clear to me why enforcing the latent space loss yields a disentangled objective.\n\n[3] Pawlowski et al. Deep Structural Causal Models for Tractable Counterfactual Inference. NeurIPS 2020.\n \n[4] Ribeiro et al. High Fidelity Image Counterfactuals with Probabilistic Causal Models. ICML 2023.\n\n[5] Komanduri et al. Causal Diffusion Autoencoders: Toward Counterfactual Generation via Diffusion Probabilistic Models. ECAI 2024.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Overall, I believe this work sets forth a very interesting formulation for counterfactual inference in generative models and criticizes the use of conditional generative models for counterfactual generation.\n- The VCI framework is formulated well and the consistency assumption to ensure counterfactuals keep individual components of the factual but also reflect changes according to a treatment is intuitive.\n- The theoretical results for ELBO derivation, identifiability, and disentanglement further strengthen the paper and properly contextualize the work within the broader literature of causal generative models.\n- The empirical results show the capability of the proposed framework to generate accurate counterfactuals on several datasets including image and single-cell perturbation datasets.", "weaknesses": "- It seems that the Effectiveness metric from [1] is utilized through the MAE, where an anti-causal predictor is trained to predict the values of thickness and intensity, and a counterfactual reconstruction is fed in for evaluation. I believe this to be a sound metric. However, I do not think the MSE between the true and generated images is necessarily a good metric to evaluate the quality of the counterfactual. The MSE evaluates a pixel-wise similarity between the two images. However, in a counterfactual, the changes are more in the abstract variables describing the image. \n- Although there are theoretical intuitions for disentanglement, there is no strict evaluation of the latent disentanglement of z. The DCI metric [2] evaluates the level of one-to-one correspondence between latent factors and the ground truth generative factors. I am not sure if this metric applies in evaluating disentanglement between z and the treatment variable t, but it would certainly be good to use a similar metric (perhaps Mean Correlation Coefficient) to evaluate disentanglement.\n\n[1] Monteiro et al. Measuring axiomatic soundness of counterfactual image models. ICLR 2023.\n\n[2] Eastwood and Williams. A Framework for the Quantitative Evaluation of Disentangled Representations. ICLR 2018.", "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": "", "extra_scores": {}}, {"review_id": "WtM1A3Z63l", "reviewer": "Reviewer_81Hc", "summary": "A framework is presented in which a counterfactual observation is estimated via a new variational formulation, in which the observed outcome and treatment are used within the conditional distribution for the unobserved (counterfactual) data. Experiments are presented on a range of applications, from biology to various types of images, with impressive quantitative performance.", "questions": "1. In the beginning of Sec. 2.1, \\Omega, \\mathcal{Y}, \\Sigma_Y, \\mathcal{X}, \\Sigma_{\\mathcal{X}}, etc are not defined. This notation is unclear, and not clear why it is needed.\n\n2. On p. 3, Assumptions 1 and 2, Remarks 1 and 2 are mentioned. The reader is very confused. None of these have been mentioned thus far, and the reader is not told such details are in the Appendix. This is unnecessarily confusing. By ICLR guidelines the main paper should be self contained. It is not.\n\n3. On p. 3 it seems you use the adversarial discriminator D(X,T,Y) to manifest the loss. This seems ad hoc. The adversarial discriminator has a connection to the loss, but this is a surrogate and it is not particularly justified. I can see how it might work in practice, but this is an important point that deserves more discussion.\n\n4. In (3), there is inconsistency in using upper case letters (Z for example) and lower case letters (e.g., z) for random variables. The way (2) useses notation on upper/lower case variables is not consistent with (3).\n\n5. Does this really have to be framed as causal analysis? Why is it not just imputation of missing (counterfactual) data? I understand why you want to connect to methods like CEVAE, but they are doing something very different, and hence it seems unnecessary.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The formulation is sound.  Figure 2 is particularly helpful in explaining the model and setup.", "weaknesses": "While the formulation is sound, the terminology of \"causal reasoning\" is questionable. One could argue that you are performing imputation of missing data, and then conditioning on the observed data makes complete sense. The criticism of the CEVAE is inappropriate, because they are solving a very different problem (traditional causal inference), and they do not condition on specific observations (the predictions are not \"personalized\"). It is ok to comment on and compare to CEVAE, but the setup and assumptions are different.\n\nAnother weakness, commented on below, is the presentation. The paper could be much better written, and thereby easier for the reader. Details below.", "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": "", "extra_scores": {}}, {"review_id": "5AHxwyZAUX", "reviewer": "Reviewer_nddM", "summary": "This article presents a variational inference model specifically designed for counterfactual generation, motivated by the goal of extracting information contained within high-dimensional potential outcomes. This concept is intriguing, and the relevant experiments demonstrate the efficacy of the proposed method.", "questions": "1.\tI don’t quite understand the sentence “hence preserving maximum individuality yet minimum treatment characteristics in counterfactual construction  g_theta(y, t')  during inference.” Could you explain it in more detail?\n\n2.\tL841. I have some questions regarding this statement. Why do you think that the performance of intervening with digits cannot be assessed? In fact, I find that evaluating the counterfactual performance of intervening with digits is more intuitive. I am also looking forward to seeing how your model performs with digit interventions.\n\n3.\tIn your paper, you discuss the introduction of a novel ELBO. I am curious about its properties in terms of tightness. Specifically, how does this new lower bound compare to existing bounds in the literature? Additionally, what methods or analyses have you employed to assess its tightness, and how might this impact the overall performance and effectiveness of your model?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThe article tackles a vital issue in counterfactual generation, contributing meaningfully to the understanding of causal relationships in diverse fields.\n\n2.\tThe focus on high-dimensional potential outcomes offers a reasonable angle, as it may capture complexities that traditional methods might miss, thus enhancing counterfactual generation.\n\n3.\tThe experimental outcomes appear encouraging, indicating that the proposed model effectively supports counterfactual generation.", "weaknesses": "1.\tThe article does not provide asymptotic experiments to assess the performance of the proposed ELBO, which limits the understanding of its behavior as the sample size increases and may raise questions about its robustness in practical applications.\n\n2.\tThe absence of a comparison between Morpho-MNIST and diffusion model-based counterfactual generation models Diff-SCM undermines the evaluation's robustness, making it difficult to assess the proposed model's effectiveness relative to state-of-the-art methods and diminishing its perceived impact within the field.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)