Wasserstein Causal Forests for Distribution-Valued Outcomes

Hugo Gobato Souto

stat.MLSep 27, 2026

Wasserstein Causal Forests for Distribution-Valued Outcomes

Authors: Hugo Gobato Souto

Organizations: University of São Paulo

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 \citep{word1990state}, 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.

Figures & tables

Figure 1: Simulation accuracy on selected designs at

n=1000

and

R=50

paired replications. Law points are mean MMD 2 ; the other panels are Monte Carlo RMSEs. Bars show one Monte Carlo standard error. LS1 changes location and shape, LS3 combines the same effect with limited overlap, S4 has a multimodal outer law with a heterogeneous location effect, and S5 has heterogeneous shape effects. All vertical axes use a log scale.

Figure 2: Project STAR calibrated functional contrasts with influence-function standard errors (bars). Open circles report the permuted-treatment placebo, and the dashed line marks the experimental literature value of about one fifth of a within-grade standard deviation for the mean effect. Skewness is dimensionless; all other contrasts are in within-grade standard deviations.

Figure 3: Project STAR arm densities relative to the pooled control benchmark, by school free-lunch quartile. Panels Q1 to Q4 and the marginal panel show unadjusted equal-unit mixtures of classroom-grade laws. The final panel compares raw and calibrated reference contrasts. The treated distribution separates most clearly in Q4, where the control law lies below the blended benchmark and treatment moves it toward that benchmark.

Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Design	Method	Law	TATE	TCATE	Ref. TATE	Ref. TCATE
LS0	WCF	0.0154(0.0002)	0.0134(0.0011)	0.0289(0.0012)	0.0156(0.0017)	0.0321(0.0020)
LS0	Causal-DRF	0.0483(0.0006)	0.1451(0.0027)	0.1452(0.0027)	0.1714(0.0036)	0.1715(0.0035)
LS0	DRF	0.0654(0.0007)	0.1396(0.0024)	0.1396(0.0024)	0.1653(0.0031)	0.1653(0.0031)
LS1	WCF	0.0176(0.0003)	0.0148(0.0010)	0.0296(0.0012)	0.0165(0.0014)	0.0311(0.0016)
LS1	Causal-DRF	0.0466(0.0006)	0.1250(0.0028)	0.1302(0.0028)	0.1457(0.0036)	0.1526(0.0037)
LS1	DRF	0.0605(0.0007)	0.1308(0.0022)	0.1347(0.0023)	0.1515(0.0029)	0.1566(0.0030)

Appendix

Table 1: Location, shape, and overlap designs at

n=1000

. Law entries are Monte Carlo means; functional-effect entries are Monte Carlo RMSEs. Parentheses give Monte Carlo SEs over fifty replications; lower is better.

Design	Method	Law	TATE	TCATE	Ref. TATE	Ref. TCATE
S1	WCF	0.0207(0.0006)	0.0152(0.0015)	0.0139(0.0006)	0.0112(0.0013)	0.0134(0.0007)
S1	Causal-DRF	0.0942(0.0020)	0.0312(0.0013)	0.0683(0.0009)	0.0113(0.0009)	0.0895(0.0011)
S1	DRF	0.1084(0.0020)	0.0179(0.0012)	0.0560(0.0011)	0.0110(0.0010)	0.0814(0.0012)
S2	WCF	0.0090(0.0001)	0.0206(0.0013)	0.0405(0.0017)	0.0207(0.0019)	0.0416(0.0017)
S2	Causal-DRF	0.0154(0.0002)	0.0483(0.0028)	0.0520(0.0026)	0.0197(0.0025)	0.0254(0.0020)
S2	DRF	0.0132(0.0001)	0.0214(0.0019)	0.0322(0.0017)	0.0194(0.0023)	0.0281(0.0018)

Appendix

Table 2: Law-separation and heterogeneous-effect designs at

n=1000

