Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
Authors: Jiaxuan Zhang
Emadeldeen Eldele
Fuyuan CAO
Yang Wang
Xiaoli Li
Jiye Liang
Venue: ICML 2025 poster (2025)
Keywords: Causal inference
individual treatment effect
counterfactual generation
contrastive learning
Primary area: general_machine_learning->causality
TL;DR: We propose a method for ITE estimation that leverages a derived error bound to ensure fine-grained alignment and robust performance, even with individual heterogeneity.

Abstract: Estimating Individual Treatment Effects (ITE) from observational data is challenging due to covariate shift and counterfactual absence. While existing methods attempt to balance distributions globally, they often lack fine-grained sample-level alignment, especially in scenarios with significant individual heterogeneity. To address these issues, we reconsider counterfactual as a proxy to emulate balanced randomization. Furthermore, we derive a theoretical bound that links the expected ITE estimation error to both factual prediction errors and representation distances between factuals and counterfactuals. Building on this theoretical foundation, we propose FCCL, a novel method designed to effectively capture the nuances of potential outcomes under different treatments by (i) generating diffeomorphic counterfactuals that adhere to the data manifold while maintaining high semantic similarity to their factual counterparts, and (ii) mitigating distribution shift via sample-level alignment grounded in our derived generalization-error bound, which considers factual-counterfactual similarity and category consistency. Extensive evaluations on benchmark datasets demonstrate that FCCL outperforms 13 state-of-the-art methods, particularly in capturing individual-level heterogeneity and handling sparse boundary samples.

