# RepFlow: Representation Enhanced Flow Matching for Causal Effect Estimation

arXiv: https://arxiv.org/abs/2605.05890 · HTML: https://mlsanity.com/paper/2605.05890/

## 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.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 23 of 23)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [Causal Effect Inference with Deep Latent-Variable Models](https://mlsanity.com/paper/1705.08821/) (2017)
- [Adapting Neural Networks for the Estimation of Treatment Effects](https://mlsanity.com/paper/1906.02120/) (2019)
- [Automated versus do-it-yourself methods for causal inference: Lessons learned from a data analysis competition](https://mlsanity.com/paper/1707.02641/) (2017)
- [Diffusion Models for Causal Discovery via Topological Ordering](https://mlsanity.com/paper/2210.06201/) (2022)
- [Semiparametric proximal causal inference](https://mlsanity.com/paper/2011.08411/) (2020)
- [Semiparametric counterfactual density estimation](https://mlsanity.com/paper/2102.12034/) (2021)
- [Optimal Transport for Treatment Effect Estimation](https://mlsanity.com/paper/2310.18286/) (2023)
- [Causal Diffusion Transformers for Generative Modeling](https://mlsanity.com/paper/2412.12095/) (2024)
- [Adversarial Balancing-based Representation Learning for Causal Effect Inference with Observational Data](https://mlsanity.com/paper/1904.13335/) (2019)
- [Convergence of Continuous Normalizing Flows for Learning Probability Distributions](https://mlsanity.com/paper/2404.00551/) (2024)
- [Learning Disentangled Representations for Counterfactual Regression via Mutual Information Minimization](https://mlsanity.com/paper/2206.01022/) (2022)
- [DiffPO: A causal diffusion model for learning distributions of potential outcomes](https://mlsanity.com/paper/2410.08924/) (2024)
- [Covariate balancing using the integral probability metric for causal inference](https://mlsanity.com/paper/2305.13715/) (2023)
- [X-Learner: Learning Cross Sources and Tasks for Universal Visual Representation](https://mlsanity.com/paper/2203.08764/) (2022)
- [Enhancing Causal Effect Estimation with Diffusion-Generated Data](https://mlsanity.com/paper/2504.03630/) (2025)
- [Wasserstein Distributionally Robust Nonparametric Regression](https://mlsanity.com/paper/2505.07967/) (2025)

## 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.