Debiased Counterfactual Generation via Flow Matching from Observations

TL;DR

This paper introduces a flow-matching approach for debiased counterfactual distribution estimation, leveraging the relationship between observational and counterfactual data to improve accuracy and robustness.

Contribution

It proposes a novel deconfounding flow method with an efficient estimator, extending to high-dimensional settings, outperforming existing methods.

Findings

Deconfounding flows outperform existing estimators.

The method effectively mitigates failure modes of flow-based approaches.

Counterfactual and observational distributions share key properties under weak confounding.

Abstract

Estimating counterfactual distributions under interventions is central to treatment risk assessment and counterfactual generation tasks. Existing approaches model the counterfactual distribution as a standalone generative target, without exploiting its relationship to the observational data. In this work, we show that under standard assumptions, observational and counterfactual outcome distributions are tightly linked: they have identical support and tail behavior, remain statistically close under weak confounding, and share any features of high-dimensional outcomes which are invariant to confounders. These properties motivate learning counterfactual distributions not from scratch, but via a deconfounding flow from the observational distribution. We formulate this problem via flow-matching and derive a semiparametrically efficient estimator based on a novel efficient influence function correction. We subsequently extend our estimator to target minimal-energy flows in high-dimensions, which we show can be especially simple targets between observational and counterfactual distributions. In experiments, deconfounding flows outperform existing debiased counterfactual distribution estimators, while also mitigating known failure modes of flow-based methods.