. Law entries are Monte Carlo means; functional-effect entries are Monte Carlo RMSEs. Parentheses give Monte Carlo SEs over fifty replications; lower is better.

Design	Method	Law	Mass	Mass TCATE
Z0	Two-part WCF	0.0136(0.0004)	0.0542(0.0016)	0.0421(0.0033)
Z0	Causal-DRF	0.0258(0.0007)	0.0969(0.0015)	0.0577(0.0037)
Z0	DRF	0.0303(0.0008)	0.1072(0.0018)	0.0585(0.0043)
Z1	Two-part WCF	0.0155(0.0004)	0.0541(0.0015)	0.0366(0.0024)
Z1	Causal-DRF	0.0254(0.0006)	0.0910(0.0013)	0.0422(0.0037)
Z1	DRF	0.0292(0.0007)	0.1026(0.0015)	0.0526(0.0038)

Appendix

Table 3: Structural-zero designs at

n=1000

. Mass is the error in the degenerate-component probability, and mass TCATE is its moderator-stratum contrast error. Entries are Monte Carlo mean (SE) over fifty replications; lower is better.

Figure 4: All ordinary designs at

n=1000

and fifty paired Monte Carlo replications. Bars are one Monte Carlo SE for each method separately. The vertical divider separates LS from S designs. Nonoverlap makes several LS and law-level rankings clear, and the narrower intervals sharpen most comparisons, but overlap remains for several S functional criteria. Because the displayed bars are not SEs of paired method differences, they should not be used as a formal paired test.

Figure 5: All structural-zero designs at

n=1000

and fifty paired Monte Carlo replications. Bars are one method-specific Monte Carlo SE. Two-part WCF is clearly separated on law and mass error; among the Z0 to Z2 mass-TCATE comparisons, Z1 remains within paired uncertainty.

Panel A: $K$ sensitivity at $M=10$
DGP	$K$	Law	Mean-quantile	Ref. TCATE plug-in	Ref. TCATE AIPW	Ref. TATE AIPW
LS0	5	0.0171(0.0003)	0.0489(0.0106)	0.0197(0.0058)	0.0341(0.0050)	0.0158(0.0056)
LS0	25	0.0154(0.0002)	0.0538(0.0098)	0.0205(0.0062)	0.0321(0.0052)	0.0156(0.0062)
LS0	49	0.0150(0.0002)	0.0486(0.0091)	0.0194(0.0058)	0.0330(0.0053)	0.0158(0.0056)
LS1	5	0.0194(0.0003)	0.1059(0.0046)	0.0655(0.0052)	0.0333(0.0048)	0.0176(0.0060)
LS1	25	0.0176(0.0003)	0.1079(0.0051)	0.0648(0.0062)	0.0311(0.0041)	0.0165(0.0050)

Appendix

Table 4: Sensitivity of WCF to the quantile grid size

K

and the particle budget

M

n=1000

for designs LS0 to LS3 under the confirmatory protocol. Panel A varies

K

M=10

; Panel B varies

M

K=25

. Law entries are Monte Carlo means and functional entries are Monte Carlo RMSEs over fifty paired replications, with Monte Carlo standard errors in parentheses; lower is better. The native-grid target changes with

K

, so Panel A is a stability check rather than a ranking of grid resolutions. WCF’s functional columns use its cross-fitted AIPW layer; the plug-in column is shown for comparison.

Panel A: LS1
$K$	$M=5$	$M=10$	$M=25$
5	0.0309(0.0042)	0.0333(0.0048)	0.0333(0.0047)
25	0.0326(0.0041)	0.0311(0.0041)	0.0315(0.0046)
49	0.0312(0.0044)	0.0307(0.0039)	0.0315(0.0042)
Panel B: LS3
5	0.0562(0.0065)	0.0540(0.0066)	0.0563(0.0067)

Appendix

Table 5: Interaction grid for the quantile grid size

K

and the particle budget

M

on LS1 and LS3 at

n=1000

