Research map / Probabilistic and scientific modeling / Treatment effects
Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes FCCL, which generates diffeomorphic counterfactuals via gradient ascent in a normalizing flow's latent space and aligns factual-counterfactual pairs with a contrastive loss for individual treatment effect estimation.
Problem
Estimating individual treatment effects from observational data under covariate shift and absent counterfactuals, where global distribution balancing lacks sample-level alignment, especially for boundary and heterogeneous samples.
Approach
Generate counterfactuals on the data manifold using a normalizing flow's latent space (instead of noisy input-space perturbations) and apply a contrastive loss between factual and counterfactual representations, motivated by a derived ITE generalization-error bound linking estimation error to factual error and factual-counterfactual representation distances.
Key results
- On IHDP out-of-sample, FCCL achieves sqrt PEHE 0.64(0.07) and ATE error 0.12(0.02), the best among 13 baselines, reducing out-of-sample ATE error by 57.1% versus CFR.
- On synthetic datasets, FCCL's out-of-sample sqrt PEHE rises from 2.58 at gamma=0.4 to 5.12 at gamma=1.2, a 12.1% smaller increase than CFR-Wass.
- Flow-based counterfactual generation gives out-of-sample sqrt PEHE 0.64 on IHDP versus 0.87 for gradient ascent in input space and 0.72 for diffusion-based generation.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use FCCL instead of CBRE | ITE estimation on IHDP · IHDP, out-of-sample | rooted PEHE: 0.66(0.07) -> 0.64(0.07) | Helped |
| use FCCL instead of CBRE | ITE estimation on IHDP · IHDP, out-of-sample | absolute ATE error: 0.13(0.02) -> 0.12(0.02) | Helped |
| use FCCL instead of CFR-Wass | ITE estimation on Synthetic · Synthetic, γ = 0.4 | out-of-sample rooted PEHE: 2.48(0.06) -> 2.58(0.06) | Hurt |
| use FCCL instead of ABCEI | ITE estimation on Synthetic · Synthetic, γ = 0.7 | out-of-sample rooted PEHE: 3.57(0.09) -> 3.50(0.09) | Helped |
| use FCCL instead of CFR-Wass | ITE estimation on Synthetic · Synthetic, γ = 1.0 | out-of-sample rooted PEHE: 4.72(0.14) -> 4.49(0.13) | Helped |
| use FCCL instead of ABCEI | ITE estimation on Synthetic · Synthetic, γ = 1.2 | out-of-sample rooted PEHE: 5.19(0.12) -> 5.12(0.12) | Helped |
| replace flow-based counterfactual generation with gradient ascent in X-space (ablation) | FCCL on IHDP · IHDP, out-of-sample | rooted PEHE: 0.64(0.07) -> 0.87(0.19) | Hurt |
| replace flow-based counterfactual generation with GAN generation (ablation) | FCCL on IHDP · IHDP, out-of-sample | rooted PEHE: 0.64(0.07) -> 0.84(0.17) | Hurt |
| replace flow-based counterfactual generation with diffusion generation (ablation) | FCCL on IHDP · IHDP, out-of-sample | rooted PEHE: 0.64(0.07) -> 0.72(0.10) | Hurt |
Surprising findings
- Diffusion-based counterfactual generation is close to FCCL in within-sample error, but is worse out-of-sample, which the authors attribute to noise-driven generation drifting from the data manifold.
Limitations
- Evaluated only on synthetic, semi-synthetic (IHDP) and one real-world (Jobs) dataset; the paper states no limitations section.
- Boundary-sample analysis relies on a fixed definition (30 farthest samples from class centers).
Experimental setup
Datasets: IHDP Jobs Synthetic (generated multivariate normal covariates)
Benchmarks: IHDP Jobs Synthetic
Metrics: rooted PEHE ATE absolute error policy risk ATT bias
Baselines: OLS-1 OLS-2 BART KNN DML CFR-Wass CFR-MMD SITE CITE GANITE ABCEI CBRE
Builds on: CFRNet SITE CITE Dombrowski et al. manifold counterfactual explanations
Closest papers by meaning
- Conditional Counterfactual Mean Embeddings: Doubly Robust Estimation and Learning Rates (NeurIPS 2026 · similarity 0.89)
- Counterfactual Generative Modeling with Variational Causal Inference (ICLR 2025 · similarity 0.88)
- Causal Contrastive Learning for Counterfactual Regression Over Time (NeurIPS 2024 · similarity 0.88)
- Continuous Treatment Effect Estimation Using Gradient Interpolation and Kernel Smoothing (AAAI 2024 · similarity 0.88)
- Learning Counterfactual Outcomes Under Rank Preservation (NeurIPS 2025 · similarity 0.88)
- Adversarially Balanced Representation for Continuous Treatment Effect Estimation (AAAI 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.