BCF posterior-mean ATE highly seed-sensitive on small, confounded data

[asmahani]

Hi maintainers,

While developing our metacausal package (a causal-ML ensemble that integrates stochtree's BCF via a thin adapter alongside ~9 other causal-ML estimators), we observed unexpectedly large seed-to-seed variability in BCF's posterior-mean ATE on small, confounded data — wider than what any of the other ~9 estimators in our default ensemble exhibits, and large enough that we'd appreciate your read on whether this is the expected small-data regime or potentially a sampler issue worth deeper investigation.

Self-contained reproducer

Runs in ~30s; no external data, no metacausal dependency:

import warnings
warnings.filterwarnings("ignore")
import numpy as np
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.model_selection import cross_val_predict
from threadpoolctl import threadpool_limits
from stochtree import BCFModel


def confounded_dgp(n=445, p=8, n_seed=0):
    """Small confounded DGP roughly mimicking LaLonde-style earnings data."""
    rng = np.random.default_rng(n_seed)
    X = rng.normal(size=(n, p))
    e = 1 / (1 + np.exp(-(1.5 * X[:, 0] - 0.7 * X[:, 1])))   # confounded propensity
    T = rng.binomial(1, e).astype(float)
    tau = 1500 + 800 * X[:, 2] - 500 * X[:, 3]               # heterogeneous, mean ~1500
    Y0 = np.maximum(np.random.default_rng(n_seed + 1).normal(
        loc=3000 + 1500 * X[:, 0], scale=4000, size=n), 0)
    Y = Y0 + T * tau
    return X.astype(float), T, Y.astype(float)


X, T, Y = confounded_dgp(n_seed=0)   # fixed dataset

for seed in [0, 1, 7, 13, 42, 100, 2026]:
    propensity = cross_val_predict(
        HistGradientBoostingClassifier(max_iter=200, random_state=int(seed)),
        X, T.astype(int), method="predict_proba", cv=5,
    )[:, 1]
    propensity = np.clip(propensity, 0.01, 0.99)

    m = BCFModel()
    with threadpool_limits(limits=1):
        m.sample(
            X_train=X, Z_train=T, y_train=Y,
            propensity_train=propensity,
            num_gfr=5, num_burnin=200, num_mcmc=200,
            general_params={"random_seed": int(seed) % (2**31)},
        )
    tau_hat = m.predict(X=X, Z=T, propensity=propensity,
                        type="posterior", terms="tau")
    print(f"seed={seed:>4}  BCF ATE={float(np.mean(tau_hat)):>+8.0f}")

Observed (stochtree 0.4.2, Python 3.13.3, macOS)

seed=   0  BCF ATE=   +2775
seed=   1  BCF ATE=   +3386
seed=   7  BCF ATE=   +2200
seed=  13  BCF ATE=   +1683
seed=  42  BCF ATE=     +87
seed= 100  BCF ATE=   +1821
seed=2026  BCF ATE=    -318

The dataset is fixed across runs; only the BCF MCMC seed (and the propensity-CV seed) varies. Range: ~$3,700 against a true ATE of ~$1,500.

The same pattern holds on the LaLonde job-training data (n=445) bundled with our package, where the per-seed BCF ATE swings from −$5,550 to +$4,085 with identical configuration. Notably, on LaLonde, going from stochtree 0.4.0 → 0.4.2 (with an unchanged metacausal call) shifted the seed=42 estimate from +$1,482 to −$5,550. The 0.4.2 release notes mention BCF propensity-handling (#334) and parametric-intercept (#326) fixes, so we believe those fixes had a real effect; the question is whether the post-fix sampler is more small-data sensitive than expected.

What we tried

• Increasing num_burnin and num_mcmc from 200 each to 1,000 each shrinks the seed range to ~$3,500 — helpful but the variability remains substantial.
• threadpool_limits(limits=1) around model.sample (above) — no change.
• A separately-seeded propensity vs. sharing the BCF seed — no change.
• Other components in our ensemble (DoubleMLIRM, CausalForestDML, several meta-learners) on the same data have seed-to-seed std in the $100-$360 range; BCF's is ~$3,300.

Question

Is this magnitude of seed-sensitivity expected for BCF in the n≈400-500 / low-overlap regime, or does it warrant deeper investigation given the recent 0.4.2 sampler fixes? Happy to run additional diagnostics — multi-chain BCF, alternative propensity configurations, longer chains, anything else that would help.

Thanks for stochtree — it's a fantastic addition to the open-source causal-ML ecosystem, and our package depends on it.

[asmahani]

Quick follow-up isolating the propensity pathway:

I re-ran the synthetic reproducer above and the same experiment on the LaLonde dataset (n=445), comparing propensity_train=cross_fit_HGB(...) vs. propensity_train=None, holding everything else (including the BCF MCMC seed) fixed.

Synthetic DGP, range across 7 seeds: $3,704 with propensity, $3,453 without

LaLonde, range across 7 seeds: $4,739 with propensity, $4,390 without

Removing propensity shifts the central tendency on the synthetic DGP (drops estimates by $200–$900, consistent with the propensity correctly debiasing) but barely affects LaLonde (Δ ranges ±$50–$350 with mixed sign). In both cases the seed-to-seed range is essentially unchanged.

So the variability is intrinsic to the sampler given a fixed dataset, not driven by the propensity input — which means the 0.4.2 propensity-handling changes (#334) are unlikely to be the proximate cause of the LaLonde 0.4.0→0.4.2 jump from +$1,482 → −$5,550. The parametric-intercept fix (#326) or some interaction of the two seems more likely. Hope this helps narrow things down.

[andrewherren]

Hi @asmahani, thank you for opening this issue! This is likely an issue with how the parametric treatment intercept term is handled when terms = "tau" is specified in the predict method, which I need to better document (I will post a vignette to the website soon).

In order to start reconciling this, can you let me know if you change the m.predict call in your provided example to specify terms="cate", does the ATE distribution stabilize?

[asmahani]

Confirmed — switching terms="tau" → terms="cate" collapses the seed-to-seed spread by an order of magnitude or more on both reproducers, with everything else held fixed:

	terms="tau" (range / std across 7 seeds)	terms="cate" (range / std)
Synthetic DGP (true ATE ≈ $1,500)	$3,414 / $1,182	$58 / $23
LaLonde (n=445)	$4,739 / $1,778	$212 / $66

The per-seed cate values land tightly around the expected magnitudes — synthetic at ~$1,400–$1,460 (within $50 of the DGP's $1,500 effect), LaLonde at ~$1,400–$1,610. So your hypothesis lines up with the observation: the parametric treatment intercept term was being excluded from tau but was the dominant contributor to the seed-to-seed swing.

I also saw #377 — looks like the right upstream fix; making tau always fold in tau_0 will close the gap for users who reach for terms="tau" expecting "the treatment effect."

In the meantime we patched our adapter on the consumer side (metacausal) to use terms="cate", which is the semantically-correct answer regardless of BCF setup (works for vanilla BCF, post-PR-#377 BCF, and BCF with random slopes on Z). Thanks for the very quick diagnosis.

[andrewherren]

Thank you for quickly closing the loop on this! Our reparameterization of BCF into parametric and nonparametric treatment effect terms definitely made the quantity returned by terms = "tau" counterintuitive in 0.4.2. #377 should make the interface clearer and more consistent going forward.

It might be a week or so before #377 is on PyPI and CRAN (as version 0.4.3), but once it's there, you should be able to revert to terms = "tau" in metacausal. For models with random slopes on Z (i.e. when rfx_group_ids_train is provided and 'model_spec': 'intercept_plus_treatment' is specified in random_effects_params), terms = "cate" will also include the random slopes on Z, so after 0.4.3, terms = "tau" will be the "correct" way to request just the covariate-dependent treatment effect function.

I am working on a few additional documentation pages on the stochtree site that spell out more precisely all of stochtree's modeling options and how they map to software flags in BART and BCF.