over fifty paired replications. Entries are end-to-end reference-TCATE Monte Carlo RMSEs with Monte Carlo standard errors in parentheses; lower is better, rows are

K

, and columns are

M

DGP	Law	Mean-quantile	Ref. TCATE plug-in	Ref. TCATE AIPW	Ref. TATE AIPW
LS0	-2.6(1.5)	-9.1(13.5)	-4.0(8.1)	2.6(3.2)	6.2(6.9)
LS1	-2.6(1.7)	-0.2(4.4)	-0.7(4.8)	-1.0(3.4)	-6.6(5.5)
LS2	-2.8(1.3)	1.3(5.9)	-4.2(3.2)	-1.1(2.8)	13.3(7.2)
LS3	-0.8(1.4)	-0.1(2.8)	2.0(2.7)	-3.6(2.9)	-1.8(5.9)

Appendix

Table 6: Seed-paired resolution differences at

n=1000

for designs LS0 to LS3, computed as

K=49

minus

K=25

M=10

over fifty paired replications. Entries are percentage changes, so a negative value means the larger grid has the smaller error; the value in parentheses is the Monte Carlo standard error of the paired percentage difference in percentage points. Law is a mean, the other columns are errors. The native-grid target changes with

K

, so these differences do not separate quadrature resolution from tail coverage.

Figure 6: Quantile-grid and particle-budget sensitivity of WCF at

n=1000

, one line per design (LS0 to LS3). The top row varies

K

M=10

and the bottom row varies

M

K=25

. Panels report native-grid law error (Monte Carlo means), mean-quantile error, and end-to-end reference-TCATE error (Monte Carlo RMSEs). Because the native target changes with

K

, the top row is a stability check rather than a ranking of grid resolutions. Bars are Monte Carlo standard errors over fifty paired replications; lower is better.

Panel A: $n=500$
DGP	WCF	Causal-DRF	DRF
SYM-RANDOM	0.0449(0.0056)	0.0522(0.0058)	0.0472(0.0051)
SYM-LIN	0.0491(0.0049)	0.0686(0.0064)	0.0592(0.0068)
SYM-NL	0.0562(0.0059)	0.0856(0.0060)	0.0720(0.0060)
SYM-MU	0.0487(0.0054)	0.0562(0.0062)	0.0502(0.0064)
Panel B: $n=1000$

Appendix

Table 7: Symmetric and nonlinear assignment designs at

n=500

and

n=1000

over fifty paired replications. Entries are end-to-end reference-TCATE Monte Carlo RMSEs with Monte Carlo standard errors in parentheses; lower is better. WCF uses its cross-fitted AIPW layer and the forest baselines use plug-in contrasts. SYM-RANDOM is a constant-propensity calibration control, SYM-LIN uses a linear assignment index, and SYM-NL and SYM-MU use nonlinear assignment with the same symmetric outcome laws.

Panel A: $n=500$
DGP	Logistic	Random forest	Gradient boosting	Oracle
SYM-NL	0.0562(0.0059)	0.0575(0.0060)	0.0575(0.0060)	0.0575(0.0060)
SYM-MU	0.0487(0.0054)	0.0487(0.0045)	0.0487(0.0045)	0.0487(0.0045)
Panel B: $n=1000$
SYM-NL	0.0420(0.0057)	0.0427(0.0055)	0.0427(0.0055)	0.0427(0.0055)
SYM-MU	0.0361(0.0045)	0.0348(0.0042)	0.0348(0.0042)	0.0348(0.0042)

Appendix

Table 8: Propensity-model sensitivity on the SYM-NL and SYM-MU designs at

n=500

and

n=1000

over fifty paired replications. Entries are end-to-end reference-TCATE Monte Carlo RMSEs with Monte Carlo standard errors in parentheses; lower is better. The columns replace the default logistic propensity with a random-forest factory, a gradient-boosting factory, and the oracle propensity, which is a diagnostic and not a feasible competitor. The outcome learner, grid, and particle budget are held fixed.

