Skip to content

Varying-coefficient treatment effects

A VC term gives a node a treatment-effect head with its own bias-variance budget. It contributes \(\beta(x)\, x_t\) with \(\beta(x) = \beta_0 + b_\Theta(x)\) to the node's shift, where \(b_\Theta\) is a deliberately small, penalized network of the modifiers. The point is not expressiveness. The point is that the flow estimates the effect function with care, not as a by-product. Every claim here runs in notebooks/varying_coefficients.py, which is the one place for the code, or in the tests that the text names.

Modifiers may appear twice in a node, as prognostic parents through CS or LS and as effect modifiers through VC. The design intends this pattern. Only the treatment named by t= owns its edge, and a second term that declares it raises an error.

Why not a joint CS over treatment and modifiers

For a binary treatment the multi-parent CS is equivalent in expressiveness. Any shift decomposes as \(s(x,t) = s(x,0) + [s(x,1) - s(x,0)]\, t\). But the CS form has no effect-specific regularization. The likelihood rewards a good average fit of \(s(x,t)\), and nothing rewards a smooth difference between the arms. So the read-out is the difference of two jointly fitted, unregularized networks. It amplifies random variation. On the vc_hetero process the VC term clears the recovery bar of the tests (see Validation). The CS reduced form stays below a correlation of 0.8 on the same protocol, although the model is in-class (test_cs_reduced_form_falls_short_of_the_bar). Causal forests and R-learners do not work because they target the effect. They work because they regularize it (Athey et al., 2019; Nie & Wager, 2021)1 2. VC brings that ingredient into the TRAM framework.

Semantics

  • Scale. \(\beta(x)\) lives on the node's latent scale, like an LS weight. A continuous node adds it, and an ordinal node subtracts it from the cutpoints. With no modifiers, VC(t="T") is LS("T") bit-exactly, so VC against LS is a nested question.
  • Penalty. The objective is \(\sum_i \mathrm{NLL}_i + \lambda \lVert b_\Theta \rVert^2\) on the total-NLL scale, with penalty= as \(\lambda\). It is a fixed Gaussian prior whose shrinkage vanishes as \(n\) grows, it never applies to \(\beta_0\), and \(\lambda \to \infty\) recovers the LS fit. Raise it when the modifiers are many or \(n\) is small.
  • Identification. A constant moves freely between \(\beta_0\) and \(b_\Theta\). The head's output layer is zero-initialized, so \(\beta(x) = \beta_0\) at step 0. After fit each node re-centers its head to mean zero over the training data, which preserves the function. \(\beta_0\) is therefore the main effect in the training population, the Colr reading when \(\beta\) is constant.
  • Warm start. flow.nodes[name].warm_start(train_df) fits the all-LS version of the node classically and copies the treatment weight into beta0. Call it before fit. Training then starts at the classical answer and learns only deviations. flow.warm_start(train_df) does this for every node with a VC term and returns the reports by node name.
  • Treatments. Continuous or binary ordinal, entering as its 0/1 level so that \(\beta\) is the identified level-1-against-0 contrast. The term is linear in the treatment.
  • Read-out. flow.varying_coef(df, node) evaluates \(\beta(x)\) in closed form, deterministic and free of the outcome. tramdag.plots.plot_varying_coef(flow, df, node, by=) draws it along one modifier. For a continuous outcome, \(\beta(x) = \operatorname{logit} P(Y \le y \mid x, do(T=1)) - \operatorname{logit} P(Y \le y \mid x, do(T=0))\) at every \(y\). This is an interventional quantity and needs no counterfactual. Because abduct inverts a continuous node exactly, \(\beta(x)\) also equals the abduction difference \(u(x,1,y) - u(x,0,y)\), which a test pins. For an ordinal outcome abduct samples the latent, so this identity does not apply.

Propensity centering: propensity="col"

