mlsanityData to 14 Oct 2026

RepFlow: Representation Enhanced Flow Matching for Causal Effect Estimation

arXiv 2605.05890 ↗First posted 7 May 2026cs.LG stat.MEcited by 0

On the map: Diffusion models → Diffusion foundations and sampling

Abstract (arXiv)

Estimating causal effects from observational data has become increasingly critical in diverse fields including healthcare, economics, and social policy. The fundamental challenge in causal inference arises from the missing counterfactuals and the selection bias. Existing methods are largely limited to point estimates and lack the capacity for distribution modeling. In this work, we propose RepFlow, a novel framework that formulates causal effect estimation as a joint optimization problem integrating representation learning with Conditional Flow Matching (CFM). RepFlow mitigates selection bias by minimizing the entropically regularized Wasserstein distance between treated and control representations. To enhance numerical stability, we further introduce an L2 normalization constraint on latent representations. This balanced representation enables the flow model to accurately capture the distribution of potential outcomes. Extensive experiments across a wide range of benchmarks demonstrate that RepFlow consistently outperforms existing methods in both point and distributional causal effect estimation.

Key facts. No citing papers in the map yet (data through 14 Oct 2026). Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models. Builds on 23 of its 48 references that are in the map. First posted on arXiv 7 May 2026.

Builds on

23 of 48 references.

Closest ancestors · 3

Same direction, newest first

  1. 2024DiffPO: A causal diffusion model for learning distributions of potential outcomes11 Oct
  2. 2022Diffusion Models for Causal Discovery via Topological Ordering12 Oct
  3. 2020Denoising Diffusion Probabilistic Models19 Jun

Other references · 20

  1. 2022Flow Matching for Generative Modelingcited by 5,702
  2. 2022Generative Adversarial Networkscited by 5,627
  3. 2018Neural Ordinary Differential Equationscited by 3,779
  4. 2018UMAP: Uniform Manifold Approximation and Projection for Dimension Reductioncited by 2,898
  5. 2015Variational Inference with Normalizing Flowscited by 2,170
  6. 2015Estimation and Inference of Heterogeneous Treatment Effects using Random Forestscited by 600
  7. 2017Causal Effect Inference with Deep Latent-Variable Modelscited by 410
  8. 2019Adapting Neural Networks for the Estimation of Treatment Effectscited by 287
  9. 2017Automated versus do-it-yourself methods for causal inference: Lessons learned from a data analysis competitioncited by 106
  10. 2020Semiparametric proximal causal inferencecited by 59
  11. 2021Semiparametric counterfactual density estimationcited by 42
  12. 2023Optimal Transport for Treatment Effect Estimationcited by 32
8 more
  1. 2024Causal Diffusion Transformers for Generative Modelingcited by 31
  2. 2019Adversarial Balancing-based Representation Learning for Causal Effect Inference with Observational Datacited by 30
  3. 2024Convergence of Continuous Normalizing Flows for Learning Probability Distributionscited by 26
  4. 2022Learning Disentangled Representations for Counterfactual Regression via Mutual Information Minimizationcited by 19
  5. 2023Covariate balancing using the integral probability metric for causal inferencecited by 9
  6. 2022X-Learner: Learning Cross Sources and Tasks for Universal Visual Representationcited by 6
  7. 2025Enhancing Causal Effect Estimation with Diffusion-Generated Datacited by 4
  8. 2025Wasserstein Distributionally Robust Nonparametric Regressioncited by 3

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2025Flow-based Generative Modeling of Potential Outcomes and Counterfactuals
  2. 2023Causal normalizing flows: from theory to practice
  3. 2025Deep Disentangled Representation Network for Treatment Effect Estimation
  4. 2026Statistical Inference via Generative Models: Flow Matching and Causal Inference
  5. 2022Personalized Public Policy Analysis in Social Sciences using Causal-Graphical Normalizing Flows
  6. 2025The Third Pillar of Causal Analysis? A Measurement Perspective on Causal Representations
  7. 2025Transportability without Graphs: A Bayesian Approach to Identifying s-Admissible Backdoor Sets
  8. 2024Conditional Generative Models are Sufficient to Sample from Any Causal Effect Estimand