examples.demo_tram_dag_colab

TRAM-DAG in 5 minutes — one causal model, all three rungs of Pearl's ladder

Open in Colab

TRAM-DAGs (Sick & Dürr, CLeaR 2025) are interpretable neural causal models: one normalizing flow wired exactly like the adjacency matrix of your causal DAG — each variable is transformed conditional on its parents. Fit it once on observational data and you can

  1. L1 sample / score the observational distribution,
  2. L2 answer interventional queries — do(...) — by graph mutilation,
  3. L3 compute individual counterfactuals ("what would have happened to this unit?") via Pearl's abduction–action–prediction.

This demo uses the benchmark the paper leads with: a 3-variable SCM whose source is bimodal — the example where a default causal normalizing flow visibly fails (paper, Fig. 4) and TRAM-DAG doesn't. Ground truth is known analytically, so every claim below is checked, not asserted.

Runs on CPU; on a Colab GPU runtime (Runtime → Change runtime type → any GPU, e.g. the free T4) the final section races the two.

import importlib.util
import subprocess
import sys
import time

if importlib.util.find_spec("tramdag") is None:  # Colab: install from PyPI
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "tramdag"])
    # to track active development instead of the latest release, use:
    #   pip install "git+https://github.com/tensorchiefs/tramdag.git@main"

import matplotlib.pyplot as plt
import numpy as np
import scipy.stats as st  # for KDE plots only (preinstalled on Colab)
import torch

from tramdag import CausalFlowDAG, ContinuousNode, I
from tramdag.simulations import VacaTriangle

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
plt.rcParams["figure.dpi"] = 110
print(f"torch {torch.__version__}  device: {DEVICE}")

torch 2.12.0+cu130  device: cpu

1. The challenge: a bimodal structural causal model

The DGP (Sánchez-Martín et al. 2022, App. E.1; TRAM-DAG paper App. C.1):

$$ \begin{aligned} x_1 &\sim \tfrac12\,\mathcal N(-2,\,1.5) + \tfrac12\,\mathcal N(1.5,\,1) &&\text{(bimodal source)}\ x_2 &= -x_1 + \mathcal N(0,1)\ x_3 &= x_1 + 0.25\,x_2 + \mathcal N(0,1) \end{aligned} $$

The DAG is $x_1 \to x_2$, $x_1 \to x_3$, $x_2 \to x_3$.

gen = VacaTriangle(seed=42)
df = gen.observational(50_000)
train, val = df.iloc[:45_000], df.iloc[45_000:]

fig, axes = plt.subplots(1, 3, figsize=(11, 3))
for ax, c in zip(axes, df.columns):
    ax.hist(df[c], bins=80, density=True, alpha=0.7)
    ax.set_title(f"observed ${c[0]}_{c[1]}$")
fig.suptitle("50,000 observational samples — note the bimodal $x_1$")
fig.tight_layout()
plt.show()

png

2. Fit the TRAM-DAG — the spec is the DAG

