The model¶

A TRAM-DAG is one triangular normalizing flow whose sparsity is a causal DAG. This page states the model in the notation of the paper and maps each part of it onto the constructors in tramdag.spec. For the symbols see notation.md; to read a fitted model see interpretation.md; for a worked fit see the demo notebook.

The transformation function¶

The variables have a causal ordering. For each variable a TRAM-DAG fits a monotone transformation function $h$. The function maps the observed value to a latent scale, conditional on the parents of that variable:

\[ \begin{aligned} u_1 &= h(x_1) \\ u_2 &= h(x_2 \mid x_1)\\ u_3 &= h(x_3 \mid x_1, x_2) \\ u_p &= h(x_p \mid x_1, x_2, \dots, x_{p-1}) \end{aligned} \]

This is the observed-to-latent direction. It is Eq. 2 of the paper, $F_{X\mid\mathrm{pa}}(x)=F_U\big(h(x\mid\mathrm{pa})\big)$, and in code it is u = h(x) + shift.

Training evaluates $h$ directly to score the likelihood, which is cheap. Sampling runs the inverse $x_i = h^{-1}(u_i \mid \mathrm{pa})$, which has no closed form: bracketed bisection solves it, so sampling costs more.

Together the $h$'s form one triangular flow. Each variable depends only on its causal parents $\mathrm{pa}(x_i)$, which is a subset of its predecessors. The Jacobian sparsity of the flow is therefore the DAG itself.

Each $h(x_i \mid \mathrm{pa})$ is one transformation model of $x_i$ given its parents. In code it is a Node, which also fits on its own, without the DAG (fitting.md).

The latents $u_1,\dots,u_p$ are standard logistic. That choice is what makes the fitted parameters interpretable, because a shift on the latent scale is then a log-odds ratio.

The components¶

The transformation decomposes additively on the latent scale, which is what keeps the interpretation valid. Each node's $h$ is

\[ u_i \;=\; h(x_i \mid \mathrm{pa}(x_i)) \;=\; \underbrace{h_{\boldsymbol{\vartheta}}(x_i)}_{\text{intercept}} \;+\; \underbrace{\textstyle\sum_j \beta_{ij}\, x_j}_{\text{linear shifts}} \;+\; \underbrace{\textstyle\sum_k g_{ik}(x_k)}_{\text{complex shifts}} \;+\; \underbrace{(\beta_0 + b_{\boldsymbol{\Theta}}(x_{\text{mod}}))\, x_t} _{\text{varying coefficient}} , \]

and every causal parent enters through exactly one term. Take $x_5$ with parents $\mathrm{pa}(x_5) = \{x_1, x_2, x_4\}$ as the example.

Simple intercept, SI. $h_{\boldsymbol{\vartheta}}(x_5)$ has constant parameters. It is a flexible monotone baseline, a Bernstein polynomial by default, and it is the same for every observation.
Complex intercept, CI. The parameters $\boldsymbol{\vartheta}$ are themselves a function of some parents. The whole transformation bends with the parent, which permits interactions that an additive shift cannot express.
Linear shift, LS. $\beta_{51} x_1 + \beta_{52} x_2$. One interpretable number per parent.
Complex shift, CS. $g(x_4)$, an unrestricted network of the parent. It stays additive on the latent scale.
Varying coefficient, VC. $(\beta_0 + b_{\boldsymbol{\Theta}}(x_{\text{mod}}))\, x_t$: a treatment $x_t$ whose effect $\beta$ varies with modifier covariates through a small penalized network. varying-coefficients.md is its guide.

Sampling has to invert only the intercept, because the shifts move to the other side:

\[ x_5 = h^{-1}(u_5 \mid x_1, x_2, x_4) = h_{\boldsymbol{\vartheta}}^{-1}\Big(u_5 - \underbrace{\beta_{51} x_1 + \beta_{52} x_2}_{\text{LS}} - \underbrace{g(x_4)}_{\text{CS}}\Big). \]

The components in code¶

Each node declares its transformation as an additive formula of terms. The formula is the node's first argument, a + sum of terms. Each constructor names the parents its term depends on. The paper's symbols I, LS, CS, VC are the classes Intercept, LinearShift, ComplexShift, VaryingCoefficient.

formula for a continuous node $x_3$	$u_3 = h(x_3 \mid \mathrm{pa})$
nothing, or `I()`	$h_{\vartheta}(x_3)$, the simple intercept
`LS("X1")`	$h_{\vartheta}(x_3) + \beta x_1$
`I("X1")`	$h_{\vartheta(x_1)}(x_3)$, the complex intercept
`CS("X1")`	$h_{\vartheta}(x_3) + g_1(x_1)$
`LS("X1") + CS("X2")`	$h_{\vartheta}(x_3) + \beta x_1 + g_2(x_2)$
`CS("X1", "X2")`	$h_{\vartheta}(x_3) + g_{12}(x_1, x_2)$, one joint network
`CS("X1") + CS("X2")`	$h_{\vartheta}(x_3) + g_1(x_1) + g_2(x_2)$, two additive networks
`I("X1", "X2")`	$h_{\vartheta(x_1, x_2)}(x_3)$, one joint network
`I("X1", "X2", allow_interaction=False)`	$h_{\vartheta(x_1) + \vartheta(x_2)}(x_3)$, one network per parent, summed in coefficient space
`CS("X1") + VC("X2", t="T")`	$h_{\vartheta}(x_3) + g_1(x_1) + \beta(x_2)\, x_T$