Figure 7: Assignment and propensity sensitivity of the end-to-end reference-TCATE error at

n=500

and

n=1000

. Left column: constant (SYM-RANDOM), linear (SYM-LIN), and nonlinear (SYM-NL, SYM-MU) assignment mechanisms under the default logistic propensity. Right column: the SYM-NL and SYM-MU designs under the default logistic propensity and under random-forest, gradient-boosting, and oracle propensity factories; the factories nearly coincide because the confirmatory AIPW layer is fixed across columns. Bars are Monte Carlo standard errors over fifty paired replications; lower is better.

Panel A: $n=500$
Target	WCF	Causal-DRF	DRF
TATE grid mean	0.0400(0.0038)	0.1684(0.0070)	0.1274(0.0074)
TCATE grid mean	0.0664(0.0028)	0.1739(0.0068)	0.1348(0.0065)
Ref. TATE	0.0256(0.0033)	0.0269(0.0033)	0.0269(0.0033)
Ref. TCATE	0.0518(0.0025)	0.0554(0.0037)	0.0624(0.0039)
Panel B: $n=1000$

Appendix

Table 9: Null companion designs with exactly zero treatment effects, so every entry is a false-effect error. Entries are end-to-end Monte Carlo RMSEs pooled over the six null companion designs and fifty paired replications, with Monte Carlo standard errors in parentheses; lower is better. WCF uses its cross-fitted AIPW layer and the forest baselines use plug-in contrasts. The first two rows use the grid-mean functional, and the last two rows use the reference-distance marginal and conditional targets.

Figure 8: False-effect error of the six null companion designs at

n=500

and

n=1000

, pooled over designs and fifty paired replications. Entries are end-to-end Monte Carlo RMSEs; every treatment effect is exactly zero, so each value is a false-effect error. Bars are Monte Carlo standard errors; lower is better.

Figure 9: Sample-size sensitivity averaged over designs within each family. Law panels report mean MMD 2 ; reference and mass-TCATE panels report Monte Carlo RMSE. The LS and S panels use ordinary WCF, whereas Z uses two-part WCF. Lines are descriptive family averages and bars are Monte Carlo standard errors over the ten paired replications; complete design-level results appear in Tables 10 to 12 .

Design	$n$	Method	Law	TATE	TCATE	Ref. TATE	Ref. TCATE
LS0	125	WCF	0.0300(0.0019)	0.0616(0.0070)	0.1264(0.0186)	0.0699(0.0111)	0.1486(0.0301)
LS0	125	Causal-DRF	0.1136(0.0035)	0.2094(0.0150)	0.2095(0.0149)	0.2497(0.0166)	0.2498(0.0165)
LS0	125	DRF	0.1520(0.0044)	0.2241(0.0203)	0.2241(0.0203)	0.2658(0.0223)	0.2658(0.0223)
LS0	250	WCF	0.0255(0.0009)	0.0369(0.0056)	0.0686(0.0064)	0.0399(0.0078)	0.0715(0.0090)
LS0	250	Causal-DRF	0.0899(0.0026)	0.2067(0.0089)	0.2070(0.0089)	0.2388(0.0114)	0.2392(0.0114)
LS0	250	DRF	0.1330(0.0031)	0.2180(0.0102)	0.2180(0.0102)	0.2518(0.0119)	0.2518(0.0119)

Appendix

Table 10: Sample-size sensitivity for the LS family. Law is mean MMD 2 ; other entries are Monte Carlo RMSE. Parentheses are Monte Carlo SEs over ten replications.