Estimating the Individual Treatment Effects (ITE) from ob- servational data is critical for personalized decision decision-making in diverse domains such as digital marketing (Li et al., 2024; Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
Jiaxuan Zhang' Emadeldeen Eldele?
Fuyuan Cao
Yang Wang'
Xiaoli Li Jiye Liang'
Abstract
Estimating Individual Treatment Effects (ITE) from observational data is challenging due to co-
member smmpees
60
contrll tretted
40
variate shift and counterfactual absence. While existing methods attempt to balance distributions globally, they often lack fine-graing mne-grained sample-level alignment, especially in scenarios with significant individual heterogeneity. To address these issues, reconsider counterfactual a proxy to ulate balanced randomization. Furthermore, we derive theoretical bound that links the expected ITE estimation error to both factual prediction errors and representation distances between fac- tuals and counterfactuals. Building on this the- oretical foundation, we propose FCCL, novel method designed to effectively capture the nu- ances of potential outcomes under different treat- ments by (i) generating diffeomorphic counter- factuals that adhere to the data manifold while maintaining high semantic similarity to their fac- tual counterparts, and (ii) mitigating distribution shift via sample-level aligmment grounded in our derived generalization-ermo bound, which consid- ers factual-counterrfatuual similarity and category consistency. Extensive evaluations on benchmark datasets demonstrate that FCCL outperforms 13 state-of-the-an methods, particularly in captur- ing individual-level heterogeneity and handling sparse boundary samples.
20
tszed sead sca anbodding bealay coazel coant mom enbelding cmaed
age
(4) Coninunder bolorce
do Diontarion belanon
1.0
Rareal repaddes
Uestteen ellier
kin sopeedas
0.6
0.2
40
60
70
T-D
(c) lietemganas Incatment effect cximation
Countrritttll gesertion
Figure I. (a) Existing balance methods align overall distribu- tional statisties (e.g., mean) of confounders (e.g., age) between treated and control groups. (b) Maximum Mean Discrepancy (MMD) minimizes group-leve discrepancies by aligning mean representations but fails to capture individual-lers heterogeneity, particularly for boundary samples. (c) Current methods effectively estimate ITE for normal responders but result in higher erross for hyper and low responders, highlighting the need for finer-grained ali alignment. (d) x is generated via gradient ascent until the treat- ment label flips from T to T O, while a represents the diffeomorphic counterfactual generated in the latent space of the normalizing flow, ensuring adherence to the data manifold.
Chu al., 2022), social sciences (Martinez-Sáncchez et al., 2024), and healthcare (Liu et al., 2021; Bica & Van der Schaar, 2022). Accurate ITE estimation guides crucial de- cisions by quantifying the causal impact of interventions. Unlike randomized controlled trials (RCTs), observational studies often suffer from non-random treatment assign- ment mechanism, leading to covariate shift (Cheng et al., 2024; Yao et al., 2021; Kong al., 2023). Accurate estima- tion of treatment effects requires mitigating this covariate shift and predicting counterfactual outcomes-sstt eetimmttiossssssss of what would have happened to an individual under a dif- ferent treatment. Nevertheless, generating realistic coun- terfactuals and achieving a balanced representation remain significant challenges.
1. Introduction
"Shool of Computer and Information Technology, Shanxi Uni- versity, Taiyuan, China Instiittee for InfoComm Research, Agency for Science, Technology and Research, Singapore "Centre for Fron- tier Al Research, Agency Science, Technology and Research, Singapore. Correspondence to: Fuyuan Cao <cty @ sxu.edu.en">
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
Covariate balance methods aim to mitigate distribution shift by aligning the distributions of the treated and control groups (Figure la). This helps balance out the influence of confounders (factors that can affect both the treatment and the outcome). Methods like Maximum Mean Discrepancy (MMD) (Jiang & Sun, 2022) minimize the distribution dis- crepancies by aligning the mean embeddings of the treated and control groups. While generally effective, these meth- ods neglect individual-level heterogeneity (Wu et al., 2023), particularly for boundary samples (Figure lb). This over- sight undermines ITE estimation accuracy and confidence, especially in scenarios with sparse or extreme samples that deviate significantly from population averages.
outcomes under different treatments and enhancing ITE estimation, particularly for boundary and heterogeneous samples (see Figure 3). Furthermore, we provide novel generalization-eri error bound that links ITE estimation error to factual prediction error and the representation distance be- tween factual and counterfactual samples, validating FCCL's efficacy.
Our main contributions are summarized as follows:
We introduce a theoretically grounded diffeomorphic counterfactual as proxy, which ensures counterfactu- als are realistic and semantically meaningful by adher- ing to the data manifold.
For example, in precision medicine, representation learn- ing methods may balance confounder distributions like age (which influences both medication allocation and health out- comes), but fail to capture critical variations in individual drug sensitivity. Patients typically fall into three response categories: normal responders, hyper-responders, and low responders (Figure 1c). While current methods may per- form adequately for the majority (normal responders), they frequently produce significant errors for hyper-responders, who risk adverse effects, and low responders, who may re- quire increased interventions dosage (Feuerriegel et al., 2024). This lack of fine-gras fineigraieed alignment can compromise treatment decisions for these vulnerable groups
We derive theoretical ITE generalization-ene error bound based on factual error and representation distances, and propose FCCL., a novel contrastive learning frame- work for ITE estimation that aligns individual sample representations across treatments.
We conduct a comprehensive evaluation of FCCL, demonstrating its superior performance over CFR- based and adversarial training- methods, with up 25% and 33% reduction in estimation error on IHDP, achieving more accurate and stable ITE estimation.
2. Related Work
To address sample-level heterogeneity, recent approaches emphasize aligning pairs of samples, particularly those in intermediate or boundary regions identified through propen- sity scores (Yao al., 2018; Li & Yao, 2022; Zhao et al., 2024). While promising, these methods achieve partial alignment and still face challenges in effectively mitigating distribution shifts. Generating instance-specific counterfac- tuals offers a potential solution, as these proxies emulate randomized treatment assignment by constructing personal- ized comparisons. However, this strategy is highly sensitive to noise (Cai et al., 2024), often producing adversarial ex- amples rather than meaningful counterfactuals (Figure Id).
Traditional treatment effect estimation methods either rely propensity score or directly optimize sample weight, in- cluding propensity score matching, debiased machine learn- ing, and entropy balancing (Li et al., 2016; Jung et al., 2021; Athey et al., 2018). However, these methods require correct model specification; otherwise, they may lead to unreliable estimates. Recently, numerous representation learning stud- ies focus on learning balanced representations of covariates (Johansson et al., 2022), which can be broadly categorized into covariate balance and adversarial balance approaches.
2.1. Covariate Balance
To overcome these limitations, we propose FCCL (Flow- based Counterfactual Contrastive Learning), a novel frame- work that integrates two key techniques: (1) Flow-based Counterfactual Generation, which utilizes normalizing flows to generate realistic counterfactuals that adhere to the data manifold while minimally altering sample semantics. This ensures that counterfactuals maintain fidelity to real-world plausibility while reversing treatment classification. (2) Sample-level Alignment for Balancing Distribution, which employs contrastive loss to optimize the representation dis- tances between factual and counterfactual samples based on category consistency and similarity.
Covariate balance methods aim to learn a representation space o(.)) that mitigates treatment selection bias. MMD and the Wasserstein metric were first introduced in TARNet (Johansson et al.. 2016) and CFRNet (Shalit et al., 2017) to minimize the distribution differences between treatment and control groups within the representation space, effectively framing causal inference as a transfer learning problem. No- tably, CFRNet provides a generalization-erro bound, show- ing that the expected ITE estimation error can be reduced by the difference between the treated and control distributions. These models inspired numerous representation balancing methods, including (Kazemi & Ester, 2024; Zhang et al., 2020). For example, the Perfect Match method (Schwab et al., 2018) extended the TARNet to scenarios with multiple
This dual strategy in FCCL enables robust sample-level alignment, effectively capturing the nuances of potential
2

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effeet Estimation
treatments and addressed varying sample sizes by augment- ing samples with propensity-matehon nearest neighbors. Huang et al. (2024) derived new ITE error bound based on H(-divergence and proposed DIGNet to mitigate the trade- off problem between the outcome prediction and covariate balance. Besides, DESCAN (Zhong et al., 2022) improved representation learning by jointly learning treatment and response functions over the entire sample space, effectively addressing both selection bias and sample imbalance.
3. Preliminary
3.1. Notations and Assumptions
Following the Neyman- Rubann potential outcome framework (Rubin, 2005; Shalit al., 2017), the covariate space X' is a bounded subset X' C R4, and the potential outcome space is Y RR. The binary treatment indicator & {0, 1} indicates whether the unit receives treatment t (e.g., medication) or serves as control t = 0 (e.g. placebo). We define y as the potential outcome (e.g. blood sugar level) for the unit under treatment t E {0. 1}. the potential outcomes yP: y E y. Given a unit æ E X with treatment assignment I, we only observe the factual outcome y', while the counterfactual outcome y remains unobserved. The observed outcome can be expressed as: y y= (1 - vd + tyy. Our goal is to estimate the ITE and evaluate its accuracy using the met- ric Precision in the expected Estimation of Heterogeneous Effect (PEHE).
While these methods focus on achieving global distribu- tional balance, they often neglect individual-level informa- tion crucial for estimating potential outcomes. SITE (Yao et al., 2018), for instance, attempted to capture sample-level variation by selecting specific sample pairs for alignment. However, these approaches only achieve partial balance and fail to effectively mitigate the distribution shift. This limited alignment prevents robust sample-wise debiasing, which is essential for accurate ITE estimation.
Definition 3.1. The individual treatment effect:
2.2. A Adversarial Balance
T (x) = E yyy - -
When datasets contain finite samples and complex rela- tionships between covariates, neural network-ber feature representations may not be accurate. Adversarial training makes factual and counterfactual distributions indistinguish- able, naturally mitigating distribution bias. GANITE (Yoon al., 2018) utilized adversarial training prevent the dis- criminator from distinguishing between true factual out- comes and generated counterfactual outcomes. SCIGAN (Bica et al., 2020) extended GAN architectures with a hierar- chical discriminator for factual-counterfactus identification and a multi-task generator, enabling outcome estimation for continuous and multi-level discrete interventions via hetero- geneous response curve learning. ABCEI (Du et al., 2021) balanced covariate distributions in the latent space using adversarial training. and addressed information loss with mutual information estimator. Unlike ABCEI, CBRE (Zhou et al., 2022) introduced an information loop to preserve pre- dictive information that might otherwise be lost during the raw-to-latent space transformation in adversarial training. NCMs (Xia et al., 2022) demonstrated that structural con- straints for counterfactual reasoning could be captured and proposed algorithms for joint identification and estimation of counterfactual distributions.
Definition 3.2. Let f f x {0, 1} + y by an hypothesis.
The estimated individual treatment effect:
fy(x) f(x, 1) f(x,
Definition 3.3. The expected Precision in Estimation of Heterogeneous Effect loss of f:
(PEHE(f)= - (f)(z) p(a))d
/
(1)
We made the following assumptions to ensure that treatment effects are identifiable:
Assumption 3.4. Consistency, Ignorability, and Overlap. Consistency: For unit with treatment assignment the ob- served outcome equals potential outcome y². Ignorability: The potential outcomes are independent of the treatment conditioning on covariates, such that (y²,y") _I 1/r. Over- lap: For any x, the probability of receiving treatment is positive. That is, 0 < P(t 1|x) 1, for Vx E X.
3.2. Preliminary Propositions
Assumption 3.5. Manifold hypothesis (Dombrowski et al., 2023). The data is assumed concentrate within a small region d around submanifold M of X:
Although adversarial training methods can learn counterfac- tual distribution, they face significant challenges in generat- ing realistic counterfactuals. Naively adding noise to inputs often produces adversarial examples rather than semanti- cally meaningful counterfactuals. In contrast, our approach generates counterfactuals that preserve semantic integrity, respecting the intrinsic structure of the data manifold. This ensures that the generated counterfactuals are both mean- ingful and reliable for causal inference tasks.
S = M x Is. X ... X ISNZ NM²
(2)
where Ia, = (- 2 4) for i E (I,...,
-Mr. is
an open interval of length ò, We assume that 6, is small, meaning the data lies approximately on the submanifold M C X, which has a much lower dimensionality NM than the dimensionality Nx of X.
3

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
Factuals
Cibi(og(=))
Engoder
P
Projection
hy
Bhared
Shured
Ce
(())
Factual counterparts
gg
Projection
Cy ah.
sample
cuntrol sample
treated
control counterfactual
minimize
Faciuals z
Normalizing New Represeniation
Comierfactual generation
Couneetbciulss
Factuals
Counterfactuals
Figure 2. The architecture of the proposed FCCL. framework. The diffeomorphic counterfactual generation module generates realistic counterfactuals via gradient ascent in the latent space of normalizing flow, adhering to the data manifold. The contrastive learning module refines sample representations under different treatments, enabling precise instance-level distribution alignment. Finally, two separate neural networks, ft (d)(x)) and (x)) are used estimate potential outcomes under different treatments.
Definition 3.6. Diffeomorphic counterfactual. The diffco- morphic counterfactual is defined as minimal modifications on the data manifold of the input reversing the classification decision & (z), represented as:
4. Methodology
We propose a novel framework, FCCL (Flow-based Counterfactual Contrastive Learning), for individual treat- ment effect estimation. FCCL grounded in theoretical insights that justify its key components: (1) flow-based coun- terfactual generation, (2) contrastive learning for sample- level alignment, and (3) predictive model for potential outcomes. Figure 2 provides an overview, integrating the theoretical foundation with methodological details of FCCL.
I = arg min dist 2) s.t. k(T) + k(x), ZeM
where M is the data manifold, dist(E, 2) denotes the dis- tance between 전 and x. The classifier & X + {0, 1} as- signs an imput X' to a class E (0, 1}, and k(z) * k(x)) indicates change in the classifier's prediction outcome.
Motivation. Randomized controlled trials (RCTs) are rec- ognized as the gold standard for causal inference due to their ability to achieve random treatment assignment independent of covariates (Ma & Zhang, 2023). Our method aims to emulate the balanced randomization inherent in RCTs by employing counterfactual contrastive learning. To the best of our knowledge, this the first work to leverage the re- lationship between factual and counterfactual samples to produce consistent representations that effectively capture the characteristies of potential outcomes under different treatments, thereby addressing distribution shifts between treated and control groups through fine tne-grained, sample- level alignment.
Definition 3.7. Let J. ) x y - R be the absolute loss or squared loss, 42, t) denote the expected loss for the unit-treatment pair tz, with representation func- tion +- X R and hypothesis which can be mathematically represented as: 1n,0 (2, t) Jy. L(v', ((((x), ())r) ())(())"(aa (x)dy'' The expected factual and counterfactual losses of h and &b are:
R. x {0, - V,
cr(h, f lx, t)p(x) t)drdt,
|xx(0.1)
ECF(h, ④) / [xx(o.1) (n, 1 t)dxdt.
4.1. Theoretical Foundations
To motivate our methodology, we first establish theoretical bounds on counterfactual loss and ITE estimation error. These bounds guide the design of FCCL and justify its components. Detailed proofs are provided in Appendix A.
Definition 3.8. The expected factual treated and control
losses are:
/ (n, 1)p² 1)p==1(x)d (x)
cF t (),
Lemma 4.1. Let + X + R be an invertible represen- tation function with inverse and h R x {0, 1] + y a hypothesis function. Define u = p(t = 1) as the treat-
/ 4n (2, 0)p²=0 (x)dx.
oF (1, 4)

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
ment proportion. For any treatment indicator & E {0,1}, let dãsc (ass((((()) (x), (3)) represent the distance between coun-
the underlying data manifold M that represents the intrinsic structure and variations of the data.
terfactual and factual samples in the representation space, Then, the counterfactual loss EC (h.
To address this limitation, we generate diffeomorphic coun- terfactual, ensuring that counterfactual F resides on the same manifold M as the original instance x; i.e., I;, 3 E M. This is formulated as an optimization problem:
4) is bounded by:
ECF (h, 2) (3 (1 (1- 4 (h, 1) + u "h, ()+
(4)
disc 4 (x)) 2' (5)) + disc 4 (3), x (x)),
min Da ă) s.t. k()) + (()),
(7)
where ECF and cp are defined in Definition 3.7 and Defini- tion 3.8.
where Das denotes the geodesic distance the manifold .M and k(-) is classifier function. By constraining E to remain on M, this formulation avoids unstructured pertur- bations while flipping the classification outcome.
This result highlights the importance of minimizing rep- resentation distances between factual and counterfactual samples to reduce counterfactual loss.
Normalizing flows provide a principled mechanism to con- struct such diffeomorphic counterfactuals. In particular, a normalizing flow maps base distribution p(z) in the latent space Z to complex target distribution p(z) via diffeo- morphism g: x = g()). The normalizing flow is trained by minimizing the Kullback-Leible (KL) divergence be- tween the true data distribution and the model distribution: (()) = DKL [p(=)]lp(z; 0]].
Theorem 4.2. Let EPEHE (h, Φ) be the estimation error with representation function 4 X + R with y being its inverse and hypothesis h: R x {0, 1) + V. For any treat- ment indicator t E {0, 13. let (isc(Q'(z) (x), (²)) be the representation distance between counterfactual and factual samples. The error is bounded by:
(PEHE(h, $) S
2( (h. 4) + (h, +) -20 2(disc( (2), क6 (i)) + disc @5 (#)) क5 (x)))
(5)
Using the learned flow g. counterfactuals are generated by optimizing the composite function k o g Z + y in the latent space Z to find latent representation غ that produces the desired classification outcome (Dombrowski et al., 2023). The iterative update rule in the latent space is given by:
This bound links ITE estimation error to the generalization error of factual predictions and representation distances, motivating our focus on minimizing these distances.
(++1) 100==00 =200) + a(k g), (2())
This theorem proves that the estimation error EPENE (h, 4) is upper bounded by two terms: the standard generalization error of factual predictions eF and the distance constraints in the representation space. These distance constraints rep- resent two types of alignment: (i) between counterfactual samples and their factual counterparts (similarity), and (ii) between counterfactual samples and those from the oppo- site group of factual samples (category consistency). These alignment errors can be minimized using contrastive loss to achieve sample-wise correspondence in the representation space, improving the accuracy of ITE estimation.
(8)
Oz
with step size A. This approach leverages the manifold structure encoded by g. ensuring that counterfactuals ă = g(3) adhere to the data manifold M. Consequently, the generated counterfactuals remain semantically meaningful, exhibit minimal deviation from the original samples, and achieve a classification flip.
Thus, diffeomorphic counterfactuals provide a theoretically sound foundation for generating meaningful sample pairs (r, .), which are subsequently utilized in the representa- tion module to achieve robust alignment and accurate ITE estimation.
4.2. Flow-based Counterfactual Generation
Existing counterfactual generation approaches typically em- play gradient ascent optimization in the input space X. Specifically, for a given step size 7) and target class f, the iterative gradient update rule is defined as:
4.3. Contrastive Learning for Sample-level Alignment
Counterfactual Contrastive Learning. For each mini- batch of training samples, we generate a diffeomor- phic counterfactual ă; for each instance 24. These fac- tual-counterfactual pairs (2,, (z); are treated as positive pairs, while instances from different samples form the can- didate set for negative pairs. To maintain computational efficiency, a fixed number of negative samples are randomly selected from this set. This sampling strategy ensures suf- ficient contrastive information while minimizing computa- tional overhead.
(1+1) (1) + " Ox
Tk:
(z. (2)))
(6)
where the process continues until the classifier ((x(++))) confidence exceeds a threshold A.
However, counterfactuals generated in this manner often result in unstructured perturbations that fail to capture mean- ingful semantic transformations, frequently deviating from
5

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
Each instance 2; and its corresponding counterfactual 2e processed through an encoder network, and the resulting representations &E and dg are processed by the projection head network to yield and ?, which facilitates the applica- tion of contrastive loss while enabling the encoder to retain more information in its learned features (Ho & Nvasconce- los, 2020). The contrastive loss is applied to learn invariant representations of G and & are closely aligned, effectively capturing the nu- ances of potential outcomes under different treatments. The contrastive loss function is defined as:
where and 8 are adjustable hyper-paramese that control the contributions of contrastive loss and regularization loss Il-l12 on model weights W to prevent overfitting.
We train our model by minimizing Equation (11) and pro- vide the detailed counterfactual contrastive learning strategy in Algorithm in the Appendix. This formulation ensures that the learned representations align factual and counterfac- tual samples while accurately predicting potential outcomes. ultimately reducing ITE estimation error.
pairs ensuring that
4.4. Why FCCL Works
exp (sim (c),2) 24) Ttemp)
Le -log Na
FCCL bridges theoretical guarantees and practical innova- tion to address key challenges in ITE estimation, partieularly for sparse and boundary samples. By generating diffeomor- phic counterfactuals that adhere to the data manifold, FCCL ensures that counterfactuals are both realistic and mean- ingful. This design reduces noise sensitivity and preserves semantic consistency, which critical for obtaining reliable sample pairs. Furthermore, contrastive learning enforces dis- tance constraints based on factual-counterfacte similarity and category consistency (as formalized in Lemma 4.1), mit- igating distribution shifts through fine- ane-grained sample-level alignment, and approximating the randomization achieved in RCTs. Together, these mechanisms reduce ITE estima- tion error, which is consistent with Theorem 4.2, and out- perform traditional methods in scenarios with significant heterogeneity limited data.
jiteep (sim (ci, (5) Ttemp))
[[e |e represents the cosine simi- larity, N₁ is the set of batch-size randomly selected negative samples, and Ttemp is the temperature coefficient.
where sim(c,, Q))
The contrastive loss aligns directly with the theoretical bounds established in Lemma 4.1, which reflects the repre- sentation distances between counterfactuals and their factual counterparts, as well as between counterfactuals and factual samples from the opposite group. By minimizing these rep- resentation distances, the counterfactual loss ECF (h, 4) re- mains bounded, detailed in Equation (4). The contrastive learning module achieves robust alignment, effectively re- ducing the error terms identified in Theorem 4.2:
5. Experiments
disc 46 (z), (i)) and disc( (isc(((().0. x (x)).
To evaluate the effectiveness of FCCL, particularly its ro- bustness under varying degrees of covariate dispersion, we employ six benchmark datasets: four synthetic, one semi- synthetic, and one real-world datasets. Comprehensive ex- periments across various datasets validate our claims, in- cluding performance evaluation, latent space analysis and boundary sample examination.
These alignment errors, which quantify the distances be- tween factual- counterfcuttr pairs and between counterfac- tual samples and opposing group samples, directly impact the ITE estimation error. By minimizing fe, the learned representations align with the theoretical bounds, improving the reliability of ITE estimation.
Prediction Head. The learned balanced representations
5.1. Datasets
(x,) are fed into two neural networks to predict poten- tial outcomes for treatment (t = 1) and control (t = 0) (Assaad et al., 2021; Huang et al., 2023). The predicted outcomes are defined as Tour = h2,) = 1) and Cout = h( = 0), respectively. The predictive loss given by:
Synthetic: We generate covariates from multivariate normal distribution N' (0,7-0) o² [pl,1] + (1 - -))I, where the scaling parameter E (0.4. 1.0, 1.2) con- trols the degree of covariate dispersion, as shown in Figure 7. Each dataset consists of 800 units split into training (63%), validation (27%). and test (10%) sets. We generate 30 in- dependent datasets for analysis. Detailed data generation steps are outlined in Appendix D.3.
C
Σ
wi L (dr (4)
(10)
11.
where w;
f1 + 144
211 and u 1 Einh 4. The to-
The Infant Health and Develop- ment Program (IHDP) dataset, modified by Hill (Hill, 2011), evaluates the effect of specialist home visits on children's 'ss cognitive scores. It includes 25 covariates describing infants and their mothers. Selection bias is simulated by systemati-
Semi-sytthetic (IHDP):
2u
tal loss L, combines predictive loss, contrastive loss, and regularization:
L. = C. + +afa + 1|,
(11)
6

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
Table 1. Within-sample and out-of-sample mean and standard er- rors for the metrics (Lower is better) on IHDP dataset.
Table 2. Out-of-sample mean and standard errors for the rooted
PEHE (V "EEH Synthetic datasets.
Method
within PEBE
withhn ATT
nut-of PEEE
and of CAE
Method
7 = 0.4
Y 0.7
7 1.0
7 1.2
OLS-I
5.83(..39)
0.73().04)
5.91(..27)
0.95(0.06)
OLS-1
8.41(0.84)
10.85(1.34)
13.00(1.68)
14.32(1.55)
OLS-2
1.4
0.14(0.02)
1.)
0.31(0.02)
OLS-2
5.94(0.60)
7.64(0.96)
1.13(1.1))
10.07(1.11)
BART
2.13(0.22)
0.24(0.05)
2.3
0.35(0.03)
BART
3.40(0.50)
4.17(0.60)
4.86(0.61)
6.14(0.58)
KNN
2.13(0.08)
0.15(0.05)
4.16(0.23)
0,80(0005)
KNN
5.50(0.62)
7.26(0.75)
9.08(1.27)
10.27(1.07)
DML
2.45(0.12)
0.20(0.05)
2.60(0.14)
0.33(0.05)
DML.
3.96(0.54)
5.30(0.62)
6.08(0.80)
8.46(0.77)
CFR- Wass
0.71(0.04)
0.27(0.03)
0.83(0.08)
0.28(0.03)
CFR-Wast Wass
2.48(0.06)
3.60(0.09)
4.72(0.14)
5.37(0.14)
CFR-MMD
0.77(0.05)
0.25(0.04)
0,92(0.09)
0,28(0.04) 04)
CFR-MMD
2.54(0.06)
3.62(0.09)
4.74(0.14)
(0.1)))4)
SITE
0.84(0.05)
0.30(0.04)
0.98(0.07)
0.32(0.05)
SITE
2.69(0.13)
4.25(0.23)
6.01(0.38)
6.21(0.17)
CITE
2.71(0.07)
3.70(0.10)
4.74(0.14)
14)
CITE
0.59(0.06)
0.11(0.02)
0.67(0.14)
0.14(0.02)
GANITE
4.69(0.06)
6.16(0.07)
7.33(0.07)
8.11(0.08)
GANITE
1.92(0.29)
0.43(0.41)
2.43(0.46)
0.49(0.38)
ABCEI
2.75(0.06)
3,57(0.09)
4.73(0.13)
5.19(0.12)
ABCEI
0.79(0.06)
0.12(0.02)
1.00(0.13)
0.15(0.03)
CBRE
2.93(0.05)
3,85(0.08)
5.02(0.12)
5.73(0.13)
CBRE
0.59(0.06)
0.11(0.02)
0.66(0.07)
0.13(0.02)
DIGNet
3.18(0.11)
3,97(.13))
5.100.17)
5.81(.115)
DIGNet
0.60(0.04)
0.15(0.02)
0.67(0.07)
0.16(0.02)
FCCL
2.58(0.06)
3.50(0.09)
4.49(0.13)
5.12(0.12)
FCCL
0.53(0.04)
0.09(0.01)
0.64(0.07)
0.12(0.02)
et al., 2017) are two methods using the Wasserstein and MMD metric for counterfactual regression, respectively; SITE ( (Yao et al., 2018), which preserves local similarity in sample representations; CITE (Li & Yao, 2022) learns representation based on propensity score; GANITE (Yoon et al., 2018) implicitly learns counterfactual distribution us- ing GANs; ABCEI (Du et al., 2021) balances distributions using adversarial learning; CBRE (Zhou et al., 2022) con- structs an information loop during adversarial training to minimize information loss; DIGNet (Huang et al., 2024) utilizes individual propensity confusion and group distance minimization to learn the balanced representation.
cally excluding subset of treated samples, resulting in 747 instances (139 treated and 608 control). We use the same 100 datasets, following the standard practice in the field.
Real-world (Jobs): This dataset examines the causal ef- feet of job training (treatment) on income and employment status. It is reformulated
binary classification task to predict unemployment (Dehejia & Wahba, 2002). This dataset contains 297 treated samples and 2,915 control sam- ples. For consistency with prior work, we average over 10 experiments. The train, validation, and test splits are set to 56%, 24%, and 20%, respectively.
5.4. Experimental Results
5.2. Metrics
We evaluate FCCL's performance on multiple metries across various datasets, focusing on robustness under varying co- variate dispersion. Additionally, we conduct latent space analysis and examine boundary sample behavior to provide deeper insights into FCCL's effectiveness. Further results, including sensitivity analysis, are presented in Appendix D.
For datasets with known true treatment effects, we evaluate using the rooted Precision in Estimation of Heterogeneous Effect VEPEHE and the absolute error of Average Treatment Effect CATE. For the Jobs, where ground-truth counterfactu- als are unavailable, we adopt the policy risk Repo (TF)) and the bias of Average Treatment Effect on the Treated predic- tion CATT. Detailed formulas are provided in Appendix C.
We compare our FCCL against the baseline methods on IHDP, Jobs and Synthetic datasets, with results presented in Table I and Table 2, and additional results are provided in the Appendix. Key findings include: (I) FCCL outperforms all baseline methods across various datasets. Notably, on the IHDP dataset, FCCL achieves the lowest VEPEHE and ATE values, reducing out-of by 57. 1% compared to CFR. (2) Methods leveraging contrastive learn- ing (c.g., CITE, FCCL) consistently outperform CFR-based methods, highlighting their strength in aligning representa- tions under diverse treatments. (3) FCCL shows resilience to covariate dispersion, as it maintains robust performance as the covariate dispersion (7) increases from 0.4 to 1.2, out-of with V EPEHE increasing from 2.58 to 5.12. This increase is 12. lower compared to CFR-Wass.
Performance Evaluation:
5.3. Baseline Approaches
We compare our FCCL against thirteen baselines, catego- rized as traditional and deep learning methods.
Traditional Methods: Ordinary least square (OLS-1), which uses treatment as a covariate to predict outcomes: (OLS-2), which predicts outcomes separately for each group; Bayesian additive regression trees (BART) lever- aging sum-of-trees structure; K-nearest neighbor (KNN) matches samples using k-nearest neighbors; Debiased ma- chine learning (DML), handling confounding bias through orthogonal residual regression. Deep learning Methods: CFR- Rass (Shalit et al., 2017) and CFR-MMD (Shalit
7

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
Table 3. ITE estimation errors with different counterfactual gener- ation strategies on IHDP dataset.
custrvd
coafril frested trasted hds.
7.5
frested
costroc
7.5
54
5.8
Method
withhn PERL
within ATT
nut-nf PEEEE
out-of
CATE
2.5
2.5
grad ase in 2''
0,60(0.06)
0.11(0.02)
0.15(0.04) 0.87(0.19)
0.0
GAN
0.57(0.05)
0.10(0.01)
0.14(0.03) 0.84(0.17)
-1.5-
-.5
diffusion
0.54(0.05)
0.09(0.01)
0.72(0.10)
0.12(0.03)
CFR
FCCL
FCCL
(.53().004
0.09(0.01)
0.64(0.07)
0.12(0.02)
-1.4
dis 15.44
-35
316
-18.9
Kl-1114
-18.9
KL
(4) Repersentation FF
Representation
samples CFR)
1.75
pos: 0.7089
1.50
1.6762
18.0
centrol
10.0
coatral
diszance
7.5
centrol Inds
7.5
condred Ids
5.8
treated Ids
5.8
trasted bdt
0.75
1.5
3.5
0.50
0.25
-1.5
-.5
center to center sample to center (a) The impact of contrastive loss
dis
dis-inn dis-sum
-500
-5.8
(b) Key dintance metrics
ABCEI
PCCL 0,2
-184
-18.8
Figure 4. (a) The impact of contrastive loss (CL) on distribution balance, comparing distances between class centroids (center-to- center) and between samples and the center of the opposite class (sample-to-center with and without CL. (b) Key distance met- rics in the representation space: factual-counterfaatuull (dis-pos), counterfactual-oppoiii factual (dis-inn), and factual-factual (dis- sam), shown with mean values for each metric. Contrastive learn- ing reduces distances, enhancing representation alignment and mitigating distribution discrepancies.
-31
-35
10
15
-20
-15
-10
10
(1) Reprovrntation ABCCE
(4) Reprovrntation of F. (Boonddryy comples from ABCE.)
Figure 3. Latent Distributions of two representative methods, CFR and ABCEI, alongside our FCCL on IHDP dataset. The dots in black and yellow represent the boundary samples from the treat- ment and control groups, respectively. The metric dis quantifies the average distance between boundary samples and correspond- ing class centers, reflecting sample heterogeneity. The metric KL divergence is used to quantify the difference between treatment and control distributions.
5.5. Components Analysis
This section explores the contribution of two key compo- nents of con.--oounerfacttulllll generation and contrastive learning-toware improving ITE estimation performance.
Latent Space Representation and Boundary Sample Analysis: Figure 3 visualizes the latent space represen- tation for CFR (covariate balance-based), ABCEI (adversar- ial baaiin,, and our proposed FCCL. method (con- trastive soaming-based). FCCL generates well-balanced and compact representations and achieves lower KL divergence than CFR and ABCEI, significantly reducing distribution discrepancies between treatment and control groups through robust sample-level alignment. Additionally, boundary and heterogeneous samples limit the model's generalizability, especially in scenarios with substantial individual hetero- geneity. A smaller "dis" value indicates that dispersion of samples in the latent representation space is reduced, which shows that there are fewer boundary samples, en- abling better model fitting and exhibiting lower ITE estima- tion bias. Motivated by this, FCCL focuses on boundary damples-defined as the 30 farthest samples from their class centers-and achieves: (1) reduced dispersion: as FCCL aligns boundary samples closer to their class centers, outper- forming CFR (dis- 15.74) and ABCEI (dis = 13.79); and (2) improved ITE estimation: analyzing ITE estimation bias (Tables 4 and 5 in Appendix D) shows that FCCL achieves lower estimation bias, for boundary samples compared to baselines.
Effectiveness of Counterfactual Generation: To validate the efficacy of the flow-based counterfactual generation, we compare it with alternative strategies, including gradient ascent in the X'ssaae, GAN-based generation (Goodfellow et al., 2020), and diffusion- based generation (Kotelnikov et al., 2023). As shown in Table 3, FCCL sigmificantly out- performs these methods, achieving a 26.4% reduction in out-of EPEHE compared to the gradient ascent in the X' d'space. While the diffusion model achieves performance compara-
ble to FCCL, yields slightly inferior result on / e This may be due to the noise-driven generation mechanism of diffusion models, which causes counterfactual to devi- ate from the sample semantic space, especially in the out- sample cases. Unlike other approaches, which introduce noise and fail to preserve the inherent structure of the data, FCCL generates diffeomorphic counterfactuals that remain consistent with the data manifold that represents the intrinsic structure of the data. This structural consistency enhances causal effect estimation, making FCCL more reliable for individual-level treatment effect predictions.
8

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effeet Estimation
Impact of Contrastive Learning: Contrastive learning plays critical role in FCCL's robust performance by align- ing representations of treated and control groups. As il- lustrated in Figure 4, contrastive loss reduces the class- center distances (center-to-center) and optimizes the dis- tances between: (1) counterfactual and corresponding fac- tual samples: where the mean positive distance (dis-pos) is 0.7089: (2) counterfactual and opposite-class factual sam- ples, where the mean inner-class distance (dis-inn) is 1.2180. This alignment effectively balances the treated and control distributions, achieving sample-level correspondence and mitigating distribution discrepancies. By ensuring category consistency and factual-countere coutterccoutt similarity, FCCL achieves superior representation balance, as theoretically justified in Section 4.4.
References
Alaa, A. M. and Van Der Schaar, M. Bayesian inference of individualized treatment effects using multi-tase gaussian processes, In Proceedings of the Advances in Neural Information Processing Systems, pp. 3427-3435, 2017.
Assaad, S., Zeng, S., Tao, C., Datta, S., Mehta, N., Henao, R., Li, F., and Carin, L. Counterfactual representation learning with balancing weights. In Proceedings of the International Conference on Artificial Intelligence and Statistics, pp. 1972-1980, 2021.
Athey, S., Imbens, G. W., and Wager, S. Approximate resid- ual balancing: debiased inference of average treatment effects in high dimensions. Journal of the Royal Sta- tistical Society Series B: Statistical Methodology, 80(4): 597-623, 2018.
6. Conclusion
In this paper, we proposed FCCL, a robust method for ITE estimation, grounded in our derived ITE estimation error bound. FCCL innovatively integrated diffcomorphic coun- terfactual generation and contrastive learning to address dis- tribution shifts between treated and control groups through fine-grained, sample-level alignment. Comprehensive periments across various datasets, including performance evaluation, latent space analysis and boundary sample ex- amination, have demonstrated that FCCL. achieves more accurate and robust ITE estimation, particularly in scenarios with significant individual heterogeneity. For future direc- tions, we aim to extend our framework accommodate time-dependent outcome variables.
Bica, I. and Van der Schaar, M. Transfer learning on hetero- geneous feature spaces for treatment effects estimation. In Proceedings of the Advances in Neural Information Processing Systems, pp. 37184-37199 2022.
Bica, I., Jordon, J., and Van der Schaar, M. Estimating the ef- fects of continuous-valued interventions using generative adversarial networks. In Proceedings of the Advances in Neural Information Processing Systems, pp. 16434 16445, 2020.
Cai, R., Zhu, Y., Qiao, J., Liang, Z., Liu, F., and Hao, Z. Where and how to attack? a cuusality-inspiree recipe for generating counterfactual adversarial examples. In Proceedings of the AAAI Conference on Artificial Intelli- gence, pp. 1 11132-111 140, 2024.
Impact Statement
FCCL represents a significant advancement in causal ef- feet estimation by providing a robust and accurate method capable of handling complex scenarios with substantial indi- vidual heterogeneity. Its ability achieve precise alignment at the sample level makes it particularly suited for applica- Potential domains
Cheng, D., Li, J., Liu, I.., Zhang, J., Liu, J., and Le, T. D. Local search for efficient causal effect estimation. IEEE Transactions on Knowledge and Data Engineering, 35 (9):8823-8833) 2022.
tions in personalized decision-making of impact include digital marketing for tailoring customer strategies, social sciences for policy evaluation, and health- care for personalized treatment planning.
Cheng, G., Hardt, M., and Mendler-Dünner, C. Causal in- ference out of control: Estimating performativity without treatment randomization. In Proceedings of the Interna- tional Conference on Machine Learning, pp. 8077-8103, 2024.
Acknowledgements
Chu, Z., Ding, H., Zeng, G., Huang, Y., Yan, T. Kang, Y., and Li, S. Hierarchical capsule prediction network for marketing campaigns effect. In Proceedings of the ACM International Conference on Information & Knowledge Management, pp. 3043-300 3051, 2022.
This work is supported by the National Natural Science Foundation of China (U24A20323, 62376145), the Science and Technology Innovation Talent Team of Shanxi Province (20220405 101200005002016) (00016)) and the Key Technologies Program of Taihang Laboratory in Shanxi Province (THYF-JSZX- 24010700).
Dehejia, R. H. and Wahba, S. Propensity score-matching methods for nonexperimental causal studies. Review of Economics and Statistics, 84(1):151-16 2002.

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effeet Estimation
Dombrowski, A.-K., Gerken, J. E., Müller, K.-R., and Kessel, P. Diffeomorphic counterfactuals with gener- ative models. IEEE Transactions Pattern Analysis and Machine Intelligence, 46(5):3257-322) 2023.
Jung, Y., Tian, J., and Bareinboim, E. Estimating identifi- able causal effects through double machine learning. In Proceedings of the AAAI Conference Artificial Intelli- gence, pp. 12113-12122, 2021.
Du, X., Sun, I.,, Duivesteijn, W., Nikolaev, A., and Pech- enizkiy, M. Adversarial balancing-based representation learning for causal effect inference with observational data. Data Mining and Knowledge Discovery, 35(4): 1713-1738, 2021.
Kazemi, A. and Ester, M. Adversarially balanced repre- sentation for continuous treatment effect estimation. In Proceedings of the AAAI Conference on Artificial Intelli- gence, pp. 13085-13093, 2024.
Kong, I., Park, Y., Jung, J., Lee, K., and Kim, Y. Covariate balancing using the integral probability metric for causal inference. In Proceedings of the International Conference on Machine Learning, pp. 17430-17461, 2023.
Feuerriegel, S., Frauen, D., Melnychuk, V., Schweisthal, J., Hess, K., Curth, A., Bauer, S., Kilbertus, N., Kohane, I. S.,, and Van der Schaar, M. Causal machine learning for predicting treatment outcomes. Nature Medicine, 30 (4):958-968 2024.
Kotelnikov, A., Baranchuk, D., Rubachev, I., and Babenko, A. Tabddpm: Modelling tabular data with diffusion mod- els. In Proceedings of the International Conference on Machine Learning, pp. 17564-17579, 17579, 2023.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial networks. Communications of the ACM, 63(11):139-14444 144,144,1114 2020.
Li, S., Vlassis, N., Kawale. J., and Fu, Y. Matching via di- mensionality reduction for estimation of treatment effects in digital marketing campaigns. In Proceedings of the International Joint Conference on Artificial Intelligence, pp. 3768-3774, 2016.
Hill, J. I.. Bayesian nonparametric modeling for causal infer- ence. Journal of Computational and Graphical Statistics, 20(1):217-240, 2011.
Li, X. and Yao, L. Contrastive individual treatment effects estimation. In Proceedings of the International Confer- ence on Data Mining, pp. 1053-1058, 2022.
Ho, C.-H. and Nvasconcelos, N. Contrastive learning with adversarial examples. In Proceedings of the Advances in Neural Information Processing Systems, pp. 17081- 17093, 2020.
Li, Y., Leung, C. H., Sun, X., Wang, C., Huang, Y., Yan, X., Wu, Q., Wang, D., and Huang, Z. The causal impact of credit lines spending distributions. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 180- 187, 2024.
Huang, Y., Leung, C. H., Ma, S., Yuan, Z., Wu, Q. Wang. S., Wang, D., and Huang, Z. Towards balanced representa- tion learning for credit policy evaluation. In Proceedings of the International Conference on Artificial Intelligence and Statistics, pp. 3677-3692, 2023.
Liu, R., Wei, I.., and Zhang, P. A deep learning framework for drug repurposing via emulating clinical trials on real- world patient data. Nature Machine Intelligence, 3(1); 68-75, 2021.
Huang, Y., Siyi, W., Leung, C. H., Qi. W., Dongdong, W., and Huang, Z. Dignet: Learning decomposed patterns in representation balancing for treatment effect estimation. Transactions on Machine Learning Research, 2024.
Ma. C. and Zhang, C. High precision causal model evalua- tion with conditional randomization. In Proceedings of the Advances in Neural Information Processing Systems, pp. 52615-52627 2023.
Jiang, S. and Sun, Y. Estimating causal effects on networked observational data via representation learning. In Proceed- ings of the ACM International Conference on Information & Knowledge Management, pp. 861, 2022.
Hartinez-Sánchez. Á., Arranz, G., and Lozano- Duránn, A. Decomposing causality into its synergistic, unique, and redundant components. Nature Communications, 15(1): 9296, 2024.
Johansson, F., Shalit, U., and Sontag, D. Learning represen- tations for counterfactual inference. In Proceedings of the International Conference on Machine Learning, pp. 3020-3029, 2016.
Reddy, A. G. and Balasubramanian, V. N. Nester: An adap- tive neurosymbolic method for causal effect estimation. In Proceedings of the AAAI Conference on Artificial In- telligence, pp. 14793-14933 14801, 2024.
Johansson, F. D., Shalit, U., Kallus, N., and Sontag, D. Generalization bounds and representation learning for es- timation of potential outcomes and causal effects. Journal of Machine Learning Research, 23(166): 2022.
Rubin, D. B. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association, 100(469):322-333, 2005.
10

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
Schwab, P., Linhardt, L., and Karlen, W. Perfect match: A simple method for learning representations for counter- factual inference with neural networks. arXiv preprint arXiv:/810.00656.. 2018.
Zhou, G., Yao, L.,, Xu, X., Wang, C., and Zhu, L. Cycle- balanced representation learning for counterfactual infer- ence. In Proceedings of the 2022 SIAM International Conference on Data Mining, pp. 442-440200 2022.
Shalit, U., Johansson, F. D., and Sontag, D. Estimating individual treatment effect: generalization bounds and al- gorithms. In Proceedings of the International Conference Machine Learning. pp. 3076-3000000 3085, 2017.
Zhu, F., Lu, J., Lin, A., Xuan, J., and Zhang, G. Direct learning with multi-task neural networks for treatment effect estimation. IEEE Transactions on Knowledge and Data Engineering, 35(3):2457-244 2470, 2021.
Wu, A., Yuan, J., Kuang, K., Li, B., Wu, R., Zhu, Q., Zhuang, Y., and Wu, Learning decomposed representa- tions for treatment effect estimation. IEEE Transactions on Knowledge and Data Engineering, 5585),999-,,00, 2022.
Wu, A., Kuang, K., Xiong, R., Li, B., and Wu, F. Stable estimation of heterogeneous treatment effects. In Pro- ceedings of the International Conference on Machine Learning, pp. 37496 37496-37510, 2023.
Xia, K., Pan, Y., and Bareinboim, E. Neural causal models for counterfactual identification and estimation. arXiv preprint ar Xix::220000000:: 2022.
Yao, L., Li, S., Li, Y., Huai, M., Gao, J., and Zhang, A. Rep- resentation learning for treatment effect estimation from observational data. In Proceedings of the Advances Neural Information Processing Systems, pp. 2638 2648, 2018.
Yao, L., Chu, Z., Li, S., Li, Y., Gao, J., and Zhang, A. A survey on causal inference. ACM Transactions on Knowledge Discovery from Data, 15(5):1 2021.
Yoon, J., Jordon, J., and Van Der Schaar, M. Ganite: Estima- tion of individualized treatment effects using generative adversarial nets. In Proceedings of the International Con- ference Learning Representations, 2018.
Zhang, Y., Bellot, A., and Schaar, M. Learning overlap- ping representations the estimation individualized treatment effects. In Proceedings of the International Conference on Artificial Intelligence and Statistics, pp. 1005-1014, 2020.
Zhao, Y., Yang. F., Yang. X., Teng. J., Zhang, X., and Li, Q. Self-supervised representation matching model for treatment effect estimation. In Proceedings of the Inter- national Conference on Medical Artificial Intelligence, pp. 607-616, 2024.
Zhong, K. Xiao, F., Ren, Y., Liang, Y., Yao, W., Yang, X., and Cen, L. Desen: Deep entire space cross networks for individual treatment effect estimation. In Proceedings of the ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 4612 4622-4620, 2022.

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effeet Estimation
A. Proof.
Definition A.1. Diffeomorphic counterfactual. The diffeomorphic counterfactual is defined as minimal modifications on the data manifold of the imput reversing the classification decision & (zz), represented as:
T
arg min dist(T;,) x) 8.1. k()) A k(x);
2EM
where M is the data manifold, dist(E, x) denotes the distance between T and x. The classifier k; X + {0, 1 assigns an input Z E X' to class E 10, and k(z) * k(x:) indicates a change in the classifier's prediction outcome.
Definition A.2. Let p0 (xx) == p(x = 1). (x) = p(r|t = 0) denote the factual treatment and control distributions, respectively, and (*) == p( = 1). (8) == p(ä]t = 0) denote respectively the counterfactual treatment and control distributions.
Definition A.3. For representation function & X + R. and for distribution p defined over X,, let pv be the distribution induced by & over R. Define p/p (r) = pa (rlt 1) and p (r) = Pp pr(rtt 0) to be the factual treatment and control distributions induced over R, 6 (F) == po (f = 1) and p (F) = po (Flt = 0), to be the counterfactual treatment and control distributions induced over R.
Lemma A.4.
(h, 4) (h 4) c7 (h, 1) + (1 u) c (h,
EC (h, (1 u) & (h, 4) & (h, D)
Define u = p(t= = 1) as the treatment proportion. The proof follows directly by noting that p(z, t) = a pp + (1 w) (0)-p==00 (x) (Shalit et al., 2017).
Definition A.5. Let L : y x y + R. be the absolute loss or squared loss, (x, t) denote the expected loss for the unit-treatment pair (xr, (). with representation function + : X - R and hypothesis h R (0, 13 - V which can be mathematically represented as: (2,, t) = =I Mly', ((((1), ())(() The expected factual and counterfactual losses of and & are:
er(h, = - (n, Or(e, ()dadl,
(cr(h, )) = I [xxco., (1, t)p(x, t ddddt
Definition A.6. The expected factual treated and control losses are:
/ In, 1,,11)pp= (x))xx,
1=1 (h, 2) =
-= - (h, 4)- { (2, 0)p'= (x)
Definition A.7. The expected counterfactual treated and control losses are:
- (2, "a)
d (h, Φ) =
[ (2, O)p' (x)
EC (h, 4)
Lemma A.8. (Lemma 4.1. main text). Let + X - be an invertible representation function with inverse V and h: R 10, + y hypothesis function. Define u u: pll 1) as the treatment proportion. For any treatment indicator E [0, let disc ( (x), (*))) represent the distance between counterfactual and factual samples in the representation space. Then, the counterfactual loss (CF(h, 4) is bounded by:
(cr ②) S (1 -)) oF (h, ②) + c (h, +)+
disc(T'(x)) (')) + disc g' (z)) T (x))
12

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
Proof.
[[1 - (11 (h, + uu.er - (h, 4)
(12)
(h, 4)]
[(1 (11-w) u) (1,0)) - (h, $)]
(1 - (c,7 (h, $1}+
" (1, 6) - of (h, $)] =
(13)
(1 -- / In, 11, 1)(x (x)) -1- (s))do+ (z) p (x (x= (x))d==
(1 1-01/(holo01100000))-->> -t)
2 (100 p' (i) 1 (x))dx
(1 1 w) / (n,, p² (z))dx+
(1 u) S (10
(ū) y² (z))do+
2 (10 - p² (x))
a 1n,0 (z) - (x)dr
(14)
(1 -u)
- Σ (h -=0 (a:)) +
(1 u) Σ (1000=1116000000 - p² (x=))
(n 0)(p' -
8
Σ (Z (2,, 0)p o² (=:))
(15)
(1 ") Σ ai (3Ta) (+)))
(1 4) Σ a:
" by (0 (xi) - x² +
4 Σ b. (p" (2.)) pp" x (z:)]] =
(1 4) Σ
at 106 6 (r.) 25 (f))) +
(1 4) Σ ai 1pb (f)) 25 (r:))
u > by (pt (r:) - p0 (f:)) +
13

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
"
Σ
by
(p
(f:)) - pΦ (r:))
(1 ") Σ a4 (x:) (i:))
(1
w) Σ ( (T:)) (c:)))
4, (8 (x:) (i)) +
2
2 Σ
(0'=) (i:)) g-=0 (z:)))
(16)
(1 4) Σ 44 (x;) 2 :5'=0 (T:)
(1 4) Σ a. av2 2 (2,)² Φ (z:)
b. 12 2
20.-1(00)) T T (i+)+
" b..2 2
20 (1.)) (1=0 (1.))(1=011))
disc(0'(x). (b)) + disc @6 (t), x² (x))
Equation (12) is by Lemma A.8. Equation (13) is by Definition A.6 and Definition A.7, the expected factual treated and control losses and the expected counterfactual treated and control losses. Equation (14) is based on the definition of integral formulation. The inequality (15) is by the absolute value inequality, which states that a b S la - bl. Equation (16) is by the law of Cosines and based on the assumption that (z)) is unit vector.
The counterfactual loss CCF (h, b) is bounded by combination of factual prediction errors and the distances between nnttull-cuunerrfactua pairs disc(("(x), (T)) and between counterfactuals and opposing group factuals disc(4'()) (z), b (x)). This Lemma highlights the importance of minimizing these representation distances to reduce counterfactual loss.
Definition A.9. The individual treatment effect:
(x)
E [y' y y01x]]
Definition A.10. Let f X x 1) y by an hypothesis. The estimated individual treatment effect:
(x) = f(x, f(x, 0).
Definition A.11. The expected Precision in Estimation of Heterogeneous Effect loss of f:
(PEHE()) = / (fy(z) - (())))-(()))))) -(x))) p(z)dr.
Following CFR (Shalit et al., 2017), we define mx) Ely' |x],, the expected variance of y' with respect distribution
p(x, t) is: o (p(x, t)) = Ixxx(U' (v' - m.(x))) (()))))((='()))))) ( dy dr, and of min { og of) (p(a, - 1 t))} For any function f X x {0, 13 + V, and distribution p(z, t) over X x f0, 1}, recall that fx (f t) m.)))) p(, t)drdt (F(f) o² (p(2, f)), fx tf) - (x)))² p(x. 1 - t)drdt = cCF(f) of (p(z, - t)). Next, we provide upper bound for the expected Precision in Estimation of Heterogeneous Effect.
Theorem A.12. (Theorem 4.2. main text). Let C 3) be the estimation error with funiieentationnuuctiii &b x + R with V being its inverse and hypothesis h : R x [0, + y. For any treatment indicator t E {0. 1}, let disc (b' (x), ('()) be the representation distance between counterfactual and factual samples. The error bounded by:
14

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
<PEHE(h, ②) - 2(c7 (h, 1) EF (h, 4) 20
2(disc(('(x) 26 (a)) + disc (T). 10 (x)))
Proof. Recall that we denote EPEHE (f) = (PEnE(,, for ((r,t))-(((()))) (())),
(PEHE(S)
(1f(,,1) f(x, 0)) (mm((>) - mo(x)))² p(x)dr=
Lx ((f(x,1)- - m1(x)) + (mo(x) f(x, 0)))² p(z)de
(17)
(1(6,,1 mz)))) (mo(x) f(xz, f(x, 01)²) p(z)dr=
2 x
(18)
2
(f(z, 1) mi (x))² p(x, 1 1)dx+
2
(mo(x) (I) p(z, t 0)dr+
2
(f(r, 1) mi mm (x))² p(x, / = 0)dx+
x
2
(mo(x) f(x,0)) plx,t t 1)dx
2 x
(f(a, - mi(x))² p(a,, t)drdt+
2 xx (f(x, 1) - m,(x))² p(x,, - t)dxdt S
2 (EF - of) + (EOF - o
where Equation (17) is because (x 4)²² < 2(x² + y²)), Equation (18) is because p(r) = p(x, / = 0) + p(,,, 1).
Then, we can prove Theorem A.12 based on Lemma A.4 and Lemma A.8:
<pEME(h, ④) s
2(Er(h, 4) + ccr (h, - s
2(47 (h, 1) + 4 - 20
(z) +is) 41 (#), (x))).
The estimation error (PEHε(h, Φ) is upper bounded by two terms: the standard generalization error of factual prediction e and the distance constraints in the representation space. These distance constraints represent two types of alignment: between counterfactual and factual samples & (xr)"''' (z), and between counterfactual and opposite group factual samples o'(a) (>))) By using contrastive loss to minimize these distances, sample-win wwsele-wiseeee alignment is achieved, which helps reduce the expected ITE estimation error by learning consist representations that effectively capture the characteristics of potential outcomes under different treatments.
15

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
B. Additional Technical Details
Based on reliable sample pairs (x, we leverage contrastive learning to effectively learn consistent representations for individuals across different treatments, further obtaining balanced representations through fine-grane ane-grained, sample-level alignment. Additionally, we use two separate neural networks to estimate potential outcomes under varying treatments. The detailed counterfactual contrastive learning strategy is provided in Algorithm 1.
Algorithm 1 Coonterfactual Contrastive Learning
Input: Original dataset (I, t, y), batch size N, normalizing flow 9. classifier k, step size X, threshold A,
temperature Ttemp. Encoder network Enc, Projection head Pro, prediction head h, hyper-parameters a, 8 2 0.
2: Output: Predicted factual yy, counterfactual Wef-
3: Split D into training set Drain and validation set Dealid-
4: # Stage I: Generate counterfactuals for each sample
5: for each mini-batch in Prain do
6: 2 x- 9 (x)
7: for: 1 to O(kog),
8:
V. 6
8g
9:
2 optimizerstep(X, v+)
10:
if (())))), A then
11:
return (()) 12: end if
13: end for
14: Obtain counterfactuals ( from D.
end for
# Stage Train encoder using contrastive loss with generated counterfactuals
17:
repeat
18:
Compute embeddings = =Pro (Fr(())((;))) é -Pro(Enc(En
19:
Calculate contrastive loss L. using Equation (9).
20:
Calculate predictive loss Lp using Equation (10), obtain predicted factual ûs and counterfactual d.f-
21:
Compute total loss L. = Cp + aL. + |||W|||-
22:
Update Encoder network and Projection head.
23:
Terminate if converged validation set Dvalid-
24:
until convergence
C. Metrics
For datasets with known true treatment effects, adopt two commonly evaluation metrics(Alaa & Van Der Schaar, 2017; Cheng et al., 2022; Wu et al., 2022), namely the rooted Precision in Estimation of Heterogeneous Effect (VErEhE) and the absolute error of Average Treatment Effect (CATE) defined as:
CPEEE
(r(x:) - +(xi))
(19)
ÉATE = AĪE ATE Σ (Tr 4.)
(20)
where 9; refer to the ground truth treatment effect of subject, è is the estimated treatment effect.
On Jobs dataset where the ground truth counterfactual outcomes are unavailable(Zhu et al., 2021; Reddy & Balasubramanian, 2024), we employ the policy risk Reo and the bias of Average Treatment Effect on the Treated prediction EATT as evaluation metries:
16

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effeet Estimation
Rool( = - ( = 1) Ely (x) = 1] +p( = 0) Ely (T+ (z)
(21)
17
17
77
cATT
1
y1
Σ
3
y3 v.2 v²)1
(22)
T1
7oo
71
where T X {0, 1} a policy induced from an ITE estimator f (-) with (x)
f(x) > and f(x) otherwise. This measures the average regret when treating with the induced policy Ty and thus can serve as a proxy the ITE estimation error. [Ti] and [Tol are the number of the units in the treatment and the control groups, respectively.
D. Additional Experiments.
D.1. Additional Experiments on IHDP dataset.
Existing covariate balance methods mitigate treatment selection bias across treatment and control groups by enforcing distributional balance, but fail to account for boundary samples, which are scarce and situated far from the main data distribution. These sparse boundary samples limit model's generalizability, which is particularly pronounced in scenarios with substantial individual heterogeneity. Focusing on and analyzing these samples significantly enhance model robustness and effectively reduce estimation errors.
We first provide formal definition of boundary samples. Let the feature matrix be X' E Raxd, where is the number of samples and d is d-dimensional feature vector (c.g., age, sex). Let t {0, 1 denote the treatment assignment for each unit i, where & = 1 indicates the treated sample, and t; = 0 indicates the sample belongs to the control group.
The treated group center in and the control group center uo are calculated as the mean values within each group:
240
Fi+
To filt; 0}.
To
iETo
81
Fet
T = Lill = 1}.
iET
For each sample x; E Rd, calculate the distances between x; and the treated group center, and the control group center respectively: (x,)) ||.; "1l and dz();) ||.,- uoll. Taking the treated group example, we first sort all samples in Ti by their distance to the control group center dp (x)) in descending order and select samples with large distances to exclude those in the overlapping region between treated and control groups. Among these candidates, we then select the top k samples with the largest distance to their own group center di (æ)) as treated group boundary samples. The value of ką is k, where k is the total number of boundary samples.
determined by the treated ratio n
[ such that kı
As previously mentioned, existing deep learning- causal inference methods can generally be grouped into two categories. The first category, exemplified by CFR, focuses on balancing the distributions of covariates between treated and control groups in the representation space to mitigate selection bias. The second category, represented by ABCEI, uses adversarial training to implicitly model the counterfactual distribution. To evaluate the robustness of different methods in estimating ITE, we identify 30 boundary samples corresponding to each of these two representative methods (CFR and ABCEI) and calculate the individual treatment effect estimation bias (EITE = | IT Etruel) for each boundary sample. Meanwhile, we also calculate the ITE estimation bias for these same samples using our FCCL. The results are shown in Table and Table 5.
We observe that our FCCL achieves significantly better performance than CFR and ABCEI. For example, compared to CFR, 22 out of 30 boundary samples achieve lower estimation errors, with an average error reduction of 81.83% Similarly, our FCCL also demonstrates impressive improvements over ABCEI, with 23 out of 30 samples achieving significantly smaller ITE estimation errors. These results show that FCCL enables robust sample-level alignment, achieving more accurate and stable ITE estimation.
17

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
Table The individual treatment effect estimation error of boundary samples for CFR-MMD and FCCL on IHDP dataset.
CFR-MMD
FOCL
CFR-MMD
FCCL
Sample
IT Faraee
Sample
IT Fitrue
Fpeed
FITE
PTEpred
6ITE
6ITE
1TEpped
KITE
3,5205
-0.7059
4.2264
3.7350
0.2144
16
3.8053
0.8315
2.9738
-4.0517
7.8569
2.8730
5.9707
3.0977
2.8136
0.0594
17
2.7568
-1.170
4.5739
-3.1549
5.9117
4.3918
4.3581
0.0337
4.3113
0.0805
0.9267
0.2564
0.6703
1.5344
0.6077
0.9640
-5.1606
6.1246
1.6058
0.6418
19
3.9234
0.9879
2.9355
4.0710
0.1477
5 -0.7087
-2.8101
2.1014
-0.4824
0.2263
20
2.8283
-1.0772
3.9056
3.0010
0.1726
6
3.1767
2.8624
0.3143
3.6967
0.5200
21
4.4578
1.3300
3.1278
4.5571
0.0993
2.1111
-0.6245
2.7356
-2.1950
4.3061
22
3,4389
2.1831
1.2558
3.5603
0.1215
4.5069
1.4256
3.0813
4.6920
0.1850
4.4578
-0.0649
4.5226
4.8985
0.4407
9 4.3352
1.7478
2.5874
4.5570
0,2218
24
4.3395
0.9931
3,3464
4.2455
0.0940
10
4.6119
2.6889
1.9231
4.8943
0.2824
25
4.5317
2.3423
2.1894
4.7399
0.2082
11
4.5069
-3.9979
7.9049
4.7954
9.3024
26
0.3966
-4.1701
4.5667
0.3130
0.0836
12
3.0626
-0.6165
3.6791
3.5241
0.4615
27
4.5836
-1.3938
5,9774
5,0885
0,5049
13
4.6119
1.5407
3.0712
4.9734
0.3615
1,9660
-0,9089
2.8749
-1.4851
3.4511
14 3.6437
-3.7587
7.4024
-3.7375
7.3812
29
4.5283
4.1638
0.3646
4.2638
0.2645
15
4.4007
-2.5881
6.9888
-4,4562
8.8569
30
4.1820
1.9005
2.2815
4.1756
0.0064
Table 5. The individual treatment effect estimation error of boundary samples for ABCEI and FCCL on HHDP dataset.
ABCEI
FCCL
ABCEI
FCCL
Sample
ITErree
Sample
IT Etvae
Epeed
EITEE
Frred
6STEl
IT
|ETTEl
IT UpEed
<ITE
4.6441
3,9641
0.6800
4.6739
0.0298
16
4.6354
7.2090
2.5736
4,7878
0.1524
4.6577
4.6809
0.0232
4.7763
0.1186
17
1.7826
1.2860
0.4967
3,1989
1.4163
4.6627
5.2865
0.6238
4.9184
0.2558
1.7884
1.2142
0.5742
2.3505
0,5620
-1.8670
-3.3532
1.4862
-2.5838
0.7168
3.4047
3.0164
0.3883
3.1130
0,2918
5
0.5711
0.5848
0.0137
-0.8183
1.3895
20
4.6441
3.0385
1.6056
4.6846
0.0405
6
4.5688
5.4534
0.8846
4.7526
0.1838
21
4.6354
-0.8909
5.5263
-4.8816
9.5170
2.1111
-1.7264
3.8375
-2.3312
4.4423
22
4.6687
3.5674
1.1013
4,8389
0.1701
4.6701
3,0477
1.6225
4.8280
0.1579
4.6687
1,9664
2.7024
4.8932
0.2244
4.5480
3.3778
1.1702
4.6724
0.1244
24
4.6701
4.1481
8.8182
-5,2006
9,8707
10 4,4591
-1,0039
5,4630
4.7291
0.2701
25
1.1268
-2.4555
3.1723
0.7695
0.3573
11
4.3369
3.8228
0.5141
4.5132
0.1763
26
3.9765
4.4113
0.4348
4,2919
0.3153
4.6577
5,4907
0.8329
4.7580
0.1003
27
4.6672
3.1701
1.4971
4,7203
0.0531
4.6701
2.9026
1.7676
4.9135
0.2434
4.6517
-5.6671
9.7188
-4.9836
9,6353
14 4.6627
1.9003
2.7324
4.9797
0,3170
29
4.4322
1.9372
2.4950
4.5838
0.1516
15
4.5855
2.2245
2.3610
4.8239
0.2384
30
4.6687
-3.1260
7.7947
-4.9109
9.5797
18

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effeet Estimation
We perform sensitivity analysis focusing on three key parameters: the weight of contrastive loss a, the weight of model complexity term 8, and the temperature coefficient Ttemp in the contrastive loss. In partieular, set (0.1, 0.3, 0.5, 0.8), E (1e-3, le-4, le-5, 1e-6), and Ttemp E {0.1, 0.3, 0.5, 0.7, 0.9]. The results in Figure 5 show that the contrastive loss, as one of the core optimization objectives in our proposed FCCL, impacts the individual treatment effect estimation performance. We find that FCCL generally robust to different settings of these parameters. However, fine-tuning these hyperparameters can still enhance ITE estimation performance.
1.0
1.0
0.20
PEHE
ATE
0.16
PEHE
ATE
PEHE
ATE
0.9
0.14
0.75
0.15
0.9
0.15
0.70
PEHE
0.8
0.12
PEHE
ATE:
PEHE
0.8
0.65
0.10
0.10
0.7
0.7
0.08
0.60
0.05
0.05
0.6
0.06
0.55
0..6
0.5
0.04
0.50
0.00
0.5
0.00
0.1
0.3
0.5
0.88
1fc-3
1de-4
1
1
0.1
0.3
0.5
0.7
0.9
a (B=1* 0.,=
B (1 mm050
Tuo (c=0.11 (=1*c-4)
(c)
Figure 5. ITE estimation performance of our FCCL under different parameters IHDD dataset.
D.2. Additional Experiments on Jobs dataset.
On the Jobs dataset, although the performance of representation learning is less satisfactory compared to adversarial training to some extent, FCCL achieves comparable results the baselines, particularly in the error metric CATT- The main reason we identified is the limited data size on the Jobs dataset, which is further evidenced by the observation that traditional methods often outperform deep learning on this dataset.
We also provide sensitivity analysis on Jobs dataset focusing on the weight of contrastive learning o, the weight of model complexity term 3, and the temperature coefficient Ttemp. In particular, we vary the weight the contrastive loss {0.2, 0.5, 0.8, 1.0, 1.2} while fixing 1c-4 and Ttemp 0.2. Similarly, we explore the effect the regularization weight (1e-3, 1e-4, 1e-5, lc-6) with a 1.0 and Ttemp 0.2. For the temperature coefficient, we evaluate Ttemp {0.1, 0.3, 0.5, 0.7, 0.9} while keeping
1.0 and 3 1e-5. From the results in Figure 6, find our model is generally robust to different a settings and the best performance is achieved with 1.0.
Table 6. sampin-sample and out-of-sample policy risk and error on the average treatment effect on the treated (ATT) for the various models on Jobs dataset.
Method
Roat of
out oott-of ATT
OLS-1
0.22(0.02)
0.01(0.00)
0.23(0.02)
0.08(0.04)
OLS- 2
0.21(0.01)
0.01(0.01)
0.24(0.03)
0.08(0.03)
BART
0.23(0.00)
0.02(0.02)
0.25(0.02)
0.08(0.04)
KNN
0.23(0.01)
0.02(0.00)
0.26(0.02)
0.13(0.05)
DML
0.23(0.01)
0.02(0.02)
0.24(0.02)
0.10(0.03)
CFR- Wass
0.23(0.01)
0.06(0.02)
0.26(0.02)
0.10(0.04)
CFR-MMD
0.22(0.00)
0.07(0.03)
0.27(0.01)
0.12(0.05)
SITE
0.23(0.01)
0.05(0.02)
0.25(0.02)
0.10(0.04)
CITE
0.23(0.00)
0.10(0.03)
0.26(0.02)
0.13(0.05)
GANITE
0.14(0.02)
0.27(0.74)
0.15(
0.31
ABCEI
0.17(0.02)
0.05(0.02)
0.22(0.02)
0.15(0.08)
CBRE
0.30(0.00)
0.08(0.03)
0.31(0.00)
0.11(0.03)
DIGNet
0.23(0.01)
0.07(0.04)
0.26(0.01)
0.13(0.04)
FCCL
0.23(0.01)
0.05(0.01)
0.25(0.02)
0.07(0.03)
19

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
0.29
0.10
0.280
0.150
0.34
Pollic risk
Bias att
0.16
Policy risk
Bias att
Policy risk
Bins att
0.14
0.28
-0.08
0.275
0.125
Pollcy risk
0.32
0.12
Blas aft
Policy risk
0.066
Policy risk
0.270
0.100
0.30
0.10
0.27
18ia
0.265
0.075
Bias
-0.04
0.28
0.260
0.050
0.26
0.06
0..02
0.255
0.025
0.26
0.04
0.25
0.00
0.250
0.000
0.2
0.5
1.2
19e-3
1
19e-
1-6-
0,5
0.9
(1=1"M-44 0.2)
(1.1.00 0.2)
Trey (a=1.0, B=1*e-5)
(4)
(b)
(c)
Figure 6. ITE estimation performance of our FCCL under different parameters on Jobs datasct.
D.3. Additional Experiments on Synthetic dataset.
We generate covariates from the multivariate normal distribution N 70, [.11pTT + (1-11111)) (1-))[p]),
where the covariance
matrix combines an all-ones matrix 1 and an identity matrix Ip. The scaling parameter {0.4, 1.0, 1.2) controls the degree of covariate dispersion, as shown in Figure 7. We sample 800 units with parameters p = 10, p p= 0.2, o² = 3, 3o = [0.2, 0.2). and 31 [1.2, .... 1.2]. For each 7, we generate 30 independent datasets, dividing them into training, validation, and test sets with ratios 63%, 27%, and 10%, respectively. The data generation process is outlined as follows:
X, N (0,7 o². 1pl,1 + (1 [p)))[p)) p)), T: X, Bernoulli (1/(1+ exp x-,))) y = X, + E.. Y! 3,X, Ei. & (0001)
Figure presents T-SNE visualizations of the covariates under different values of y. At = 0.4, the treated and control groups exhibit relatively concentrated distributions. y increases 1.0 and 1.2, the covariate distributions become significantly more dispersed, leading greater number of boundary and extreme samples. This trend highlights how larger ? values represent greater heterogeneity, providing various scenarios for the individual treatment effect estimation.
y-0.4
y-0.7
Treatment Control
Treatment Control
4-
2-
2
0-
-2
-2
-4
6
-4
-2
0
4
-4
-2
4
y=1.0
y=1.2
Treatment Control
Treatment
4-
4
Control
2-
0-
-2-
6
-4
-2
0
4
-4
4
Figure 7. T-SNE visualizations of the covariates as varies.
We report additional experimental results as supplement to the main text. Table presents the results, where FCCL
20

---

Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
Table 7. Additional experimental results Synthetic datasets.
0.4
7 0.7
7 1.0
? 1.2
Method
out-of
ond ord-oo of
out-of
ond out-o1
EPENE
EEEE
PEEE
PEHE
EEEEE
SEEHE
CEEHE
РEEE
OLS-1
8.39(0.38)
8.41(0.84)
10.86(0.43)
30.85(1.34)
(2.89(0.59)
(3.00(1.68)
14.21(0.59)
14.32(1.35)
OLS-2
5.92(0.27)
5.94(0.60)
7.64(0.30)
7.64(0.96)
9.05(0.40)
9.13(1.18)
9.97(0.41)
10.07(1.11)
BART
4.04(0.22)
3.40(0.50)
4.70(0.20)
4.17(0.60)
5.36(0.31)
4.86(0.61)
5.81(0.26)
6.14(0.58)
KNN
4.55(0.33)
5.50(0.62)
6.48(0.37)
7.26(0.75)
7.91(0.38)
9.08(1.27)
8.76(0.46)
10.27(1.07)
DML
4.10(0.26)
3.96(0.54)
4,96(0.26)
5.3000.62)
5.72(0.52)
6.08(0.80)
7.76(0.62)
8.46(0.77)
CFR-Wass
2.48(0.05)
2.48(0.06)
3.73(0.05)
3.60(0.09)
4.68(0.07)
4.72(0.14) 14)
5.34(0.07)
5.37(0.14)
CFR-MMD
2.54(0.05)
2.54(0.06)
3.75(0.05)
3.62(0.09)
4.70(0.07)
4.74(0.14)
5.37(0.08)
5.41(0.14)
SITE
2.68(0.11)
2.69(0.13)
4.17(0.1))
4.25(0.23)
5.98(..34)
6.01(0.38)
6.19(0.15)
6.21(0.17)
CITE
2.69(0.04)
2.71(0.07)
3.81(0.06)
3.70(0.10)
4.68(0.06)
4.74(0.14)
5.39(0.07)
5.41(0.14)
GANITE
4.66(0.03)
4.69(0.06)
6,20(0.03)
6.16(0.07)
7.32(0.03)
7.33(0.07)
8.08(0.03)
8.11(0.08)
ABCEI
2.73(0.03)
2.75(0.06)
3.74(0.04)
3,57(0.09)
4.61(0.05)
4.73(0.13)
5.19(0.06)
5.19(0.12)
CBRE
2.91(0.03)
2.93(0.05)
4.01(0.04)
3.85(0.08)
4.95(0.05)
5.02(0.12)
5.77(0.06)
5.73(0.13)
DIGNet
3.17(0.09)
3.18(0.11)
4.09(0.10)
3.97(0.13)
5.05(0.10)
5.10(0.17)
5.78(0.09)
5.81(0.15)
FCCL
2.56(0.04)
2.58(0.06)
3.65(0.05)
3.50(0.09)
4.40(0.06)
4.49(0.13)
5.10(0.06)
5.12(0.12)
achieves smaller /EEEHE than baselines, especially when covariate distributions become increasingly dispersed (e.g..
? 1.2). This demonstrates our method's superior capability in capturing individual-level heterogeneity, especially under high covariate dispersion with sparse boundary samples.
21


Reviews:
[{"review_id": "DrFCqVp89j", "reviewer": "Reviewer_PLY3", "summary": "The paper proposes FCCL  framework for the ITE estimation.  The proposed method can generate realistic counterfactuals by leveraging normalizing flows to ensure adherence to the data manifold, preserving semantic similarity to factual samples. The authors also derive a new generalization bound connecting ITE estimation error to factual prediction errors and representation distances between factual-counterfactual pairs, providing theoretical grounding for their proposed sample-level alignment method.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your time and your thoughtful and encouraging comments. We hope our responses can resolve your concern.\n\n$\\textbf{Q1:}$ I suggest authors give some theoretical or empirical evidence to show that the bound is tight (at least tighter than the bound proposed in CFR).\n\n$\\textbf{R1:}$ Following your helpful suggestion, we conduct an experimental analysis to compare the generalization error bounds of our proposed FCCL (Equation (14) in Appendix A) with the CFR bound (Equation (8) in Appendix A.2, [1]) under the same conditions.\n\nAs shown in Table 1, we track the generalization error bounds at different iterations. The results demonstrate that as the number of iterations increases, FCCL achieves a significantly tighter bound than that of CFR.\n\nTable1 Generalization-error bound comparison on IHDP.\n| Iterations | 400      | 800      | 1200     | 1600     | 2000     |\n|------------|----------|----------|----------|----------|----------|\n| CFR | 13.24 | 10.84 | 10.06 | 9.64 | 9.91 |\n| FCCL | 8.54 | 3.74 | 2.93 | 2.77 | 2.67 |\n\nIn addition, we would like to clarify that our FCCL provides a different theoretical perspective from CFR. Our generalization error bound links ITE estimation error to the generalization error of factual predictions and representation distances, motivating our focus on minimizing these distances via sample-level alignment. Instead, the bound of CFR shows ITE estimation error can be reduced by the difference between the treated and control distributions.\n\n[1] Shalit, U., et al. Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pp. 3076–3085. PMLR, 2017.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes.", "experimental_designs_or_analyses": "The experimental design and comparisons are proper.", "relation_to_broader_scientific_literature": "ITE model development.", "essential_references_not_discussed": "Some papers also discuss the new upper bound of PEHE. Can include them in the literature review.", "other_strengths_and_weaknesses": "Strengths:\n\nThe motivation, presentation, and experimental comparisons are good.\n\nWeaknesses:\n\nI only have one concern about the theoretical parts. One of the key contributions is that the authors claim they propose a new ITE error bound for their sample alignment method. However, anyone can propose an ITE error bound and minimize such a bound to learn the ITE model. The point is, how can you make sure the proposed ITE bound is tight or not? If the bound is too loose, such a theoretical bound may give bad guidance for training the ITE model.\n\nI thus suggest authors give some theoretical or empirical evidence to show that the bound is tight (at least tighter than the bound proposed in CFRWASS)", "other_comments_or_suggestions": "None.", "questions_for_authors": "Please see Strengths And Weaknesses", "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": "M8uXDPVj7H", "reviewer": "Reviewer_HBA5", "summary": "The paper introduces Flow-based Counterfactual Contrastive Learning (FCCL), a novel approach for Individual Treatment Effect (ITE) estimation that integrates normalizing flows for realistic counterfactual generation and contrastive learning for fine-grained sample alignment. It derives a theoretical generalization-error bound linking ITE estimation error to factual prediction error and representation distances. Empirical evaluations on synthetic, semi-synthetic (IHDP), and real-world (Jobs) datasets demonstrate FCCL’s superior performance", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful suggestions. We have addressed the comments related to the counterfactual generation and model robustness evaluation. Please see our responses below.\n\n$\\textbf{Q1:}$ How does FCCL compare to diffusion-based counterfactual generators?\n\n$\\textbf{R1:}$ We thank the reviewer for the question. We add experiments on diffusion-based counterfactual generation (Table1), which shows performance comparable to FCCL, but with slightly inferior result on $\\epsilon _{PEHE}^{out-of}$. This may be due to the noise-driven mechanism of the diffusion model [1], which causes counterfactual to deviate from the sample semantic space, especially in the out-sample cases. In contrast, the flow-based model ensures that counterfactuals reside on the same manifold as the original instances, which generates both meaningful and reliable counterfactuals, making FCCL more reliable for individual-level treatment effect predictions. We will include this discussion in the revision.\n\nTable1 ITE estimation errors (std) with different generation methods on IHDP.\n| Method    | $\\epsilon _{PEHE}^{within}$ | $\\epsilon _{ATE}^{within}$ | $\\epsilon _{PEHE}^{out-of}$ | $\\epsilon _{ATE}^{out-of}$ |\n|-----------|-----------------------------|----------------------------|-----------------------------|----------------------------|\n| diffusion-based | 0.54(0.05)   | 0.09(0.01)   | 0.72(0.10)  | 0.12(0.03) |\n| FCCL     | 0.53(0.04)  | 0.09(0.01)   | 0.64(0.07)  | 0.12(0.02) |\n\n$\\textbf{Q2:}$ How does the contrastive loss temperature $\\tau $ affect performance?\n\n$\\textbf{R2:}$ We perform sensitivity analysis focusing on the temperature coefficient $\\tau$ (see Table 2). We observe minimal variation in performance with respect to $\\tau $, which demonstrates our model's general robustness.\n\nTable2 ITE estimation errors (std) with different values of $\\tau $ on IHDP.\n| $\\tau$ | $\\epsilon_{PEHE}^{within}$ | $\\epsilon_{ATE}^{within}$ | $\\epsilon_{PEHE}^{out-of}$ | $\\epsilon_{ATE}^{out-of}$ |\n|--------|----------------------------|----------------------------|-----------------------------|----------------------------|\n| 0.1  | 0.54(0.04)   | 0.11(0.02)  | 0.64(0.06)  | 0.13(0.02) |\n| 0.3  | 0.56(0.04)   | 0.10(0.02)  | 0.70(0.08)  | 0.14(0.02) |\n| 0.5  | 0.58(0.06)   | 0.10(0.02)  | 0.76(0.16)  | 0.11(0.02) |\n| 0.7  | 0.61(0.06)   | 0.10(0.02)  | 0.79(0.15)  | 0.11(0.02) |\n| 0.9  | 0.56(0.04)   | 0.10(0.01)  | 0.69(0.07)  | 0.12(0.02) |\n\n$\\textbf{Q3:}$ Can FCCL generalize to continuous treatments? \n\n$\\textbf{R3:}$ This is a valuable point. FCCL currently focuses on binary treatments, we can extend it to continuous treatments in future work. The most direct approach would involve discretizing the continuous treatment variable[2]. Specifically, we can split treatments into $E$ heads, each assigned a dose level from the range$[a, b]$, which is divided into $E$ equal intervals of width $(b-a)/E$. We will explore better approaches to handle continuous treatments directly.\n\n$\\textbf{Q4 (Claim2):}$ Additional ablation studies isolating the contrastive loss would be beneficial. \n\n$\\textbf{R4:}$ We discussed the impact of contrastive loss, thank you for going into Figure 4 in the main text. \n\n$\\textbf{Q5 (Claim3):}$ Empirical validation of whether this ITE bound holds in practice is not explicitly tested.\n\n$\\textbf{R5:}$ Our empirical analysis demonstrates that the proposed ITE bound effectively guides ITE model training and becomes tighter than the bound proposed in CFR as iterations increase under identical conditions[3] (Table 3). \n\nTable3 Generalization-error bound comparison on IHDP.\n| Iterations | 400 | 800 | 1200 | 1600 | 2000 |\n|------------|---------|---------|---------|---------|---------|\n| CFR  | 13.24 | 10.84 | 10.06 | 9.64 | 9.91 |\n| FCCL | 8.54 | 3.74 | 2.93 | 2.77 | 2.67 |\n\n$\\textbf{Q6 (Methods And Evaluation Criteria):}$ Further quantitative measures of alignment would reinforce the claims (e.g., KL divergence).\n\n$\\textbf{R6:}$ Thank you for your suggestion. We add KL divergence comparisons between treatment and control distributions for two typical methods (see Table 4), which demonstrates that our FCCL can better address distribution shifts. We will include this metric in Figure 3 in the next version.\n\nTable4 KL divergence mean (std) comparison on IHDP.\n| Method | CFR  | ABCEI | FCCL |\n|--------|-----------|-----------|-----------|\n| KL | 0.14(0.08) | 0.20(0.30) | 0.09(0.05) |\n\n[1] Kotelnikov Akim, et al. Tabddpm: Modelling tabular data with diffusion models. In International Conference on Machine Learning, pp. 17564–17579. PMLR, 2023.\n\n[2] Schwab et al. Learning counterfactual representations for estimating individual dose-response curves. Proceedings of the AAAI Conference on Artificial Intelligence, pp. 5612-5619, 2020.\n\n[3] Shalit, U., et al. Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pp. 3076–3085. PMLR, 2017.", "claims_and_evidence": "Claim 1: FCCL generates realistic counterfactuals that adhere to the data manifold.\n\nSupported. Through the use of normalizing flows, which enforce structure on counterfactual transformations. The authors present visualization results showing that FCCL maintains sample-level semantic consistency better than baseline methods.\n\nClaim 2: FCCL reduces ITE estimation error via sample-level alignment.\n\nPartially supported. While contrastive learning improves alignment, the empirical evidence (particularly in Tables 1 and 2) primarily shows improvement in error metrics rather than a direct measure of sample alignment. Additional ablation studies isolating the contrastive loss would be beneficial.\n\nClaim 3: The proposed theoretical generalization-error bound justifies FCCL’s effectiveness.\n\nSupported. The bound is mathematically derived (Theorem 4.2) and aligns with the contrastive loss objective. However, empirical validation of whether this bound holds in practice is not explicitly tested.\n\nClaim 4: FCCL significantly outperforms state-of-the-art baselines.\n\nMostly supported. The results show FCCL achieving the best ϵ_PEHE and ϵ_ATE scores across datasets. However, the magnitude of improvement varies, and for some cases (e.g., IHDP out-of-sample ATE error), the advantage is marginal.", "methods_and_evaluation_criteria": "The use of benchmark datasets (IHDP, Jobs, synthetic data) is appropriate for treatment effect estimation.\nBaselines (OLS, CFR, GANITE, ABCEI, etc.) are well-chosen, covering both traditional and deep learning-based ITE estimation methods.\nEvaluation metrics (ϵ_PEHE, ϵ_ATE, ATT) are standard, but additional fairness metrics (e.g., subgroup fairness or bias analysis) could strengthen the evaluation.\n\nThe use of latent space visualizations (Figure 3) is insightful, but further quantitative measures of alignment (e.g., KL divergence or propensity score matching quality) would reinforce the claims.", "theoretical_claims": "Correctness of Proofs: The theoretical derivations appear correct and align with existing literature on treatment effect bounds (e.g., Shalit et al., 2017).\n\nMissing Considerations: The assumptions regarding the invertibility of representations (Φ) and the geodesic distance formulation could be further justified. Additionally, the bound does not account for sample sparsity effects, which are critical in real-world datasets.", "experimental_designs_or_analyses": "The experiment design is generally well-structured, with appropriate train-test splits and multiple trials.\n\nMissing Analyses: The impact of hyperparameters (especially temperature τ in contrastive loss) is not explored. Additionally, no robustness checks (e.g., sensitivity to noise or dataset shifts) are provided.", "relation_to_broader_scientific_literature": "FCCL extends prior work on representation learning for ITE estimation (e.g., CFR, SITE, CITE) to contrastive learning for causal inference, aligning with recent trends in self-supervised learning for structured data.\nThe approach could be adapted to multi-treatment and continuous intervention settings, similar to recent developments in continuous treatment estimation (e.g., Kazemi & Ester, 2024).", "essential_references_not_discussed": "The paper extensively cites prior work but does not discuss alternatives to normalizing flows for counterfactual generation, such as energy-based models (Du et al., 2021) or diffusion-based generative approaches (Song et al., 2021).", "other_strengths_and_weaknesses": "no", "other_comments_or_suggestions": "no", "questions_for_authors": "1) How does FCCL compare to diffusion-based counterfactual generators? Diffusion models are gaining traction in structured data. Would FCCL’s theoretical framework extend to them?\n2) How stable is FCCL under different hyperparameter choices? Particularly, how does the contrastive loss temperature τ affect performance?\n3) Can FCCL generalize to continuous treatments? The paper focuses on binary treatments, but many real-world applications require handling continuous interventions.", "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": "hAt4DVw0ZD", "reviewer": "Reviewer_N1DA", "summary": "This paper presents FCCL, an ITE estimation method. FCCL integrates diffeomorphic counterfactual generation and contrastive learning to align treatment and control groups at a fine-grained, sample level, mitigating distribution shifts and approximating RCT randomization. By ensuring realistic counterfactuals and enforcing semantic consistency, FCCL lowers ITE estimation error. Experiments indicate superior performance in heterogeneous and data-scarce settings.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and for recognizing the importance of heterogeneous treatment effect estimation. Please see below answers to the questions.\n\n$\\textbf{Q1:}$ Provide a proper definition of boundary samples and discuss why handling them is challenging.\n\n$\\textbf{R1:}$ We thank the reviewer for consideration. We provide a formal definition of boundary samples in Appendix D.1, and will include this definition in the main text in the revision. \n\nFurthermore, we analyze the challenges of handling boundary samples from two perspectives:\n\n(1) Many scenarios demand flexible investigation of effect heterogeneity across individuals. However, boundary samples, due to their distinctive characteristics, reside in low-probability density regions distant from class centroids, which poses challenges for accurate ITE estimation.\n\n(2) Existing methods, such as MMD, minimize the distribution discrepancy between treated and control groups by aligning their mean representations $\\frac{1}{m\\_1}\\textstyle\\sum_{i=1}^{m\\_1}{\\phi}(x\\_{i}^{t})$ and $\\frac{1}{m\\_2}\\textstyle\\sum_{j=1}^{m\\_2}{\\phi}(x_{j}^{c})$. However, such methods overlook samples near the boundaries of the treatment and control distributions.\n\nOur method achieves robust performance via sample-level alignment, especially in scenarios where individual differences are significant. Besides, we also provide boundary sample analysis by listing the estimation errors of the first five samples from Table 4 in Appendix D, as shown in Table 1. More detailed results are presented in Appendix D.\n\nTable1 The ITE estimation error $|\\epsilon_{ITE}|$ of boundary samples on IHDP.\n| Sample | CFR-MMD | FCCL   |\n|--------|---------|--------|\n| 1      | 4.2264  | 0.2144 |\n| 2      | 3.0977  | 0.0594 |\n| 3      | 0.0337  | 0.0805 |\n| 4      | 6.1246  | 0.6418 |\n| 5      | 2.1014  | 0.2263 |\n\n$\\textbf{Q2:}$ Is optimization end-to-end? or are diffeomorphic counterfactuals found first and then the representation and prediction modules are trained?\n\n$\\textbf{R2:}$ Our approach falls into the latter: we first generate the diffeomorphic counterfactuals, and then train the representation and prediction modules.\n\n$\\textbf{Q3:}$ Why the smaller evaluation metric “dis” is desirable?\n\n$\\textbf{R3:}$ A smaller \"dis\" value indicates that dispersion of samples in the latent representation space is reduced, which further reflects a smaller distance between boundary samples and their corresponding class centers. This shows that there are fewer boundary samples, enabling better model fitting and exhibiting lower ITE estimation bias, as demonstrated by our results. We will make this point clearer in the next version.\n\n$\\textbf{Q4 (Weaknesses2):}$ The evaluation could benefit from inclusion of more challenging datasets (e.g., ACIC).\n\n$\\textbf{R4:}$ Thank you for your valuable suggestions. For validation, we additionally evaluate our method against several representative baselines on ACIC (Table 2). We observe that our method still outperforms the baselines.\n\nTable2 Within-sample and out-of-sample mean (std) for the metrics on ACIC.\n| Method  | $\\epsilon_{PEHE}^{within}$ | $\\epsilon_{ATE}^{within}$ | $\\epsilon_{PEHE}^{out-of}$ | $\\epsilon_{ATE}^{out-of}$ |\n|---------|----------------------------|----------------------------|-----------------------------|----------------------------|\n| CFR-MMD | 1.70(0.38)                 | 0.29(0.12)                 | 2.36(0.59)                  | 0.30(0.12)                 |\n| SITE    | 1.71(0.39)                 | 0.38(0.14)                 | 2.33(0.58)                  | 0.39(0.15)                 |\n| ABCEI   | 1.93(0.46)                 | 0.17(0.07)                 | 2.49(0.64)                  | 0.18(0.07)                 |\n| CBRE    | 1.69(0.38)                 | 0.12(0.05)                 | 2.31(0.56)                  | 0.14(0.06)                 |\n| DIGNet  | 1.66(0.35)                 | 0.23(0.07)                 | 2.30(0.54)                  | 0.24(0.07)                 |\n| FCCL    | 1.51(0.40)                 | 0.16(0.05)                 | 2.16(0.56)                  | 0.17(0.05)                 |", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "I looked at Lemma 4.1 and Theorem 4.2. They seem to be correct.", "experimental_designs_or_analyses": "- The experiments are valid; although it would have been nice to include more complicated datasets such as ACIC. This is important to evaluate robustness and generalization of the proposed method in large scale settings.\n- The evaluation metric “dis” in Fig. 3 is not well-motivated (see section \"Questions For Authors” for a specific question).", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths:\n- The paper provides strong theoretical grounding.\n- Innovative combination of normalizing flows (to maintain semantic meaning in counterfactual samples) and contrastive learning (to ensure robust alignment), addressing core challenges in causal inference.\n\nWeaknesses:\n- The paper could improve in terms of clarity and flow. Specifically, theory is presented before providing a clear intuition or practical context, which makes it hard to follow the core logic. As a result, readers might struggle with intuitive understanding without first seeing practical motivation.\n- The evaluation could benefit from inclusion of more challenging datasets (e.g., ACIC).\n- The metric \"dis\" used in evaluations lacks clear motivation.", "other_comments_or_suggestions": "NA", "questions_for_authors": "1. It is stated that the proposed method is specifically useful for “handling sparse boundary samples.” However, neither a proper definition is provided, nor a discussion of why handling them is challenging.\n2. Is optimization end-to-end? or are diffeomorphic counterfactuals found first and then the representation and prediction modules are trained?\n3. The evaluation metric “dis” (average distance between boundary samples and corresponding class centers) in Fig. 3 is stated to be reflecting sample heterogeneity. However, it’s not mentioned why a smaller dis is desirable.", "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": "kH4VFwDjJg", "reviewer": "Reviewer_Zyty", "summary": "The paper points out that the prediction of the individual treatment effect (ITE) is crucial for personalized therapy planning and proposes a contrastive learning approach (along the lines of SIMCLR) to estimate it. The accuracy of an estimator is measured in terms of the expected squared error of the ITE estimates with respect to the ground-truth. However, ground-truth ITEs are not accessible in practice (since the counter-factual outcome is unknown), and existing sample alignment methods are not good enough for applications with high-dimensional covariates and considerable individual heterogeneity. To circumvent this difficulty, the authors derive a tractable contrastive upper bound for the squared loss that justifies their choice of learning method. Experimental comparisons with several baselines demonstrate the superior behavior of the new method.\n\nCrucially, the authors explain that contrastive training examples cannot be generated by standard data augmentations or adversarial training in data space, because these samples are not sufficiently realistic -- they are too far from the manifold of typical covariate samples. Instead, they propose to first train a normalizing flow to represent the data distribution around the manifold, and then perform a gradient search for contrastive samples in the latent space, where the gradient is calculated from the Jacobian of the decoder. The search for contrastive examples thus stays near the manifold, and the generated counter-factual training data are much more realistic.\n\nEDIT after rebuttal: The authors appropriately addressed my questions and concerns, so I've increased my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate your high recognition of our work, and are eager to share our thoughts with you.\n\n$\\textbf{Q1:}$ Relation to debiased ML.\n\n$\\textbf{R1:}$ Thanks for your insightful idea. We would like to clarify that our work primarily focuses on the challenges of deep learning-based causal effect estimation. As debiased machine learning (DML) is a traditional method, we did not explicitly discuss it in the literature review. We will include this discussion in the revision. \n\nThen, we discuss the differences and connections between DML and our FCCL. \n\nFirst, we explain the different mechanisms that DML and our proposed FCCL use to address confounding bias. DML handles confounding bias through orthogonal residual regression, i.e., $Y_{i}-M_{y}(X_{i}) =\\tau(X_{i})(T_{i}-M_{t}(X_{i}))+\\epsilon_{i} $, where $M_{y}(X_{i})$ eliminates the influence of the confounder $X$ on $Y$ and $M_{t}(X_{i})$ eliminates the influence of the confounder $X$ on $T$[1]. The way creates statistical independence between treatment and confounder to simulate randomization conditions. Instead, our method captures the characteristics of potential outcomes under different treatments, mitigating distribution shifts through sample-level alignment to emulate RCTs, which better captures individual-level heterogeneity. \n\nSecond, the DML way to addressing bias is highly instructive, particularly its unbiasedness. Therefore, we see great potential in incorporating DML into our future work, as R-learners draw inspiration from DML's approach to handling confounding bias, $\\tau(\\cdot)=argmin_{\\tau}\\frac{1}{n} \\sum_{i=1}^{n} ((Y_{i} -M_{y}(X_{i}))-\\tau(X_{i})(T_{i} -M_{t}(X_{i})) )^{2}  $. Thanks again. \n\nBesides, we add comparison experiments with DML (Table1) and the results show that FCCL achieves lower estimation errors.\n\nTable1 ITE estimation errors (std) comparison with DML on IHDP.\n| Method | $\\epsilon_{PEHE}^{within}$ | $\\epsilon_{ATE}^{within}$ | $\\epsilon_{PEHE}^{out-of}$ | $\\epsilon_{ATE}^{out-of}$ |\n|--------|----------------------------|----------------------------|-----------------------------|----------------------------|\n| DML    | 2.87(0.09)                 | 0.30(0.05)                 | 2.95(0.14)                  | 0.35(0.03)                 |\n| FCCL   | 0.53(0.04)                 | 0.09(0.01)                 | 0.64(0.07)                  | 0.12(0.02)                 |\n\n$\\textbf{Q2:}$ How to fix presentation shortcomings?\n\n$\\textbf{R2:}$ Thank you for pointing out the typos, we will correct them in next version.\n\n( 1 ) We will add a formal definition of $k(x)$ in $\\textbf{Definition 3.6}$.\n\n( 2 ) We thank the reviewer for the suggestion. Indeed, $dis(\\cdot ,\\cdot )$ denotes the inner product of vectors in the representation space,i.e.,\n    $disc(\\Phi^{t}(x),\\Phi^{t}(\\tilde{x}))=u \\sum_{i} b_{i} \\sqrt{2-2{{\\Phi }^{t=1}({x}\\_{i})}^{T}{\\Phi }^{t=1}({\\tilde{x} }\\_{i})}+(1-u)\\sum_{i} {a}_{i}\\sqrt{2-2{{\\Phi }^{t=0}({x}\\_{i})}^{T}{\\Phi }^{t=0}({\\tilde{x} }\\_{i})}$. Detailed proofs are provided in Appendix A. $\\textbf{Theorem 4.2} $ shows that the estimation error $\\epsilon _{PEHE}(h,\\Phi )$ is upper bounded by the distance constraints in the representation space. Therefore, we implement the optimization via contrastive loss, specifically by leveraging the cosine similarity $sim(\\cdot ,\\cdot )$.    \n\n( 3 ) We apologize and will thoroughly go over this for the revision, specifically changing $[\\rho \\mathbf{1}{p} \\mathbf{1}{p}^{\\prime}+(1-\\rho) \\mathbf{I}{p}]$ to $[\\rho \\mathbf{1}\\_{p} \\mathbf{1}\\_{p}^{\\prime}+(1-\\rho) \\mathbf{I}\\_{p}]$, where $\\mathbf{1}\\_{p}$ denotes the $p$-dimensional all-ones vector and $\\mathbf{I}\\_{p}$ denotes the identity matrix of size $p$.\n\n$\\textbf{Q3:}$ Experiments demonstrating the superiority of flow-based counter-factual generation. For example, one could use the MMD to quantify the difference between the distributions of real and synthetic (counter-factual) covariantes.\n\n$\\textbf{R3:}$ Following your valuable suggestion, we use the MMD to quantify the difference between the distributions of factual and counterfactual covariates (Table 2) and the result shows that the MMD is considerably smaller for our FCCL. Besides, we compare the effect of alternative counterfactual generation strategies on ITE estimation error in Table 3 in the main text. These results show the superiority of flow-based counterfactual generation.\n\nTable2 MMD mean (std) comparison on IHDP.\n| Method | grad asc in $\\mathcal{X}$ | GAN   | FCCL  |\n|--------|---------------------------|-------|-------|\n| MMD    | 0.13(0.16)                | 0.52(0.55) | 0.09(0.14) |\n\n[1] https://matheusfacure.github.io/python-causality-handbook/22-Debiased-Orthogonal-Machine-Learning.html#non-parametric-double-debiased-ml", "claims_and_evidence": "The claims are largely clear and well supported. \n\nHowever, I am surprised that debiased (\"double\") machine learning (e.g. [1]) is neither mentioned as related work nor included in the comparison. In my understanding, debiased ML is a relatively simple method that can estimate ITEs in the absence of randomized controlled experiments, where selection bias of treatment decisions would otherwise contaminate the predictions. I'm curious how the authors asses this possibility.\n\n[1] https://matheusfacure.github.io/python-causality-handbook/22-Debiased-Orthogonal-Machine-Learning.html#non-parametric-double-debiased-ml", "methods_and_evaluation_criteria": "Evaluation follows standard practices in the field and standard benchmarks. Results appear trustworthy.", "theoretical_claims": "The claims in Lemma 4.1 and Theorem 4.2 are plausible, but I did not check the proofs.", "experimental_designs_or_analyses": "Evaluation follows standard practices in the field and standard benchmarks. Results appear trustworthy.\n\nIt would have been helpful to include an experiment demonstrating that standard data augmentation and adversarial training fail and how the flow-based method fixes this. For example, one could use the Maximum Mean Discrepancy (MMD) to quantify the difference between the distributions of real and synthetic (counter-factual) covariantes. Then, the MMD should be considerably smaller for the proposed method. \n\nThe latter experiment could be augmented by an ablation study, showing how ITE results deteriorate when flow-based counter-factual generation is not used.", "relation_to_broader_scientific_literature": "The literature is reviewed well and compared under standard protocols with the proposed method, except for the possible omission of debiased ML mentioned above.", "essential_references_not_discussed": "see above.", "other_strengths_and_weaknesses": "The main weakness of the paper are some minor presentation glitches.\n\n* Definition 3.6 (and repeatedly further down the paper, e.g. equations (6) - (8)): The text refers to \"the classification decision k(x)\", but k(x) is undefined.\n\n* Lemma 4.1, Theorem 4.2: disc(., .) (the distance in representation space) is never concretely defined, nor are the pros and cons of possible choices discussed. How is it related to the similarity sim(.,.) in equation (9)?\n\n* Second line of section 5.1: possible typo [rho 1 p 1 p'] => [rho 1 p 1^T p' ] ?", "other_comments_or_suggestions": "none", "questions_for_authors": "Please answer the following questions (I'm willing to increase my score):\n\n* relation to debiased ML\n\n* how to fix presentation shortcomings\n\n* experiments demonstrating the superiority of flow-based counter-factual generation", "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:

Meta-review: The reviewers generally appreciated the work in combining normalizing flows and contrastive learning to better estimate individual treatment effects (ITE). They generally felt that both the theory and experiments were useful and were explained well.

The authors should aim to include several points of clarification and explanation brought up during the review process in the final version.
Decision: Accept (poster)