Each node gets a monotone Bernstein transform; the edge labels say how parents enter (ci = the transform's parameters depend on the parents — maximal flexibility). Training maximizes the exact joint likelihood with one Adam; the new schedule="plateau" + freeze_patience lets every node decay its learning rate independently and drop out of training once converged — the fit stops itself.

spec = {
    "x1": ContinuousNode(),  # source
    "x2": ContinuousNode(terms=[I("x1")]),
    "x3": ContinuousNode(terms=[I("x1", "x2")]),
}

torch.manual_seed(0)
flow = CausalFlowDAG(spec, device=DEVICE)
t0 = time.perf_counter()
flow.fit(
    train,
    val,
    epochs=400,
    learning_rate=1e-1,
    batch_size=4096,
    verbose=50,
    schedule="plateau",
    plateau_patience=10,
    freeze_patience=30,
    marginal_init=True,
)
t_fit = time.perf_counter() - t0
print(
    f"\nfitted on {DEVICE} in {t_fit:.1f}s "
    f"({len(flow.history['val'])} epochs, then froze itself)"
)

[epoch     1/400] train NLL 6.3241  val NLL 5.3017


[epoch    51/400] train NLL 4.9131  val NLL 4.9197


[epoch    90] all nodes frozen — stopping.

fitted on cpu in 15.0s (90 epochs, then froze itself)

Training diagnostics come for free — fit records per-node train/val NLL, learning rates, wall-clock time and the freeze epochs. Watch the nodes drop out of training one by one:

hist = flow.history
ep = np.arange(1, len(hist["val"]) + 1)
tot_tr = np.array([sum(d.values()) for d in hist["train"]])
tot_va = np.array([sum(d.values()) for d in hist["val"]])
fig, ax = plt.subplots(figsize=(7.5, 3.6))
ax.plot(ep, tot_tr, label="train NLL (total)")
ax.plot(ep, tot_va, label="val NLL (total)")
for i, (name, e) in enumerate(
    sorted(hist.get("frozen", {}).items(), key=lambda kv: kv[1])
):
    ax.axvline(e, ls="--", lw=1, color="gray")
    ax.annotate(
        f" {name} frozen",
        (e, ax.get_ylim()[1]),
        rotation=90,
        va="top",
        fontsize=8,
        color="gray",
    )
ax.set_ylim(tot_va.min() - 0.02, tot_va.min() + 0.6)  # zoom past the initial drop
ax.set_xlabel("epoch"), ax.set_ylabel("NLL"), ax.legend()
ax.set_title("training curve — per-node plateau decay, then self-freezing")
fig.tight_layout()
plt.show()

png

3. Rung 1 — does it actually fit? (the plot the CNF baseline fails)

samp = flow.sample(5 * len(df), seed=1)
cols = list(df.columns)
fig, axes = plt.subplots(3, 3, figsize=(9, 8.5))
for i, ci in enumerate(cols):
    for j, cj in enumerate(cols):
        ax = axes[i][j]
        if i == j:
            bins = np.linspace(df[ci].quantile(0.001), df[ci].quantile(0.999), 70)
            # Plot DGP histogram
            ax.hist(df[ci], bins=bins, density=True, alpha=0.5, label="DGP")
            # KDE for TRAM-DAG samples
            kde = st.gaussian_kde(samp[ci][:50_000])
            x_eval = np.linspace(bins[0], bins[-1], 300)
            ax.plot(x_eval, kde(x_eval), color="C3", lw=1.8, label="TRAM-DAG KDE")
        else:
            ax.scatter(df[cj][:1500], df[ci][:1500], s=2, alpha=0.3)
            ax.scatter(samp[cj][:1500], samp[ci][:1500], s=2, alpha=0.3, color="C3")
        if i == 2:
            ax.set_xlabel(cj)
        if j == 0:
            ax.set_ylabel(ci)
axes[0][0].legend(fontsize=8)
fig.suptitle("L1: observational joint — DGP (blue) vs fitted TRAM-DAG (red, KDE)")
fig.tight_layout()
plt.show()

png

4. Rung 2 — interventions: do(x2 = a)

Graph mutilation: clamp $x_2$, cut its incoming edge, resample. Under the DGP, $x_3\,|\,do(x_2{=}a) = x_1 + 0.25a + \mathcal N(0,1)$, so $\mathbb E[x_3] = -0.25 + 0.25a$ analytically — a hard number to be wrong about.

fig, axes = plt.subplots(1, 3, figsize=(11, 3.2), sharey=True)
print("E[x3 | do(x2=a)]:   analytic    TRAM-DAG")
for ax, a in zip(axes, (-3.0, -1.0, 0.0)):
    truth = gen.interventional(50_000, {"x2": a})
    fl = flow.sample(50_000, do={"x2": a}, seed=2)
    bins = np.linspace(truth["x3"].quantile(0.001), truth["x3"].quantile(0.999), 70)
    # DGP histogram
    ax.hist(truth["x3"], bins=bins, density=True, alpha=0.5, label="DGP")
    # Flow density: KDE for a smoother estimate
    kde = st.gaussian_kde(fl["x3"].iloc[:50_000])
    x_eval = np.linspace(bins[0], bins[-1], 300)
    ax.plot(x_eval, kde(x_eval), color="C3", lw=1.8, label="TRAM-DAG density")
    ax.set_title(f"$p(x_3 \\mid do(x_2={a:+.0f}))$")
    print(
        f"   a = {a:+.0f}:          {-0.25 + 0.25 * a:+.3f}      {fl['x3'].mean():+.3f}"
    )
axes[0].legend()
fig.tight_layout()
plt.show()

E[x3 | do(x2=a)]:   analytic    TRAM-DAG


   a = -3:          -1.000      -1.019


   a = -1:          -0.500      -0.511


   a = +0:          -0.250      -0.270

png

5. Rung 3 — the counterfactual magic trick

Take 1,000 held-out individuals. Step 1 (abduction): invert the flow to recover each individual's latent noise $u$ — everything about them the model doesn't attribute to their parents. Sanity check: pushing $u$ back through the flow must reproduce the observed data exactly.

Step 2+3 (action + prediction): rerun history with $do(x_1 = 0)$ — same $u$, mutilated graph. Because the DGP is fully continuous, the true individual counterfactuals are known (the simulator keeps its noise), so we can score the flow per individual — the strictest test in causal inference, impossible with real data.

rng = np.random.default_rng(7)
lat = gen.draw_latents(1_000, rng)
factual = gen.simulate(latents=lat)
cf_true = gen.simulate(latents=lat, do={"x1": 0.0})

u = flow.abduct(factual)
recon = flow.sample(u=u)
err = np.abs(recon.to_numpy() - factual.to_numpy()).max()
print(f"abduction -> reconstruction: max |error| = {err:.2e}  (exact recovery)")

cf_flow = flow.sample(do={"x1": 0.0}, u=u)
fig, axes = plt.subplots(1, 2, figsize=(9, 3.6))
for ax, c in zip(axes, ["x2", "x3"]):
    ax.scatter(cf_true[c], cf_flow[c], s=4, alpha=0.4)
    lims = [cf_true[c].min(), cf_true[c].max()]
    ax.plot(lims, lims, "k--", lw=1)
    r = np.corrcoef(cf_true[c], cf_flow[c])[0, 1]
    ax.set_title(
        f"counterfactual ${c[0]}_{c[1]}$ under $do(x_1{{=}}0)$   (r = {r:.4f})"
    )
    ax.set_xlabel("true (DGP, shared noise)"), ax.set_ylabel("TRAM-DAG")
fig.suptitle("L3: individual counterfactuals, scored unit by unit")
fig.tight_layout()
plt.show()

abduction -> reconstruction: max |error| = 4.57e-06  (exact recovery)

png

6. Swapping the transform: Bernstein vs spline vs affine

Each continuous node owns a monotone 1-D transform — that's where the distributional flexibility lives. One constructor argument switches it: "bernstein" (default, TRAM-faithful polynomial), "spline" (monotone rational-quadratic, the neural-spline-flow building block), or "affine" (location–scale only → every node-conditional is forced to be a logistic — essentially a classical GLM). Same DAG, same training, three model families:

def make_spec(transform):
    return {
        "x1": ContinuousNode(transform=transform),
        "x2": ContinuousNode(transform=transform, terms=[I("x1")]),
        "x3": ContinuousNode(transform=transform, terms=[I("x1", "x2")]),
    }


fits = {"bernstein": flow}  # already trained above
for tr in ["spline", "affine"]:
    torch.manual_seed(0)
    f = CausalFlowDAG(make_spec(tr), device=DEVICE)
    f.fit(
        train,
        val,
        epochs=400,
        learning_rate=1e-2,
        batch_size=4096,
        verbose=0,
        schedule="plateau",
        plateau_patience=10,
        freeze_patience=30,
    )
    fits[tr] = f
print("held-out NLL (lower is better):")
for tr, f in fits.items():
    print(f"  {tr:9s}: {sum(f.nll(val).values()):.4f}")

held-out NLL (lower is better):
  bernstein: 4.9149
  spline   : 5.6333
  affine   : 5.0393

bins = np.linspace(df["x1"].quantile(0.001), df["x1"].quantile(0.999), 70)
fig, ax = plt.subplots(figsize=(7, 3.4))
ax.hist(df["x1"], bins=bins, density=True, alpha=0.4, color="gray", label="data")
for tr, color in [("bernstein", "C3"), ("spline", "C0"), ("affine", "C2")]:
    ax.hist(
        fits[tr].sample(len(df), seed=3)["x1"],
        bins=bins,
        density=True,
        histtype="step",
        lw=1.8,
        color=color,
        label=tr,
    )
ax.set_title("the affine (GLM-like) transform cannot bend into two modes")
ax.set_xlabel("$x_1$"), ax.legend()
fig.tight_layout()
plt.show()

png

Reading the table: affine pays exactly where you'd expect — a location–scale transform cannot produce a bimodal $x_1$ (the same failure mode as the inflexible CNF in the paper's Fig. 4). The RQ-spline is expressive enough in principle, but with the small TRAM-DAG parameter heads it consistently trains to a worse optimum on this target (same result for 8–32 bins, lr 0.01–0.1, up to 2000 epochs) — an honest empirical reason why Bernstein (whose monotone softplus-cumsum parametrization is easier to optimize) is the TRAM-faithful default. Swap per node any time via ContinuousNode(transform="spline", transform_kwargs={"bins": 16}).

