Wasserstein Causal Forests for Distribution-Valued Outcomes

arXiv:2609.35898 · stat.ML, cs.LG · Submitted 2026-09-27 · Read on arXiv

stat.ML, cs.LG

Submitted: 2026-09-27

Updated: 2026-09-27

License: http://creativecommons.org/licenses/by/4.0/

The gist: This paper proposes Wasserstein Causal Forests (WCF) for settings in which each unit's outcome is itself a probability distribution.

Terminology

Abstract

This paper proposes Wasserstein Causal Forests (WCF) for settings in which each unit's outcome is itself a probability distribution. This study also defines finite-grid transformed average and conditional average treatment effects, including a reference-distance contrast that asks whether treatment moves unit-level distributions toward a prespecified benchmark. Simulations cover null effects, location and shape changes, limited overlap, equal-mean but different laws, heterogeneous effects, multimodality, and structural zeros. WCF is most accurate on the conditional-law metric in most reported designs and sharply improves reference-effect estimation in the principal location-and-shape settings, but it is less accurate than the forest baselines for multimodal settings. WCF is applied to the famous Project STAR, revealing that small classes alter more than the mean: they raise within-grade mathematics achievement by 0.161 standard deviations on average (SE 0.028); but the gain is not a uniform location shift, it is larger in the upper part of the classroom score distribution (+0.179 at the ninetieth percentile versus +0.091 at the tenth) and, in descriptive stratum estimates, largest in the schools serving the most economically disadvantaged students (highest free-lunch quartile, +0.262, versus +0.092 to +0.164 elsewhere), while overall dispersion is essentially unchanged.

Sources

Related papers