Design	$n$	Method	Law	TATE	TCATE	Ref. TATE	Ref. TCATE
S1	125	WCF	0.0877(0.0043)	0.0372(0.0087)	0.0620(0.0065)	0.0182(0.0059)	0.0702(0.0079)
S1	125	Causal-DRF	0.2911(0.0121)	0.1899(0.0159)	0.2750(0.0111)	0.0420(0.0078)	0.1804(0.0026)
S1	125	DRF	0.5065(0.0079)	0.2448(0.0228)	0.4023(0.0138)	0.0419(0.0078)	0.1925(0.0018)
S1	250	WCF	0.0573(0.0033)	0.0147(0.0047)	0.0325(0.0035)	0.0228(0.0043)	0.0363(0.0036)
S1	250	Causal-DRF	0.1930(0.0068)	0.0948(0.0091)	0.1532(0.0078)	0.0329(0.0090)	0.1499(0.0019)
S1	250	DRF	0.3209(0.0095)	0.1113(0.0123)	0.2980(0.0386)	0.0376(0.0095)	0.1873(0.0028)

Appendix

Table 11: Sample-size sensitivity for the S family. Law is mean MMD 2 ; other entries are Monte Carlo RMSE. Parentheses are Monte Carlo SEs over ten replications.

Design	$n$	Method	Law	Mass	Mass TCATE
Z1	125	Two-part WCF	0.0549(0.0047)	0.1465(0.0076)	0.1187(0.0204)
Z1	125	Causal-DRF	0.0732(0.0028)	0.1659(0.0031)	0.1125(0.0188)
Z1	125	DRF	0.1000(0.0041)	0.1958(0.0043)	0.1218(0.0259)
Z1	250	Two-part WCF	0.0380(0.0041)	0.1150(0.0090)	0.1050(0.0214)
Z1	250	Causal-DRF	0.0556(0.0040)	0.1422(0.0071)	0.1043(0.0308)
Z1	250	DRF	0.0875(0.0031)	0.1831(0.0038)	0.1193(0.0245)

Appendix

Table 12: Sample-size sensitivity for the Z family. Law is mean MMD 2 ; other entries are Monte Carlo RMSE. Parentheses are Monte Carlo SEs over ten replications.

Study	Unit	$n$ (treated/control)	Outcome ( $K{=}25$ )	Treatment	Benchmark	Outside $[0.05,0.95]$
Project STAR	classroom-grade	1,333 (522/811)	within-grade $z$ math scores	small class vs. regular/aide	pooled control classrooms	0%
State minimum wage	state-year	867 (117/750)	hourly wages, 2016 USD	real rise $>\$ 0.25$ , own-rate	Nordic (scaled); US 2016	12.3%
Kenya cash transfers	village	653 (328/325)	per-capita nondurable consumption	village randomization	pooled control villages	0%

Appendix

Table 13: Design summary of the three applied studies. Each outcome is represented by

K=25

midpoint quantiles, and the final column reports the share of fitted propensities outside the overlap band

[0.05,0.95]

Figure 10: Minimum-wage arm densities for the full sample and for the sample retained by the standard overlap band

[0.05,0.95]

. Each panel shows mixture densities of the unit laws, constructed as in Figure 3 , with the Nordic benchmark (dashed) and the US 2016 national distribution (dotted) on a log scale;

d_{0}

and

d_{1}

are the mean unit-level distances to the Nordic benchmark and

\Delta d=d_{1}-d_{0}

. Restricting to the band leaves the mixture shapes and the benchmark distances close to their full-sample values, and the wage distributions remain far apart on the dollar scale.

Figure 11: State minimum-wage calibrated contrasts by sample, with influence-function standard errors from the first fitting replicate. Diamonds report the sample retained by the standard overlap band

[0.05,0.95]

; the two reference rows use the Nordic benchmark and the US 2016 distribution. Contrasts are in 2016 USD per hour except skewness, which is dimensionless.

Figure 12: Kenya cash-transfer calibrated contrasts with influence-function standard errors (bars) and unadjusted raw contrasts (open diamonds). The right-hand column reports the two-sided permutation

p

-value for each raw contrast. Estimates are in KSH per capita except skewness, which is dimensionless; the last two rows are the reference contrasts to the pooled control benchmark capped at its

75