With a propensity= column, the term contributes \(\beta(x)\,(t - \hat e(x))\) instead of \(\beta(x)\, t\). This is Robinson's R-learner orthogonalization (Nie & Wager, 2021; Robinson, 1988)2 3 inside the likelihood. In model-based forests, this ingredient is decisive for effect estimation under confounding (Dandl et al., 2024)4. That finding reproduces here on the confounded process, where the model deliberately under-specifies its prognostic part. The uncentered \(\hat\beta\) absorbs the misfit, and centering brings it back near the truth. tests/test_vc_centered.py requires at least a factor of two. If the prognostic part is correctly specified, centering changes little.

Centering guards against a misspecified prognostic part. This also applies to randomized data: with \(\hat e = 0.5\) the regressor is \(\pm\tfrac12\), and a misfit that both arms share cannot follow it. Centering does not guard against unmeasured confounding. The propensity uses only the observed parents of the treatment, so a hidden common cause of \(T\) and \(Y\) biases \(\beta(x)\) with or without centering.

The design is two-stage and frozen, because the naive versions are wrong.

  • Training uses out-of-fold propensities, one value per row. Compute them and pass them as the training-frame column that propensity= names. Any propensity model works if a fit that never saw a fold predicts that fold. This is the cross-fitting requirement of double machine learning (Chernozhukov et al., 2018)5. In-sample propensities reintroduce the own-observation bias and can be worse than no centering. The values enter the loss as frozen data, so no gradient reaches the treatment node and the per-node factorization stays intact. fit refuses a centered spec whose training frame lacks the column. The flow's validation NLL uses the live propensity of the fitted treatment node.
  • Inference, in log_prob, sample, abduct, pmf and scores, recomputes \(\hat e\) from the flow's own fitted treatment node on the current parent values, detached. Under do(T=t) the regressor becomes \(t - \hat e(x)\). The flow caches nothing.
  • Interpretation. \(\beta_0\) is the mean of \(\beta(x)\) over the training rows, with or without centering, because fit re-centers the head in both cases. varying_coef stays the same, because centering moves the regressor and not \(\beta\). Centering requires a binary ordinal treatment.

Validation

The vc_hetero process in tests/conftest.py is a logistic-shift SCM with a known \(\beta(x) = -1 + 0.8\,X_2 - 0.6\,X_3\). It has a nonlinear prognostic part and confounded assignment, and \(X_2\) is confounder and modifier at once. In this configuration the CS reduced form fails hardest. tests/test_vc_term.py requires three things. It requires a recovery correlation of at least 0.9 with 4500 training rows. It requires a fitted \(\beta_0\) that matches fit_classical under a large penalty. It requires the read-out identities. tests/test_vc_centered.py tests the centering claims on the confounded process. scores.md covers the scan that shortlists modifiers before you declare a VC term.

  1. Athey, S., Tibshirani, J., & Wager, S. (2019). Generalized random forests. The Annals of Statistics, 47(2). https://doi.org/10.1214/18-AOS1709 ↩

  2. Nie, X., & Wager, S. (2021). Quasi-oracle estimation of heterogeneous treatment effects. Biometrika, 108(2), 299–319. https://doi.org/10.1093/biomet/asaa076 ↩↩

  3. Robinson, P. M. (1988). Root-N-consistent semiparametric regression. Econometrica, 56(4). https://doi.org/10.2307/1912705 ↩

  4. Dandl, S., Haslinger, C., Hothorn, T., Seibold, H., Sverdrup, E., Wager, S., & Zeileis, A. (2024). What makes forest-based heterogeneous treatment effect estimators work? The Annals of Applied Statistics, 18(1). https://doi.org/10.1214/23-AOAS1799 ↩

  5. Chernozhukov, V., Chetverikov, D., Demirer, M., Duflo, E., Hansen, C., Newey, W., & Robins, J. (2018). Double/debiased machine learning for treatment and structural parameters. The Econometrics Journal, 21(1), C1–C68. https://doi.org/10.1111/ectj.12097 ↩