Three rules follow from the model.

Exactly one intercept, first. A formula written without one gets I() prepended. I() without parents is the simple intercept, I(...) with parents the complex one; SI() and CI(...) spell that out with the arity checked.
Joint versus additive is argument grouping. Several parents in one term form one network over all of them, an interaction. The same parents in separate terms act additively. A node takes at most one intercept with parents, so for the intercept allow_interaction sets the grouping.
Every parent enters through exactly one edge-owning term. The VC modifiers are the one exception: CS("X2") + VC("X2", t="T") is the intended pattern. There $x_2$ acts prognostically through the shift and also modifies the treatment effect.

The three knobs on a term¶

transform= on I picks the class of $h_{\boldsymbol{\vartheta}}$ for a continuous node:

"bernstein", the default: a Bernstein polynomial whose tails extrapolate along the boundary slope;
"spline": a monotone rational-quadratic spline, whose tail slope has a fixed value (zuko-upstream.md explains the consequence);
"affine": location and scale only, so the node-conditional is a logistic GLM.

The keyword arguments of the transform class go in transform_kwargs, I(transform="spline", transform_kwargs={"bins": 16}). Each transform pre-scales the data from the training range_q and $1 - $range_q quantiles onto a fixed domain; range_q=0 uses the minimum and maximum. Ordinal nodes have no transform to pick.

input_transform= on I, CS, VC transforms that term's continuous network inputs, with statistics frozen at the first fit: "minmax", "standardize", or a callable fn(x, train). LS and the VC treatment stay raw, so their coefficients keep their units.

units=, activation=, batch_norm= on I, CS, VC size the term's network. The defaults are in the API pages of the terms.

Ordinal nodes¶

An ordinal node has no monotone transform to choose. Its intercept is the cutpoint vector of an ordered logit,

\[ P(Y \le k \mid \mathrm{pa}) = \sigma(\vartheta_k - s(\mathrm{pa})), \]

with increasing cutpoints $\vartheta$ and $s$ the total shift. The model subtracts the shift where a continuous node adds it. notation.md states the convention. interpretation.md shows what it does to the sign of a coefficient.

The probability of an observed level is the difference of two sigmoids. The code evaluates it in log space, which keeps a gradient where both sigmoids saturate (ordinal_log_prob).

Parents enter every network as features: a continuous parent raw, in one column; an ordinal parent one-hot, in one column per level. Abduction is exact for continuous nodes and truncated-logistic for ordinal ones, so flow.sample(u=flow.abduct(df)) reproduces df exactly, level-exactly for ordinal nodes. A counterfactual takes three steps: abduction, action and prediction (Pearl, 2009)¹. abduct gives the latents, do= sets the intervention, and sample(u=..., do=...) predicts.

What the model cannot do¶

The model fixes the latent distribution to standard logistic. It does not learn it.
A VC treatment is continuous or a binary ordinal node; a multi-level ordinal treatment is not supported. Propensity centering needs a binary ordinal treatment.
The DAG is an input. tramdag fits the mechanisms of a graph you supply, and does not discover the graph.
Identification is the usual causal one. No hidden confounding, and the graph has to be right.

Pearl, J. (2009). Causality: Models, reasoning, and inference (2nd ed.). Cambridge University Press. https://doi.org/10.1017/CBO9780511803161 ↩

formula for a continuous node \(x_3\)	\(u_3 = h(x_3 \mid \mathrm{pa})\)
nothing, or `I()`	\(h_{\vartheta}(x_3)\), the simple intercept
`LS("X1")`	\(h_{\vartheta}(x_3) + \beta x_1\)
`I("X1")`	\(h_{\vartheta(x_1)}(x_3)\), the complex intercept
`CS("X1")`	\(h_{\vartheta}(x_3) + g_1(x_1)\)
`LS("X1") + CS("X2")`	\(h_{\vartheta}(x_3) + \beta x_1 + g_2(x_2)\)
`CS("X1", "X2")`	\(h_{\vartheta}(x_3) + g_{12}(x_1, x_2)\), one joint network
`CS("X1") + CS("X2")`	\(h_{\vartheta}(x_3) + g_1(x_1) + g_2(x_2)\), two additive networks
`I("X1", "X2")`	\(h_{\vartheta(x_1, x_2)}(x_3)\), one joint network
`I("X1", "X2", allow_interaction=False)`	\(h_{\vartheta(x_1) + \vartheta(x_2)}(x_3)\), one network per parent, summed in coefficient space
`CS("X1") + VC("X2", t="T")`	\(h_{\vartheta}(x_3) + g_1(x_1) + \beta(x_2)\, x_T\)