7. GPU vs CPU

The whole flow is plain PyTorch, so it runs anywhere. Same 60-epoch fit, both devices (on a CPU-only runtime this just reports CPU):

def timed_fit(device, epochs=60):
    torch.manual_seed(0)
    f = CausalFlowDAG(
        {
            "x1": ContinuousNode(),
            "x2": ContinuousNode(terms=[I("x1")]),
            "x3": ContinuousNode(terms=[I("x1", "x2")]),
        },
        device=device,
    )
    t0 = time.perf_counter()
    f.fit(train, val, epochs=epochs, learning_rate=1e-2, batch_size=4096, verbose=0)
    return time.perf_counter() - t0


timings = {"cpu": timed_fit("cpu")}
if torch.cuda.is_available():
    timed_fit("cuda", epochs=5)  # warm-up (cuda kernel compilation)
    timings["cuda"] = timed_fit("cuda")
for dev, t in timings.items():
    print(f"{dev:5s}: {t:6.1f}s for 60 epochs @ n=45,000")
if len(timings) > 1:
    fig, ax = plt.subplots(figsize=(4, 2.8))
    ax.barh(list(timings), list(timings.values()), color=["C0", "C2"])
    ax.set_xlabel("seconds (60 epochs)"), ax.set_title("same model, same data")
    fig.tight_layout()
    plt.show()

cpu  :   10.0s for 60 epochs @ n=45,000

What you just saw

rung query call checked against
L1 observational joint flow.sample(n) 50k DGP samples (bimodal $x_1$ ✓)
L2 $p(x_3 \mid do(x_2{=}a))$ flow.sample(n, do=...) analytic $\mathbb E[x_3] = -0.25 + 0.25a$
L3 individual counterfactuals flow.abduct(df) + flow.sample(do=..., u=u) per-unit DGP truth, r ≈ 0.99+

And the same model family stays interpretable when you want it to be: declare an edge "ls" instead of "ci" and its coefficient is a log-odds ratio you can read off after training (that's the actual point of the paper).

More: repo · paper · didactic walkthrough: notebooks/intro_tram_dag.py

1""".. include:: ../../nbmd/demo_tram_dag_colab.md"""