th percentile (calibrated layer only) and to the external KIHBS 2015/16 benchmark, the latter fitted as a separate variant.

Figure 13: Kenya cash-transfer village densities. Each arm curve is a Gaussian kernel density estimate of the equal-unit mixture of the village laws, formed from the pooled village quantile coordinates; benchmark curves apply the same estimator to their quantile vectors interpolated to a fine level grid. The bandwidth is Silverman’s rule applied to the pooled arm coordinates and is common to all curves in a panel. Densities are scaled to a common maximum in each panel and are unadjusted for covariates. Panel (a) overlays the two arms, panel (b) fades the control density and adds the pooled control-village benchmark, and panel (c) repeats panel (b) with the external national KIHBS 2015/16 benchmark. The control benchmark overlaps both arms, while the national benchmark lies to the right with a long right tail; the raw reference contrast is negative for both and is larger for the national benchmark.

Causal Effect Estimation

Causal effect estimation measures how outcomes would change under interventions, using experiments or observational data and explicit assumptions.

See all 266 papers

Extended Wasserstein-GAN Approach to Causal Distribution Learning: Density-Free Estimation and Minimax Optimality
May 11, 2026Shu Tamano, Masaaki ImaizumiCausal Effect Estimation Causal Counterfactual Generation
1The University of Tokyo · 2Japan Institute for Health Security · 3RIKEN Center for Advanced Intelligence Project +1
Wasserstein Policy Learning for Distributional Outcomes
Jun 17, 2026Yiyan Huang, Cheuk Hang Leung, Qi Wu +1Offline RL Wasserstein Barycenters
School of Statistics and Data Science & Institute of Big Data Research, Shanghai University of Finance and Economics, Shanghai, China
Transfer learning for causal forest
Jun 5, 2026Bérénice-Alexia Jocteur, Véronique Maume-Deschamps, Pierre RibereauCausal Effect Estimation Conditional Average Treatment Effect Estimation
ICJ, PSPM · Universit´e Lyon 1, Centrale Lyon, INSA Lyon · Universit´e Jean Monnet, CNRS, ICJ UMR5208, 69622 Villeurbanne, France. +1

May 11, 2026J/K move · Enter open · S save

Extended Wasserstein-GAN Approach to Causal Distribution Learning: Density-Free Estimation and Minimax Optimality

Shu Tamano, Masaaki Imaizumi

1The University of Tokyo · 2Japan Institute for Health Security · 3RIKEN Center for Advanced Intelligence Project+1

Causal Effect Estimation Causal Counterfactual Generation Minimax Estimation Causal Inference Generative Adversarial Networks Wasserstein Distance

Open paper AskPDF arXiv

Distributional causal inference requires estimating not only average treatment effects but also interventional outcome distributions, including quantiles, tail risks, and policy-dependent uncertainty. As a method for distributional causal inference, generative adversarial network (GAN)-based counterfactual methods are flexible tools for this task. However, these methods have several limitations. First, the objectives of certain techniques do not coincide with the statistical risk of the identifiable causal target, and therefore provide limited theoretical guarantees regarding estimable counterfactual distributions or optimality. Second, they tend to rely on unstable density-based methods, such as density ratio estimation. In this paper, we propose GANICE (GAN for Interventional Conditional Estimation) with several advantages: it (i) clarifies the conditional interventional distribution for each treatment--covariate state as the causal estimation target; (ii) estimates the conditional distribution such that its averaged Wasserstein risk is minimized; (iii) establishes minimax optimality. GANICE achieves these advantages through the introduction of the extended Wasserstein distance, the incorporation of a cellwise critic in its dual, and an optimality proof based on Besov space theory. Our experiments demonstrate that GANICE consistently outperforms existing methods.

Abstract

Causal Effect Estimation

Explore similar work

Extended Wasserstein-GAN Approach to Causal Distribution Learning: Density-Free Estimation and Minimax Optimality

Wasserstein Policy Learning for Distributional Outcomes

Transfer learning for causal forest