Chapter 14
Bayesian deep learning and causal AI

Lectures 30–33 · Bayesian neural networks as causal models, transformers as Bayesian machines, learning causes instead of shortcuts, and teaching AI to model causation

Chapter 13 ended with three requirements for any learner that is to model causation rather than association: causal assumptions, written down; variation that comes from interventions or from several environments; and honest uncertainty that widens where the data cannot speak. This final chapter takes the three into deep learning, the technology behind modern AI, and asks how far each can be met.

The answer is more encouraging than one might fear and less magical than one might hope. Neural networks can be fitted the Bayesian way, and when they are used as the structural equations of a causal model their flexibility is a real asset. Transformers, the architecture of large language models, turn out to be machines that can learn Bayesian inference, and even causal inference, from simulated worlds. But no architecture removes the need for the assumptions of Chapter 13. A network trained on the wrong prior, or on data with a hidden confounder, learns the wrong answer confidently. The chapter's thesis is that making AI smarter about causes means building the causal assumptions into what it is trained on, how it is trained and what it is asked to report.

• Lecture 30: Bayesian neural networks for causal questions. Priors over functions; HMC, ensembles and their uncertainty; a two-headed network for individual treatment effects on the IHDP benchmark of Chapter 13.
• Lecture 31: transformers as Bayesian machines. Attention as kernel regression; a one-head attention model fitted in Stan, then a small transformer trained in PyTorch, that learn to do in-context Bayesian regression; TabPFN applied to the treatment-effect benchmark of Lecture 30; prior-fitted networks and in-context learning.
• Lecture 32: learning causes, not shortcuts. Shortcut learning as a causal failure; invariance across environments; an invariance-screened detector of AI-written text, tested on kinds of question it never saw.
• Lecture 33: teaching AI to model causation. Whether language models can climb the ladder, with two open models tested on the CLadder benchmark; a network trained on simulated causal worlds to estimate effects in context; interventions inside networks; and what it would take.

Lecture 30 — Bayesian neural networks for causal questions

14.1 Networks as priors over functions

A neural network with one hidden layer of $H$ units computes

where $\varphi$ is a fixed nonlinearity such as $\tanh$, and the weights $w_h$, biases $b_h$ and output weights $v_h$ are parameters. Deep networks stack several such layers. Machine learning usually fits the weights by minimising a loss. The Bayesian alternative puts a prior on the weights, and since every draw of the weights is a function, a prior on the weights is a prior on functions. Radford Neal (1996) showed what that prior looks like for a wide network: if the output weights have prior standard deviation proportional to $1/\sqrt{H}$, then as $H\to \infty$ the prior over functions becomes a Gaussian process (Chapter 8), by the central limit theorem applied to the sum over units. A finite network's prior over functions is close to a Gaussian process but not exactly one: it has extra, non-Gaussian structure. The widget draws functions from the prior.

The posterior over the weights is hard to work with. It is high-dimensional, and it has exact symmetries: swapping two hidden units, or flipping the signs of a $\tanh$ unit's input and output weights, leaves the function unchanged. So every posterior mode has $H!2^H$ copies, and different chains land in different ones. This is the label switching of Chapter 11, multiplied. The practical consequence is that the convergence of a network's chains must be judged on functions (predictions and effects), which are identified, not on the weights, which are not.

There are four common ways to approximate the posterior, and they differ mainly in how much of its spread they capture.

• Hamiltonian Monte Carlo (Neal, 1996), feasible for small networks such as the one we fit below.
• Deep ensembles (Lakshminarayanan, Pritzel and Blundell, 2017): train the network several times from different random starting weights and treat the trained networks as draws from the posterior. Each run lands in a different mode, which captures the multimodality, but not the spread within each mode.
• The Laplace approximation (MacKay, 1992; Daxberger et al., 2021): a Gaussian around one trained network, cheap and surprisingly effective.
• Variational inference (Chapter 9) and Monte Carlo dropout (Gal and Ghahramani, 2016), which is approximately variational, scale to the largest networks and tend to be overconfident.

The widget trains an ensemble of small networks in your browser on data with a gap.

The networks agree where there are data and disagree in the gap and beyond the ends, so the ensemble's spread grows where it should. But the band is often too narrow in the gap compared with what a Gaussian process with the same smoothness would give, and with few networks it varies from one ensemble to the next. Ensembles describe the disagreement between plausible functions, which is a large part of epistemic uncertainty, but not all of it.

14.2 A Bayesian network for individual effects

Chapter 13 used BART as the outcome model in g-computation and scored it on the IHDP benchmark (Hill, 2011; replication 1 of those distributed by Louizos et al., 2017), where every child's true effect is known. A neural network can play the same role. A good architecture for the job is the two-headed TARNet of Shalit, Johansson and Sontag (2017): a shared hidden layer learns a representation $\Phi (x)$ of the covariates, and two separate output heads predict the outcome under control and under treatment, ${\mu }_0(x)=h_0(\Phi (x))$ and ${\mu }_1(x)=h_1(\Phi (x))$. Each head is trained only on the units that received its treatment, while the representation is shared. The individual effect is $\tau (x)={\mu }_1(x)-{\mu }_0(x)$. The causal graph does not change: the network is a flexible structural equation for $Y$ given $T$ and $X$, and identification still rests on no unmeasured confounding and on overlap.

We write the two-headed network directly in Stan, with 8 hidden $\tanh$ units shared by both heads, normal priors scaled by $1/\sqrt{\text{fan-in}}$, and sample it with HMC: 226 weights and biases in all, plus the noise standard deviation. On two cores it takes about three minutes.

library(tidyverse)
library(brms)
library(posterior)
library(jsonlite)
library(rstan)
library(dbarts)
library(nnet)
options(mc.cores = parallel::detectCores())

harbour <- c(ink = "#10191d", muted = "#4a5f68", grid = "#d5e3e6",
             sky = "#1d5187", amber = "#805907", rose = "#ad3b28",
             moss = "#1d6b54", lilac = "#5b3a94")
theme_harbour <- function(base_size = 11) {
  theme_minimal(base_size = base_size) +
    theme(text = element_text(colour = harbour[["ink"]]),
          axis.text = element_text(colour = harbour[["muted"]]),
          panel.grid.major = element_line(colour = harbour[["grid"]]),
          panel.grid.minor = element_blank(),
          strip.text = element_text(colour = harbour[["ink"]], hjust = 0),
          legend.position = "top")
}
theme_set(theme_harbour())

ihdp <- read.csv("https://raw.githubusercontent.com/AMLab-Amsterdam/CEVAE/master/datasets/IHDP/csv/ihdp_npci_1.csv",
                 header = FALSE, col.names = c("treat", "y", "y_cf", "mu0", "mu1", paste0("x", 1:25)))
tau <- ihdp$mu1 - ihdp$mu0                                    # true individual effects (simulated, so known)
Xs <- scale(as.matrix(select(ihdp, starts_with("x"))))
ps <- fitted(glm(ihdp$treat ~ Xs, family = binomial))         # propensity scores, for checking overlap
c(children = nrow(ihdp), treated = sum(ihdp$treat), true_ATE = round(mean(tau), 3), true_effect_sd = round(sd(tau), 3),
  propensity_below_0.05 = sum(ps < 0.05))

             children               treated              true_ATE 
              747.000               139.000                 4.016 
       true_effect_sd propensity_below_0.05 
                0.860               110.000 

tarnet_code <- "
data { int<lower=1> N; int<lower=1> P; int<lower=1> H; matrix[N, P] X; vector[N] t; vector[N] y; }
parameters {
  matrix[P, H] W; vector[H] b;                    // shared representation layer
  vector[H] v0; vector[H] v1; real a0; real a1;   // two heads: outcome under control and under treatment
  real<lower=0> sigma;
}
model {
  matrix[N, H] Z = tanh(X * W + rep_matrix(b', N));
  to_vector(W) ~ normal(0, 1 / sqrt(P)); b ~ normal(0, 1);
  v0 ~ normal(0, 1 / sqrt(H)); v1 ~ normal(0, 1 / sqrt(H)); a0 ~ normal(0, 2); a1 ~ normal(0, 2);
  sigma ~ exponential(1);
  y ~ normal(t .* (Z * v1 + a1) + (1 - t) .* (Z * v0 + a0), sigma);   // each head sees only its own units
}
generated quantities {
  vector[N] tau;                                  // individual effect: difference of the two heads
  { matrix[N, H] Z = tanh(X * W + rep_matrix(b', N)); tau = Z * (v1 - v0) + a1 - a0; }
}"
path <- "fits/ch14_tarnet.rds"
if (file.exists(path)) fit_bnn <- readRDS(path) else {
  fit_bnn <- stan(model_code = tarnet_code, data = list(N = nrow(ihdp), P = ncol(Xs), H = 8, X = Xs, t = ihdp$treat, y = ihdp$y),
                  chains = 4, iter = 1000, seed = 14, refresh = 0)
  saveRDS(fit_bnn, path)
}
tau_bnn <- rstan::extract(fit_bnn, "tau")$tau                 # draws x children
rh <- summary(fit_bnn)$summary[, "Rhat"]
c(max_rhat_weights = round(max(rh[grep("^W\\[", names(rh))]), 2), max_rhat_effects = round(max(rh[grep("^tau\\[", names(rh))]), 3),
  divergences = get_num_divergent(fit_bnn))

max_rhat_weights max_rhat_effects      divergences 
           1.100            1.006            0.000 

The four chains agree on what matters: the largest $\widehat{R}$ among the 747 children's individual effects is 1.006, and HMC reported no divergent transitions. The largest $\widehat{R}$ among the weights is 1.10, which by the rules of Chapter 5 would be a failure. It is not one here: the chains have settled in different, equivalent arrangements of the hidden units, which give the same functions. Had we checked only the weights, we would have thrown away a good fit.

For comparison we fit a deep ensemble: ten networks with the same 8 hidden units, trained by maximum likelihood (with a small weight penalty) from different random starts using R's nnet package, with the treatment as an extra input. The spread of the ten networks' effects is the ensemble's uncertainty. We also refit BART as in Chapter 13.

set.seed(14)
Xt <- cbind(Xs, treat = ihdp$treat); Xt1 <- Xt; Xt1[, "treat"] <- 1; Xt0 <- Xt; Xt0[, "treat"] <- 0
ens <- sapply(1:10, function(k) { net <- nnet(Xt, ihdp$y, size = 8, linout = TRUE, decay = 0.05, maxit = 1000, trace = FALSE)
  predict(net, Xt1) - predict(net, Xt0) })                   # children x networks
fit_bart <- bart(Xt, ihdp$y, rbind(Xt1, Xt0), ndpost = 2000, nskip = 1000, verbose = FALSE)
n <- nrow(ihdp); tau_bart <- fit_bart$yhat.test[, 1:n] - fit_bart$yhat.test[, n + 1:n]
summ <- function(draws) {                                     # draws x children (or children x members, transposed)
  est <- colMeans(draws); lo <- apply(draws, 2, quantile, 0.05); hi <- apply(draws, 2, quantile, 0.95)
  c(ATE = mean(est), sqrt_PEHE = sqrt(mean((est - tau)^2)), coverage_90 = mean(tau > lo & tau < hi), mean_width = mean(hi - lo))
}
ens_draws <- t(ens)
ens_norm <- sapply(1:n, function(i) rnorm(4000, mean(ens[i, ]), sd(ens[i, ])))   # ensemble as a normal: mean ± spread
rbind(truth = c(mean(tau), 0, NA, NA), BNN_HMC = summ(tau_bnn), deep_ensemble = summ(ens_norm), BART = summ(tau_bart)) |> round(3)

                ATE sqrt_PEHE coverage_90 mean_width
truth         4.016     0.000          NA         NA
BNN_HMC       3.990     0.261       0.987      1.187
deep_ensemble 3.925     0.567       0.971      2.736
BART          4.020     0.387       0.965      1.151

The Bayesian network estimates the average effect at 3.99 (truth 4.02) and recovers the individual effects with $\sqrt{\text{PEHE}}=0.26$, against BART's 0.39 and the deep ensemble's 0.57. Its 90% intervals contain 98.7% of the true effects with a mean width of 1.19, about as wide as BART's (1.15, with 96.5% coverage). The ensemble's intervals are more than twice as wide (2.74) for 97.1% coverage: ten networks trained to maximum likelihood disagree a great deal about individual effects, a mixture of genuine uncertainty and the noise of optimisation from random starts. These are results for one replication of one benchmark with untuned settings, an anecdote in the sense of Chapter 13's Exercise 5. What they show is that a small Bayesian network, sampled properly, is a competitive and well-calibrated outcome model, not that networks beat trees.

tibble(ps = rep(ps, 3), method = rep(c("BNN (HMC)", "deep ensemble", "BART"), each = n),
       width = c(apply(tau_bnn, 2, quantile, 0.95) - apply(tau_bnn, 2, quantile, 0.05),
                 apply(ens_norm, 2, quantile, 0.95) - apply(ens_norm, 2, quantile, 0.05),
                 apply(tau_bart, 2, quantile, 0.95) - apply(tau_bart, 2, quantile, 0.05))) |>
  mutate(method = factor(method, c("BNN (HMC)", "deep ensemble", "BART"))) |>
  ggplot(aes(ps, width)) + geom_point(size = 0.7, alpha = 0.5, colour = harbour[["sky"]]) +
  geom_smooth(se = FALSE, method = "loess", formula = y ~ x, colour = harbour[["rose"]]) +
  facet_wrap(~ method) + scale_x_sqrt() + labs(x = "propensity score (square-root scale)", y = "90% interval width") +
  theme(panel.spacing = unit(1.5, "lines"))
w <- function(d) apply(d, 2, quantile, 0.95) - apply(d, 2, quantile, 0.05)
sapply(list(BNN = tau_bnn, ensemble = ens_norm, BART = tau_bart), function(d) round(c(low_overlap = mean(w(d)[ps < 0.02]), rest = mean(w(d)[ps >= 0.02])), 3))

              BNN ensemble  BART
low_overlap 1.407    3.093 1.272
rest        1.181    2.728 1.148

All three methods give wider intervals for the children with the lowest propensity scores (below 0.02), who have few treated counterparts: 1.41 against 1.18 for the rest with the Bayesian network, 1.27 against 1.15 with BART, and 3.09 against 2.73 with the ensemble. The widening is real but modest. A smooth model fitted to 139 treated children extrapolates smoothly into regions with few of them, and how much it admits not knowing there is set by its prior and its architecture, not by the data. A model's own uncertainty is no substitute for checking overlap directly, as in Chapter 13.

Lecture 31 — Transformers as Bayesian machines

14.3 Attention is kernel regression

The transformer (Vaswani et al., 2017) is the architecture behind every large language model. Its central operation, attention, is older and simpler than its reputation. Given a query $q$ and a set of items, each with a key $k_i$ and a value $v_i$, attention returns a weighted average of the values, with weights from a softmax of query–key similarities:

In a transformer, queries, keys and values are learned linear projections of the token representations, and $s$ is a scaled dot product. Statisticians will recognise the formula: with $s(q,k_i)=-\beta (q-k_i{)}^2$ it is exactly the Nadaraya–Watson kernel regression estimator of 1964, a locally weighted average with a Gaussian kernel of bandwidth $1/\sqrt{2\beta }$. And a dot product of suitably augmented features can represent this squared distance exactly, since $-\beta (q-k{)}^2=2\beta qk-\beta k^2-\beta q^2$, and the last term is the same for every item and cancels in the softmax. A single attention head is a kernel smoother with a learned kernel. A transformer stacks many such heads, interleaved with small feed-forward networks, so that each layer re-weights and recombines the items in context.

Kernel smoothers are close relatives of the Gaussian processes of Chapter 8: the GP posterior mean is also a weighted sum of the observed outcomes, with weights from the kernel. The difference is that the GP's weights account for how the observed points relate to each other (two nearby points share their weight rather than each taking a full vote), and they shrink toward the prior mean when no data are nearby. Plain attention does neither. A common addition helps with the second: an extra "sink" item with a fixed key and a value of zero, which absorbs the attention when nothing in the context is relevant, so that the output falls back toward zero. The widget compares the two on a function drawn from a GP prior.

14.4 Learning to do inference

The widget's attention head was not tuned by hand. Its two parameters came from fitting it, in Stan, to simulated tasks: 400 data sets, each with twelve noisy observations of a different function drawn from a Gaussian-process prior and one held-out query point whose value the model must predict from the twelve. The model sees nothing but the twelve context points of each task; it has no notion of a GP. Trained to predict well across many tasks drawn from the prior, the attention head learns an approximation to the posterior predictive distribution under that prior, and it can then be applied to a new data set instantly, without any fitting. This is the idea of prior-fitted networks (Müller et al., 2022): to make a network do Bayesian inference, train it on data simulated from the prior. TabPFN (Hollmann et al., 2025), a transformer trained on millions of simulated tabular data sets drawn from a prior over structural causal models, predicts on small real data sets in a single forward pass, competitively with tuned gradient-boosted trees. Ours is the same idea with one attention head and three parameters (bandwidth, sink and a value scale), small enough to fit by HMC so that we also get a posterior over the attention parameters themselves.

set.seed(14)
sim_tasks <- function(M, K = 12, ell = 0.15, noise = 0.1) {   # M tasks: functions from a GP prior on [0, 1]
  xs <- matrix(runif(M * (K + 1)), M); ys <- matrix(0, M, K + 1)
  for (m in 1:M) { S <- exp(-outer(xs[m, ], xs[m, ], "-")^2 / (2 * ell^2)) + diag(1e-8, K + 1)
    ys[m, ] <- t(chol(S)) %*% rnorm(K + 1) + rnorm(K + 1, 0, noise) }
  list(xc = xs[, 1:K], yc = ys[, 1:K], xq = xs[, K + 1], yq = ys[, K + 1])
}
attn_code <- "
data { int<lower=1> M; int<lower=1> K; matrix[M, K] xc; matrix[M, K] yc; vector[M] xq; vector[M] yq; }
parameters { real log_beta; real b0; real<lower=0> c; real<lower=0> sigma; }
model {
  log_beta ~ normal(3, 2); b0 ~ normal(0, 3); c ~ normal(1, 1); sigma ~ exponential(2);
  for (m in 1:M) {
    vector[K + 1] s; vector[K + 1] v;
    for (k in 1:K) { s[k] = -exp(log_beta) * square(xq[m] - xc[m, k]); v[k] = c * yc[m, k]; }
    s[K + 1] = b0; v[K + 1] = 0;                        // the sink: attended to when no context point is close
    yq[m] ~ normal(dot_product(softmax(s), v), sigma);
  }
}"
train_tasks <- sim_tasks(400); test_tasks <- sim_tasks(2000)
fit_attn <- stan(model_code = attn_code, data = c(list(M = 400, K = 12), train_tasks), chains = 4, iter = 1000, seed = 14, refresh = 0)
dr <- rstan::extract(fit_attn)
round(summary(fit_attn, pars = c("log_beta", "b0", "c", "sigma"))$summary[, c("mean", "2.5%", "97.5%", "Rhat")], 3)
attend <- function(lb, b0, scale, xc, yc, xq) { s <- c(-exp(lb) * (xq - xc)^2, b0); a <- exp(s - max(s)); sum(a[seq_along(xc)] * scale * yc) / sum(a) }
gp_mean <- function(xc, yc, xq, ell = 0.15, noise = 0.1) {
  S <- exp(-outer(xc, xc, "-")^2 / (2 * ell^2)) + diag(noise^2, length(xc)); sum(exp(-(xq - xc)^2 / (2 * ell^2)) * solve(S, yc)) }
use <- seq(1, length(dr$b0), by = 20)                          # 100 posterior draws
pred <- with(test_tasks, sapply(1:2000, function(m) {
  c(attention = mean(sapply(use, function(d) attend(dr$log_beta[d], dr$b0[d], dr$c[d], xc[m, ], yc[m, ], xq[m]))),
    no_sink = attend(mean(dr$log_beta), -Inf, 1, xc[m, ], yc[m, ], xq[m]),
    gp_bayes = gp_mean(xc[m, ], yc[m, ], xq[m]), prior_mean = 0) }))
rmse <- apply(pred, 1, function(p) sqrt(mean((p - test_tasks$yq)^2)))
round(c(rmse, cor_attention_gp = cor(pred["attention", ], pred["gp_bayes", ])), 3)

           mean    2.5%  97.5%  Rhat
log_beta  5.645   5.381  5.915 1.003
b0       -7.974 -10.758 -5.617 1.003
c         0.989   0.954  1.023 1.004
sigma     0.326   0.305  0.348 1.000
       attention          no_sink         gp_bayes       prior_mean 
           0.330            0.349            0.235            0.996 
cor_attention_gp 
           0.969 

The posterior is tight: $\log \beta =5.65$ (95% interval 5.38 to 5.92), so $\beta \approx 280$, a kernel bandwidth of about 0.04, much narrower than the GP's length-scale of 0.15: with little noise, the best smoother leans heavily on the nearest points. The sink score is $b_0=-8.0$, so the sink takes over when the nearest context point is farther than about $\sqrt{8/280}\approx 0.17$, roughly the GP's length-scale, which is where the GP posterior mean also starts to fall back to the prior mean of zero. The value scale is 0.99. On 2,000 new tasks the attention head predicts with an RMSE of 0.330, against 0.235 for the exact GP posterior mean, 0.349 for the same head without a sink, and 0.996 for always predicting zero; its predictions correlate 0.97 with the GP posterior mean. A three-parameter head has learned most of what Bayesian inference under this prior does. The remaining gap is what a kernel smoother cannot represent: the GP accounts for how the context points relate to each other, sharing weight within clusters, and it continues trends a little way beyond the data. Those are the things that more heads and more layers let a transformer learn.

One head is not a transformer. With PyTorch, the same idea can be tried at scale: a transformer with three layers of four attention heads (about 159,000 parameters) is trained by stochastic gradient descent on freshly simulated tasks from the same Gaussian-process prior, each with twelve context points and eight query points, to output a mean and a standard deviation for every query. Following Müller et al. (2022), every token may attend only to the context, so each query is predicted from the context alone. Training uses 6,000 batches of 256 tasks, about 1.5 million simulated data sets, and takes a few minutes on two CPU cores; the trained weights are cached. We then evaluate it on the same 2,000 test tasks as the Stan model.

run_python <- function(code) stopifnot(system2("python3", "-", input = strsplit(code, "\n")[[1]]) == 0)   # Python as a separate process
with(test_tasks, write.csv(data.frame(xc, yc, xq, yq), "fits/ch14_test_tasks.csv", row.names = FALSE))
run_python("
import torch, torch.nn as nn, numpy as np, os
torch.manual_seed(14); torch.set_num_threads(2)
K, Q, ELL, NOISE = 12, 8, 0.15, 0.1
def tasks(B):                                               # B tasks from the same GP prior as the Stan model
    x = torch.rand(B, K + Q, 1)
    S = torch.exp(-(x - x.transpose(1, 2))**2 / (2 * ELL**2)) + 1e-6 * torch.eye(K + Q)
    y = torch.linalg.cholesky(S) @ torch.randn(B, K + Q, 1) + NOISE * torch.randn(B, K + Q, 1)
    return x[:, :K], y[:, :K], x[:, K:], y[:, K:]
class PFN(nn.Module):
    def __init__(s, d=64, layers=3, heads=4):
        super().__init__()
        s.ctx = nn.Sequential(nn.Linear(2, d), nn.GELU(), nn.Linear(d, d))   # embeds a context pair (x, y)
        s.qry = nn.Sequential(nn.Linear(1, d), nn.GELU(), nn.Linear(d, d))   # embeds a query x
        layer = nn.TransformerEncoderLayer(d, heads, 4 * d, dropout=0.0, batch_first=True, norm_first=True)
        s.enc = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False); s.out = nn.Linear(d, 2)
    def forward(s, xc, yc, xq):
        k, q = xc.shape[1], xq.shape[1]
        h = torch.cat([s.ctx(torch.cat([xc, yc], -1)), s.qry(xq)], 1)
        mask = torch.ones(k + q, k + q, dtype=torch.bool); mask[:, :k] = False   # attend to context only
        o = s.out(s.enc(h, mask=mask)[:, k:])
        return o[..., 0], nn.functional.softplus(o[..., 1]) + 1e-3          # predictive mean and sd
net = PFN()
if os.path.exists('fits/ch14_pfn.pt'):
    net.load_state_dict(torch.load('fits/ch14_pfn.pt'))
else:                                                      # minimise the Gaussian log loss over prior draws
    opt = torch.optim.AdamW(net.parameters(), lr=1e-3, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=1e-3, total_steps=6000)
    for step in range(6000):
        xc, yc, xq, yq = tasks(256); mu, sd = net(xc, yc, xq)
        loss = (torch.log(sd) + 0.5 * ((yq[..., 0] - mu) / sd)**2).mean()
        opt.zero_grad(); loss.backward(); opt.step(); sched.step()
    torch.save(net.state_dict(), 'fits/ch14_pfn.pt')
net.eval()
T = torch.tensor(np.loadtxt('fits/ch14_test_tasks.csv', delimiter=',', skiprows=1), dtype=torch.float32)   # the R test tasks
with torch.no_grad():
    mu, sd = net(T[:, :K, None], T[:, K:2 * K, None], T[:, 2 * K].reshape(-1, 1, 1))
np.savetxt('fits/ch14_pfn_pred.csv', np.column_stack([mu[:, 0], sd[:, 0]]), delimiter=',')
")
pfn <- read.csv("fits/ch14_pfn_pred.csv", header = FALSE, col.names = c("mean", "sd"))
gp_sd <- function(xc, xq, ell = 0.15, noise = 0.1) {          # exact GP posterior predictive sd
  S <- exp(-outer(xc, xc, "-")^2 / (2 * ell^2)) + diag(noise^2, length(xc)); k <- exp(-(xq - xc)^2 / (2 * ell^2))
  sqrt(1 - sum(k * solve(S, k)) + noise^2) }
gsd <- with(test_tasks, sapply(1:2000, function(m) gp_sd(xc[m, ], xq[m])))
y <- test_tasks$yq; cover <- function(m, s) mean(abs(y - m) < qnorm(0.95) * s)
c(rmse_transformer = sqrt(mean((pfn$mean - y)^2)), rmse_one_head = rmse[["attention"]],
  rmse_gp_bayes = rmse[["gp_bayes"]], cor_transformer_gp = cor(pfn$mean, pred["gp_bayes", ]),
  cover90_transformer = cover(pfn$mean, pfn$sd), cover90_gp = cover(pred["gp_bayes", ], gsd),
  mean_sd_transformer = mean(pfn$sd), mean_sd_gp = mean(gsd)) |> round(3)

   rmse_transformer       rmse_one_head       rmse_gp_bayes  cor_transformer_gp 
              0.243               0.330               0.235               0.998 
cover90_transformer          cover90_gp mean_sd_transformer          mean_sd_gp 
              0.901               0.895               0.197               0.191 

The transformer predicts the 2,000 test tasks with an RMSE of 0.243, against 0.330 for the one-head Stan model and 0.235 for the exact Gaussian-process posterior mean, and its predictions correlate 0.998 with that posterior mean. It has learned almost exactly what Bayesian inference under this prior does, including the parts a single kernel smoother could not represent. It has also learned the uncertainty: its 90% predictive intervals contain 90.1% of the held-out values, against 89.5% for the exact posterior, and its intervals are on average only slightly wider (standard deviation 0.197 against 0.191). Nothing in the network's design mentions Gaussian processes, kernels or Bayes' rule; it learned the posterior predictive distribution from 1.5 million examples of the prior. That is the whole idea of prior-fitted networks, and the reason they are only as good as the prior they were trained on.

Why should training on simulated tasks produce Bayesian inference? Because the expected squared error over tasks drawn from the prior is minimised by the posterior mean under that prior: the best possible prediction of $y_q$ from the context is $𝔼[y_q\mid \text{context}]$, averaged over the prior's functions. A model trained to minimise that error, if flexible enough, converges to it. The same argument with the log loss gives the whole posterior predictive distribution. Xie et al. (2022) used a related argument to explain in-context learning in large language models, their ability to pick up a task from a few examples in the prompt: if the training text is a mixture of documents generated by many latent "concepts", next-word prediction forces the model to infer the concept from the context, which is implicit Bayesian inference. On this view a language model's pretraining data is its prior, a perspective that will matter in Section 14.8.

TabPFN can now be put to the test of Section 14.2. It was trained on simulated tables, not on any medical data, and is applied to the IHDP benchmark as an "S-learner": one regression of the outcome on the covariates and the treatment, from which each child's effect is the difference between the predictions with the treatment switched on and off. There is no prior to choose, no sampler to tune and no training on IHDP beyond conditioning on its 747 rows. We use the openly released TabPFN version 2 weights (Built with PriorLabs-TabPFN; Prior Labs License 1.1).

path <- "fits/ch14_tabpfn.rds"
if (file.exists(path)) tau_tabpfn <- readRDS(path) else {
  write.csv(data.frame(Xs, treat = ihdp$treat, y = ihdp$y), "fits/ch14_ihdp.csv", row.names = FALSE)
  run_python("
from tabpfn import TabPFNRegressor                          # pretrained prior-fitted network for tables
import numpy as np, torch
torch.set_num_threads(2)
D = np.loadtxt('fits/ch14_ihdp.csv', delimiter=',', skiprows=1); X, y = D[:, :-1], D[:, -1]   # covariates and treatment; outcome
X1, X0 = X.copy(), X.copy(); X1[:, -1] = 1; X0[:, -1] = 0
m = TabPFNRegressor(device='cpu', random_state=14).fit(X, y)  # 'fitting' is conditioning: one forward pass
np.savetxt('fits/ch14_tabpfn.csv', m.predict(X1) - m.predict(X0), delimiter=',')
")
  tau_tabpfn <- scan("fits/ch14_tabpfn.csv", quiet = TRUE); saveRDS(tau_tabpfn, path)
}
rbind(BNN_HMC = summ(tau_bnn)[1:2], BART = summ(tau_bart)[1:2], deep_ensemble = summ(ens_norm)[1:2],
      TabPFN = c(ATE = mean(tau_tabpfn), sqrt_PEHE = sqrt(mean((tau_tabpfn - tau)^2)))) |> round(3)

                ATE sqrt_PEHE
BNN_HMC       3.990     0.261
BART          4.020     0.387
deep_ensemble 3.925     0.567
TabPFN        3.917     0.297

TabPFN estimates the individual effects with $\sqrt{\text{PEHE}}=0.297$, better than BART (0.387) and the deep ensemble (0.567) and close to the Bayesian network fitted by HMC (0.261), with an average effect of 3.92 against the true 4.02. A network that has never seen IHDP, given only its 747 rows, does nearly as well as a model built and sampled for it. Two cautions keep this in proportion. TabPFN's predictive distribution describes outcomes, not effects, so unlike the HMC network and BART it gives no interval for an individual child's effect without further work. And its prior over structural causal models, not anything in the data, is what makes the S-learner work here; on a problem outside that prior, such as one with strong confounding by an unrecorded variable, it has no way to know (Section 14.7).

Large models also need to say when they do not know. Language models produce fluent answers whether or not the answer is grounded, and their token-level probabilities are poor guides to factual accuracy. Farquhar et al. (2024) proposed semantic entropy: sample several answers to the same question, cluster them by meaning (two answers are equivalent if each entails the other), and compute the entropy over meaning clusters. High semantic entropy flagged confabulations far better than naive measures. It is the posterior predictive check of this course applied to a language model: draw from the predictive distribution and look at its spread, in the space that matters.

Lecture 32 — Learning causes, not shortcuts

14.5 Shortcuts are causal failures

A classifier trained to minimise error uses whatever predicts the label in its training data. Some predictive features are causes or reliable consequences of the label; others are accidents of how the data were collected. Geirhos et al. (2020) called the habit of relying on the second kind shortcut learning, and collected examples: image classifiers that recognise cows by the green pasture behind them, pneumonia detectors that learned which hospital's scanner took the X-ray, and language models that answer reading-comprehension questions from superficial word overlap. Chapter 12 added one of our own: in the finance answers of HC3 (Guo et al., 2023), a formatting artefact separated ChatGPT from people almost perfectly.

In causal terms, a shortcut is a feature connected to the label by a path that runs through the data-collection process, a path that need not exist, or may point the other way, in new data. The causal graph of Chapter 12 already said this. What it adds is a remedy: if the mechanism by which the true causes produce the label is the same everywhere, while the shortcuts vary from one data source to another, then features whose relation to the label is stable across sources are more likely to be causal, and a model built on them is more likely to keep working on a source it has never seen. This is the principle of invariance.

HC3 lets us test it. Besides finance and open-domain questions, it has medical questions (1,246 with both kinds of answer) and computer-science questions from Wikipedia (842). We treat the four subsets as four environments: different questions, different human sources and different collection details, but the same kind of author on each side.

read_hc3 <- function(subset) {
  local <- file.path("data", "hc3", paste0(subset, ".jsonl"))      # download once, then reuse the local copy
  if (!file.exists(local)) { dir.create(dirname(local), recursive = TRUE, showWarnings = FALSE)
    download.file(sprintf("https://huggingface.co/datasets/Hello-SimpleAI/HC3/resolve/main/%s.jsonl", subset), local, quiet = TRUE) }
  raw <- lapply(readLines(local, encoding = "UTF-8"), fromJSON, simplifyVector = FALSE)
  first <- function(x) if (length(x) > 0) x[[1]] else NA_character_
  tibble(q = seq_along(raw), human = map_chr(raw, ~ first(.x$human_answers)), chatgpt = map_chr(raw, ~ first(.x$chatgpt_answers))) |>
    pivot_longer(-q, names_to = "author", values_to = "text") |>
    group_by(q) |> filter(all(!is.na(text) & nchar(text) > 0)) |> ungroup() |> mutate(ai = as.integer(author == "chatgpt"))
}
envs <- c("finance", "medicine", "open_qa", "wiki_csai")
hc3 <- map_dfr(set_names(envs), read_hc3, .id = "env")
hc3 |> mutate(seam = str_detect(text, "[a-z]\\.[A-Z]"), words = str_count(text, "\\S+")) |>
  group_by(env, author) |> summarise(answers = n(), seam = round(mean(seam), 3), median_words = median(words), .groups = "drop")

# A tibble: 8 x 5
  env       author  answers  seam median_words
  <chr>     <chr>     <int> <dbl>        <dbl>
1 finance   chatgpt    3933 0.902         206 
2 finance   human      3933 0.005         129 
3 medicine  chatgpt    1246 0.002         185 
4 medicine  human      1246 0.52           74 
5 open_qa   chatgpt    1187 0.002         104 
6 open_qa   human      1187 0.071          27 
7 wiki_csai chatgpt     842 0.021         188.
8 wiki_csai human       842 0.004         166 

The formatting artefact behaves differently in every environment. In finance, 90% of ChatGPT's answers contain the glued full stop against 0.5% of the human ones (the shortcut of Chapter 12); in medicine it reverses, appearing in 52% of human answers and 0.2% of ChatGPT's; in the other two environments it is rare on both sides. Length is less extreme but also unstable: ChatGPT writes longer answers everywhere, but the gap in median length ranges from 22 words in Wikipedia computer science (188 against 166) to 111 words in medicine (185 against 74), and the human open-domain answers are very short (27 words). A detector that leans on either feature has learned something about how each subset was collected, not about how ChatGPT writes.

14.6 Invariance across environments

Peters, Bühlmann and Meinshausen (2016) turned the invariance principle into a method, invariant causal prediction: if a set of predictors contains the causes of the outcome, the conditional distribution of the outcome given those predictors is the same in every environment, and sets that fail this test can be rejected. Arjovsky et al. (2019) proposed a version for deep networks, invariant risk minimisation, which looks for a representation on which the optimal classifier is the same across environments. Later work found that invariant risk minimisation can fail in simple settings, especially with few environments (Rosenfeld, Ravikumar and Risteski, 2021). The principle is sound; turning it into a reliable algorithm is hard. The widget shows the principle in its simplest form.

When the spurious association is the same in both training environments, nothing in the training data distinguishes it from a cause, and the pooled classifier uses it, with a large gain in accuracy if the test environment resembles the training ones and a large loss if it reverses. Only when the training environments disagree does the spurious feature reveal itself. Invariance needs variation: environments that differ in the shortcuts and agree in the mechanism.

We apply a simple, transparent version to HC3. The training environments are finance and medicine; open-domain and Wikipedia computer-science answers are held out entirely as test environments. The candidate features are the 300 most widespread words (as log counts, standardised) plus two non-lexical features, the glued full stop ("seam") and the log length of the answer. For each feature and each training environment we fit a one-feature Bayesian logistic regression (normal(0, 1) prior, Laplace approximation as in Chapter 9). We call a feature stable if its association has the same sign in both environments, is clearly non-zero in both (posterior mean more than two standard deviations from zero) and differs in size by less than a factor of two. Then we compare four detectors, each a Bayesian logistic regression on standardised features, trained on finance and medicine together: all features; all features except the two non-lexical artefacts; the stable features only; and, for reference, all features trained on finance alone.

tokenize <- function(x) str_split(str_to_lower(x), "[^a-z']+")
is_train <- hc3$env %in% c("finance", "medicine"); tok <- tokenize(hc3$text)
dfreq <- table(unlist(map(tok[is_train], unique))); dfreq <- dfreq[names(dfreq) != ""]
vocab <- names(sort(dfreq, decreasing = TRUE))[1:300]          # chosen from training environments, without labels
Fx <- cbind(log1p(t(sapply(tok, function(w) tabulate(match(w, vocab), 300)))),
            seam = as.numeric(str_detect(hc3$text, "[a-z]\\.[A-Z]")), loglen = log1p(str_count(hc3$text, "\\S+")))
colnames(Fx)[1:300] <- paste0("w_", vocab)
sdv <- apply(Fx[is_train, ], 2, sd); sdv[sdv == 0] <- 1
Fx <- scale(Fx, colMeans(Fx[is_train, ]), sdv); y <- hc3$ai
laplace_logit <- function(X, y, prior_sd = 1) {                 # MAP and Gaussian (Laplace) posterior
  X1 <- cbind(1, X); pv <- c(2.5, rep(prior_sd, ncol(X)))^2
  nlp <- function(b) { e <- X1 %*% b; -sum(y * e - log1p(exp(e))) + sum(b^2 / (2 * pv)) }
  gr <- function(b) -as.vector(t(X1) %*% (y - plogis(X1 %*% b))) + b / pv
  b <- optim(rep(0, ncol(X1)), nlp, gr, method = "L-BFGS-B", control = list(maxit = 2000))$par
  p <- as.vector(plogis(X1 %*% b))
  list(mean = setNames(b, c("(Intercept)", colnames(X))), cov = solve(crossprod(X1 * sqrt(p * (1 - p))) + diag(1 / pv, ncol(X1))))
}
predict_lap <- function(fit, X) {                               # predictive probability, probit approximation
  X1 <- cbind(1, X[, names(fit$mean)[-1], drop = FALSE]); v <- rowSums((X1 %*% fit$cov) * X1)
  as.vector(plogis((X1 %*% fit$mean) / sqrt(1 + pi * v / 8)))
}
one_feature <- function(e) t(sapply(colnames(Fx), function(f) { k <- hc3$env == e; fit <- laplace_logit(Fx[k, f, drop = FALSE], y[k])
  c(b = fit$mean[[2]], s = sqrt(fit$cov[2, 2])) }))
u_fin <- one_feature("finance"); u_med <- one_feature("medicine")
stable <- sign(u_fin[, "b"]) == sign(u_med[, "b"]) & abs(u_fin[, "b"] / u_fin[, "s"]) > 2 & abs(u_med[, "b"] / u_med[, "s"]) > 2 &
  pmin(abs(u_fin[, "b"]), abs(u_med[, "b"])) / pmax(abs(u_fin[, "b"]), abs(u_med[, "b"])) > 0.5
words_only <- grep("^w_", colnames(Fx), value = TRUE)
detectors <- list(all_features = laplace_logit(Fx[is_train, ], y[is_train]),
                  no_artefacts = laplace_logit(Fx[is_train, words_only], y[is_train]),
                  stable_only = laplace_logit(Fx[is_train, stable], y[is_train]),
                  finance_only = laplace_logit(Fx[hc3$env == "finance", ], y[hc3$env == "finance"]))
score_env <- function(fit) map_dfr(set_names(envs), function(e) { k <- hc3$env == e; p <- pmin(pmax(predict_lap(fit, Fx[k, ]), 1e-12), 1 - 1e-12)
  tibble(accuracy = mean((p > 0.5) == y[k]), log_loss = -mean(y[k] * log(p) + (1 - y[k]) * log(1 - p))) }, .id = "env")
inv_results <- map_dfr(detectors, score_env, .id = "detector")
c(stable_features = sum(stable), seam_stable = stable[["seam"]], length_stable = stable[["loglen"]])
inv_results |> select(-log_loss) |> pivot_wider(names_from = env, values_from = accuracy) |> mutate(across(where(is.numeric), ~ round(.x, 3)))
inv_results |> select(-accuracy) |> pivot_wider(names_from = env, values_from = log_loss) |> mutate(across(where(is.numeric), ~ round(.x, 3)))

stable_features     seam_stable   length_stable 
             72               0               0 
# A tibble: 4 x 5
  detector     finance medicine open_qa wiki_csai
  <chr>          <dbl>    <dbl>   <dbl>     <dbl>
1 all_features   0.99     0.99    0.696     0.726
2 no_artefacts   0.982    0.995   0.769     0.708
3 stable_only    0.949    0.978   0.739     0.711
4 finance_only   0.996    0.833   0.663     0.72 
# A tibble: 4 x 5
  detector     finance medicine open_qa wiki_csai
  <chr>          <dbl>    <dbl>   <dbl>     <dbl>
1 all_features   0.041    0.035   0.828     0.797
2 no_artefacts   0.06     0.028   0.598     0.97 
3 stable_only    0.133    0.077   0.622     0.753
4 finance_only   0.04     0.412   0.929     0.643

Seventy-two features pass the stability screen; neither the seam nor the length does. On the environments they were trained on, all four detectors are nearly perfect (95% to 100% accuracy), apart from the finance-only detector on medicine (83%), which the reversed artefact misleads. On the two held-out environments, every detector falls to between 66% and 77%. Removing the two artefacts helps most on open-domain answers (77% against 70% with all features, and a log loss of 0.60 against 0.83); the stability screen gives 74%. On Wikipedia computer science the detectors are within two points of each other (71% to 73%), and the artefact-free detector's log loss worsens to 0.97, overconfident on text unlike its training data. Pooling two environments helps on open-domain answers, where the finance-only detector gets 66%, but not clearly on Wikipedia computer science, where it gets 72%. The honest summary is that invariance thinking found and removed the shortcut, which was the largest gain; the formal screen did no better than removing the known artefacts by hand; and nothing came close to in-domain accuracy. Two training environments are few, and features stable across finance and medicine need not be stable everywhere.

coefs <- tibble(feature = rownames(u_fin), finance = u_fin[, "b"], medicine = u_med[, "b"], stable = stable) |>
  mutate(label = if_else(feature %in% c("seam", "loglen", "w_important", "w_i", "w_was", "w_treatment", "w_help"),
                         recode(str_remove(feature, "^w_"), loglen = "log length"), NA_character_))
ggplot(coefs, aes(finance, medicine)) + geom_abline(colour = harbour[["muted"]], linetype = 2) + geom_hline(yintercept = 0, colour = harbour[["grid"]]) +
  geom_vline(xintercept = 0, colour = harbour[["grid"]]) +
  geom_point(aes(colour = stable), size = 1.3, alpha = 0.7) +
  geom_text(aes(label = label), size = 3, nudge_y = 0.18, colour = harbour[["ink"]], na.rm = TRUE) +
  scale_colour_manual(values = c(`TRUE` = harbour[["moss"]], `FALSE` = harbour[["muted"]]), labels = c("unstable", "stable")) +
  labs(x = "association in finance answers", y = "association in medicine answers", colour = NULL)

Most words lie near the diagonal: the first-person "i" marks human writing in both environments, and "important" marks ChatGPT in both, consistently with Chapter 12. The seam sits far off the diagonal in the lower right, strongly ChatGPT in finance and strongly human in medicine, and the log length, while positive in both, is four times stronger in medicine (2.90 against 0.71). Their instability is what disqualifies them, and it is visible only because there are two environments to compare.

Two further tools make text models more causal. Counterfactual data augmentation (Kaushik, Hovy and Lipton, 2020) asks people to make minimal edits that flip a document's label and adds the edited documents to the training data. A model trained on such pairs cannot rely on features that the edits left unchanged, because those features no longer predict the label. For AI-text detection, the analogue is paraphrasing machine text into a human register, and human text into a machine register. When text is a confounder, treatment or outcome in a causal analysis, as in "does the tone of a complaint affect how fast it is answered?", Veitch, Sridhar and Blei (2020) and Egami et al. (2022) show how to adjust for text with learned representations, and how not to: fitting the representation and estimating the effect on the same documents invites overfitting that looks like an effect. The common thread is causal representation learning (Schölkopf et al., 2021): finding variables in raw data, such as pixels or words, that behave like the variables of a causal graph, with stable mechanisms between them.

Lecture 33 — Causal AI: what can be learned, and what must be assumed

14.7 Learning causal inference from simulated worlds

Section 14.4 trained a model on simulated data sets to make it do Bayesian prediction. The same trick can be pointed at causal questions: simulate many worlds from a prior over structural causal models, record each world's data and its true causal effect, and train a network to map data to effect. Applied to a real data set, the network returns an approximation to the posterior of the effect under that prior over worlds, in a single forward pass. This is amortised causal inference: the cost of inference is paid once, in training, and then spread ("amortised") across every future data set. Do-PFN (Robertson et al., 2025) applies it with a transformer pretrained on synthetic causal models, including their interventions, and reports accurate effect estimates from observational data on synthetic and semi-synthetic benchmarks without being told the causal graph.

It is worth asking what such a network can possibly learn, and a small example makes the answer clear. Each simulated world is the confounded triangle of Chapter 13: a covariate $C$ affects treatment $T$ with coefficient $a$ and outcome $Y$ with coefficient $b$, and $T$ affects $Y$ with the causal effect $\tau$. The prior draws $a,b\sim \text{Normal}(0,1)$, $\tau \sim \text{Normal}(0,1)$ and noise standard deviations between 0.5 and 1.5, and each world has 200 observations. We summarise each world's data by its sample covariances and train small networks (20 hidden units, nnet) on 20,000 worlds to predict $\tau$. One network sees all six covariances of $(C,T,Y)$; the other sees only the three covariances of $(T,Y)$, as if $C$ were unmeasured. Both are compared, on 2,000 fresh worlds from the same prior, with the regressions of Chapter 13: $Y$ on $T$ and $C$ (the back-door adjustment), and $Y$ on $T$ alone (naive).

set.seed(1414)
sim_scm <- function(R, n = 200, conf_sd = 1) {          # R worlds from a prior over linear SCMs: C -> T, C -> Y, T -> Y
  map_dfr(1:R, function(r) {
    a <- rnorm(1, 0, conf_sd); b <- rnorm(1, 0, conf_sd); tau <- rnorm(1)
    C <- rnorm(n); Tr <- a * C + rnorm(n, 0, runif(1, 0.5, 1.5)); Y <- tau * Tr + b * C + rnorm(n, 0, runif(1, 0.5, 1.5))
    S <- cov(cbind(C, Tr, Y))
    tibble(vC = S[1, 1], cCT = S[1, 2], cCY = S[1, 3], vT = S[2, 2], cTY = S[2, 3], vY = S[3, 3], tau = tau,
           adjusted = (S[1, 1] * S[2, 3] - S[1, 2] * S[1, 3]) / (S[1, 1] * S[2, 2] - S[1, 2]^2),   # Y ~ T + C
           naive = S[2, 3] / S[2, 2])                                                             # Y ~ T
  })
}
train_worlds <- sim_scm(20000); test_worlds <- sim_scm(2000)
inp_obs <- c("vC", "cCT", "cCY", "vT", "cTY", "vY"); inp_hid <- c("vT", "cTY", "vY")
fit_net <- function(d, inp) nnet(as.matrix(d[, inp]), d$tau, size = 20, linout = TRUE, decay = 1e-3, maxit = 2000, trace = FALSE)
net_obs <- fit_net(train_worlds, inp_obs); net_hid <- fit_net(train_worlds, inp_hid)
rmse <- function(est, truth) sqrt(mean((est - truth)^2))
score_worlds <- function(d) c(net_C_observed = rmse(predict(net_obs, as.matrix(d[, inp_obs])), d$tau),
  adjusted_regression = rmse(d$adjusted, d$tau), net_C_hidden = rmse(predict(net_hid, as.matrix(d[, inp_hid])), d$tau),
  naive_regression = rmse(d$naive, d$tau), prior_mean = rmse(0, d$tau))
round(score_worlds(test_worlds), 3)

     net_C_observed adjusted_regression        net_C_hidden    naive_regression 
              0.100               0.085               0.399               0.466 
         prior_mean 
              0.988 

With the confounder observed, the network's RMSE is 0.100, close to the 0.085 of the adjusted regression, which is the correct estimator for these worlds. With the confounder hidden, the network's error is 0.399: a little better than the naive regression (0.466), and much better than the prior's 0.988, but nearly five times the error of adjustment.

The prior matters as much here as anywhere in the course. The network was trained on worlds in which confounding coefficients have standard deviation 1. The figure tests both networks on worlds with weaker and stronger confounding than they were trained on.

conf_grid <- c(0.25, 0.5, 1, 1.5, 2, 3)
grid_res <- map_dfr(conf_grid, function(s) { d <- sim_scm(500, conf_sd = s); as_tibble(as.list(score_worlds(d))) |> mutate(conf_sd = s) }) |>
  select(-prior_mean) |> pivot_longer(-conf_sd, names_to = "estimator", values_to = "rmse") |>
  mutate(estimator = factor(estimator, levels = c("adjusted_regression", "net_C_observed", "net_C_hidden", "naive_regression"),
                            labels = c("regression, C adjusted", "network, C observed", "network, C hidden", "regression, naive")))
ggplot(grid_res, aes(conf_sd, rmse, colour = estimator)) + geom_vline(xintercept = 1, linetype = 3, colour = harbour[["muted"]]) +
  geom_line(linewidth = 0.9) + geom_point(size = 2) +
  scale_colour_manual(values = unname(harbour[c("moss", "sky", "amber", "rose")])) +
  labs(x = "confounding strength in test worlds (sd of a and b)", y = "RMSE of effect estimate", colour = NULL)
grid_res |> pivot_wider(names_from = estimator, values_from = rmse) |> mutate(across(-conf_sd, ~ round(.x, 3))) |> as.data.frame()

  conf_sd network, C observed regression, C adjusted network, C hidden
1    0.25               0.086                  0.084             0.161
2    0.50               0.091                  0.084             0.239
3    1.00               0.094                  0.082             0.389
4    1.50               0.151                  0.086             0.524
5    2.00               0.299                  0.082             0.656
6    3.00               0.669                  0.082             0.858
  regression, naive
1             0.116
2             0.230
3             0.444
4             0.674
5             0.855
6             1.106

Within the range its prior covers well (standard deviations up to 1), the network with $C$ observed is within about 0.01 of the adjusted regression. Beyond it the network degrades: its error is 0.15 at standard deviation 1.5, 0.30 at 2 and 0.67 at 3, while the regression stays at 0.08 throughout. The network with $C$ hidden beats the naive regression when confounding is moderate or strong (standard deviation 1 or more), because it has learned to expect confounding and partly corrects for it on average; but when confounding is weak it is worse than the naive regression (0.16 against 0.12 at 0.25, 0.24 against 0.23 at 0.5), because it corrects for confounding that is not there. Both networks are approximating the posterior under their training prior, and both inherit that prior's assumptions.

Three lessons follow, and each is a lesson of Chapter 13 in a new costume. First, with the confounder observed, the network learned to adjust: it approximates the back-door estimate without ever having been told the back-door criterion, because the simulated worlds contained the information that adjustment is the right thing to do. Second, with the confounder hidden, no amount of training helps: the network does a little better than the naive regression, because it has learned how much confounding to expect on average under the prior, but the effect is not identified, and the network cannot extract information the data do not contain. Third, the network is only as good as its prior over worlds: outside the range of confounding it was trained on, it degrades, while the regression, which encodes the identification argument directly, does not. In a prototype, a network trained on worlds with weak confounding (standard deviation 0.2) and tested on strong confounding (standard deviation 2) had an RMSE of 0.68, nearly as bad as the naive regression (0.83) and eight times the adjusted regression's (0.08): it had never seen confounding that strong, so it had never learned to remove it. Amortisation buys speed and flexibility, not identification.

14.8 Can language models climb the ladder?

Large language models answer causal questions fluently. Asked whether altitude causes temperature or temperature causes altitude, they answer correctly; asked what would have happened had a patient not taken a drug, they produce plausible counterfactuals. Kıcıman et al. (2023) found GPT-4 answering pairwise "which causes which" questions with 97% accuracy and counterfactual-reasoning questions with 92%, well above earlier algorithms, and argued that language models carry a large store of causal knowledge extracted from text: the kind of domain knowledge that analysts have always used to draw DAGs.

Knowledge is not reasoning. The pairwise questions concern variables whose causal relation is written about at length, so a model can answer them from association with the text describing them. Zečević et al. (2023) called such models causal parrots: they repeat causal facts that appear in their training data, which is association at the level of text. CLadder (Jin et al., 2023) tests the distinction directly with more than 10,000 questions generated from formal causal graphs and queries on all three rungs, including stories whose causal relations are deliberately anti-commonsensical or nonsensical, so that memorised facts cannot help and only the stated graph and probabilities can. GPT-4 answered about 62% correctly, where 50% is chance; a prompting strategy that walks the model through the formal steps (extracting the graph, choosing the estimand, applying the do-calculus, then computing) raised it to 70%.

Since CLadder is public, we can run a small version of its test ourselves. We take 450 questions, 50 from each rung in each of the three kinds of story (commonsensical, anti-commonsensical and nonsensical), balanced between "yes" and "no" answers, and put them to two small open language models, Qwen2.5-Instruct with 0.5 and 1.5 billion parameters (Qwen Team, 2024), small enough to run on two CPU cores. Each model reads the question and we compare the probability it gives to answering "Yes" with the probability of "No". These models are about a thousand times smaller than GPT-4, so the test is a demonstration of method, not a measure of the state of the art.

path <- "fits/ch14_cladder.csv"
if (!file.exists(path) || length(unique(read.csv(path)$model)) < 2) run_python("
import torch, pandas as pd
from transformers import AutoTokenizer, AutoModelForCausalLM
torch.set_num_threads(2); rows = []
qs = pd.concat([pd.read_csv(f'data/cladder/test-{kind}-v1.5.csv').assign(story=kind)
                  .groupby(['rung', 'label'], group_keys=False).sample(25, random_state=14)
                for kind in ['commonsense', 'anticommonsense', 'noncommonsense']])
for name in ['Qwen/Qwen2.5-0.5B-Instruct', 'Qwen/Qwen2.5-1.5B-Instruct']:
    tok = AutoTokenizer.from_pretrained(name); lm = AutoModelForCausalLM.from_pretrained(name, dtype=torch.bfloat16).eval()   # half precision: fits in 8 GB
    yes = [tok.encode(w, add_special_tokens=False)[0] for w in ['Yes', 'yes', ' Yes']]
    no = [tok.encode(w, add_special_tokens=False)[0] for w in ['No', 'no', ' No']]
    for q in qs.itertuples():
        chat = [{'role': 'user', 'content': q.prompt + '\\nAnswer with Yes or No only.'}]
        ids = tok(tok.apply_chat_template(chat, add_generation_prompt=True, tokenize=False), return_tensors='pt').input_ids
        with torch.no_grad(): lp = torch.log_softmax(lm(ids, logits_to_keep=1).logits[0, -1], -1)
        score = float(torch.logsumexp(lp[yes], 0) - torch.logsumexp(lp[no], 0))   # log odds of answering yes
        rows.append(dict(model=name.split('/')[1], story=q.story, rung=q.rung, label=q.label, score=score))
    pd.DataFrame(rows).to_csv('fits/ch14_cladder.csv', index=False)          # saved after each model
    del lm; import gc; gc.collect()
")
cl <- read_csv(path, show_col_types = FALSE) |>
  mutate(story = factor(story, c("commonsense", "anticommonsense", "noncommonsense"), c("commonsensical", "anti-commonsensical", "nonsensical")),
         model = str_remove(model, "-Instruct"), says_yes = score > 0, correct = says_yes == (label == "yes")) |>
  group_by(model) |> mutate(correct_debiased = (score > median(score)) == (label == "yes")) |> ungroup()
cl |> group_by(model) |> summarise(questions = n(), accuracy = mean(correct), says_yes = mean(says_yes),
                                   accuracy_after_removing_yes_bias = mean(correct_debiased)) |> mutate(across(where(is.double), ~ round(.x, 3)))
cl |> group_by(model, story) |> summarise(accuracy = mean(correct), debiased = mean(correct_debiased), .groups = "drop") |>
  pivot_wider(names_from = model, values_from = c(accuracy, debiased)) |> mutate(across(where(is.double), ~ round(.x, 3))) |> as.data.frame()

# A tibble: 2 x 5
  model        questions accuracy says_yes accuracy_after_removing_yes_bias
  <chr>            <int>    <dbl>    <dbl>                            <dbl>
1 Qwen2.5-0.5B       450    0.498    0.78                             0.527
2 Qwen2.5-1.5B       450    0.471    0.602                            0.484
                story accuracy_Qwen2.5-0.5B accuracy_Qwen2.5-1.5B
1      commonsensical                 0.513                 0.493
2 anti-commonsensical                 0.513                 0.453
3         nonsensical                 0.467                 0.467
  debiased_Qwen2.5-0.5B debiased_Qwen2.5-1.5B
1                 0.540                 0.460
2                 0.513                 0.487
3                 0.527                 0.507

cl |> group_by(model, story, rung) |> summarise(k = sum(correct_debiased), n = n(), .groups = "drop") |>
  mutate(est = (k + 1) / (n + 2), lo = qbeta(0.05, k + 1, n - k + 1), hi = qbeta(0.95, k + 1, n - k + 1)) |>
  ggplot(aes(factor(rung), est, colour = story)) + geom_hline(yintercept = 0.5, linetype = 2, colour = harbour[["muted"]]) +
  geom_pointrange(aes(ymin = lo, ymax = hi), position = position_dodge(width = 0.5), size = 0.3) + facet_wrap(~ model) +
  scale_colour_manual(values = unname(harbour[c("moss", "rose", "sky")])) + scale_y_continuous(limits = c(0.25, 0.85)) +
  labs(x = "rung of the ladder (1 association, 2 intervention, 3 counterfactual)", y = "accuracy", colour = NULL)

Both models are at chance. The 0.5-billion-parameter model answers 49.8% of the questions correctly and the 1.5-billion-parameter model 47.1%, where always answering "yes" would score 50%. Both lean heavily towards "yes" (78% and 60% of their answers), but even after removing that lean, by answering "yes" only when a question's score is above the model's median, their accuracies are 52.7% and 48.4%, and in the figure 16 of the 18 intervals, by model, rung and kind of story, include chance; the other two, one a little above chance and one below it, are about what 18 intervals at the 90% level would produce by chance alone. They also do no better on commonsensical stories than on anti-commonsensical ones, so at this scale there is neither causal reasoning nor much stored causal knowledge to draw on. The contrast with GPT-4's 62%, and with the 70% that step-by-step causal prompting achieves, is a reminder that what language models can do with causal questions depends steeply on scale and training, and has to be measured for each model rather than assumed. This is a design for such a measurement: a public benchmark, a balanced sample, scores rather than generated text, a correction for response bias, and intervals.

In the terms of this chapter, a language model's pretraining text is its prior over worlds (Section 14.4), and that prior contains a great deal of causal knowledge written down by people. It can help choose a DAG, propose confounders and instruments, and critique an analysis plan. What it cannot do is extract an identified effect from data whose structure leaves the effect unidentified (Section 14.7), and it may confidently report what its training text associates rather than what the data and assumptions in front of it imply. The productive division of labour, suggested by Kıcıman et al. themselves, is to let language models propose the assumptions and let formal methods, of the kind in Chapters 13 and 14, carry out the inference and report its uncertainty.

14.9 Interventions inside networks

Causal reasoning can also be turned on the network itself. A trained network is a structural causal model: inputs cause hidden activations, which cause outputs, by known equations. We can therefore intervene on a hidden unit, setting its activation to the value it would have taken on a different input, and observe the effect on the output. This is activation patching, and it is causal mediation analysis (Chapter 13) applied to a network's internal variables: the change in output when a unit is patched is the indirect effect of the input change transmitted through that unit (Vig et al., 2020). Meng et al. (2022) used it to locate where factual associations are stored in GPT-style models; Geiger et al. (2021) used interventions to test whether a network implements a hypothesised high-level causal model. The field is called mechanistic interpretability.

We apply it to the network of Section 14.7 that sees $C$. Take the worlds where adjustment matters (where the adjusted and naive regressions differ by more than 0.2). For each, compare the network's output on the real input with its output on a counterfactual input in which the two covariances involving $C$ are set to zero, as if there were no confounding; the network then returns roughly the naive answer. The difference is the total effect of the confounding information. Then patch hidden units, one at a time and in groups, with their counterfactual values, and ask how much of the total each carries.

nin <- length(inp_obs); H <- net_obs$n[2]; nW <- H * (nin + 1)
hidden <- function(X) plogis(cbind(1, X) %*% matrix(net_obs$wts[1:nW], nrow = nin + 1))   # the network's hidden layer
readout <- function(Hm) as.vector(cbind(1, Hm) %*% net_obs$wts[-(1:nW)])                  # and its output layer
X <- as.matrix(test_worlds[, inp_obs]); stopifnot(max(abs(readout(hidden(X)) - predict(net_obs, X))) < 1e-4)   # matches nnet
matters <- abs(test_worlds$adjusted - test_worlds$naive) > 0.2
X_real <- X[matters, ]; X_cf <- X_real; X_cf[, c("cCT", "cCY")] <- 0      # counterfactual input: no confounding information
H_real <- hidden(X_real); H_cf <- hidden(X_cf); total <- readout(H_real) - readout(H_cf)
patch <- function(units) { Hp <- H_real; Hp[, units] <- H_cf[, units]; readout(Hp) }
single <- sapply(1:H, function(j) mean(abs(readout(H_real) - patch(j))) / mean(abs(total)))    # one unit at a time
ord <- order(single, decreasing = TRUE)
carried <- sapply(1:H, function(k) 1 - mean(abs(patch(ord[1:k]) - readout(H_cf))) / mean(abs(total)))   # top k units together
c(worlds = sum(matters), sum_of_single_unit_effects = round(sum(single), 2))
tibble(units_patched = 1:H, share_of_total_effect = round(carried, 2)) |> filter(units_patched %in% c(1, 2, 3, 5, 10, 15, 20))

                    worlds sum_of_single_unit_effects 
                   1059.00                       8.75 
# A tibble: 7 x 2
  units_patched share_of_total_effect
          <int>                 <dbl>
1             1                 -0.91
2             2                  0   
3             3                  0.11
4             5                  0.27
5            10                  0.42
6            15                  0.61
7            20                  1   

Among the 1,059 test worlds where adjustment matters, the network's computation of the adjustment is spread across the whole hidden layer, and its units interact. Patching single units one at a time changes the output by amounts that add up to 8.75 times the total effect, so their contributions are far from additive. Patching only the most influential unit moves the output 91% of the way in the wrong direction; the top five units together carry 27% of the effect, the top ten 42%, the top fifteen 61%, and all twenty are needed for the whole. No small circuit "does the adjustment".

Even in a network with 20 hidden units trained on one clean task, "where the adjustment happens" has no simple answer, and a single-unit patching table would mislead. Large models have millions of hidden units and billions of weights, so interpretability research looks for the right variables to intervene on: directions in activation space, or groups of units, that behave like the variables of a simple causal model (Geiger et al., 2021). This is causal representation learning (Section 14.6) seen from the inside.

14.10 What it would take

Chapter 13 ended with three things that more data from the same source cannot supply. This chapter has shown each of them at work inside deep learning, and they make a short programme for causal AI.

1. Causal assumptions as structure. TARNet's two heads build the potential-outcomes structure into the architecture; the amortised network learned adjustment only because the simulated worlds encoded which variables were confounders. Assumptions can be built into architectures, priors over worlds, or prompts, but they must be stated somewhere, and they must be the right ones.
2. Interventional and multi-environment data. Invariance needs environments that differ; counterfactual augmentation creates such differences deliberately; experiments supply them directly. Models trained only on observational data from one source learn its shortcuts.
3. Honest uncertainty. The Bayesian network's intervals widened, if only a little, where treated and control children did not overlap; semantic entropy flags confabulation; a prior-fitted network's posterior is only as honest as its prior. Uncertainty must be checked against truth where truth is available (semi-synthetic benchmarks, held-out environments) before it is trusted where it is not.

None of these is specific to neural networks, and that is the point of the course's last chapter: the ideas of Bayesian inference and causal reasoning that work for a dozen variables are the ideas that make the largest models more reliable. A network is a flexible likelihood; the rest of the analysis, the estimand, the graph, the prior, the check against reality, is still ours to do.

14.11 Exercises

Each exercise has a hint and an answer button.

Further reading and references

1. Arjovsky, M., Bottou, L., Gulrajani, I. and Lopez-Paz, D. (2019). Invariant risk minimization. arXiv:1907.02893.
2. Bahdanau, D., Cho, K. and Bengio, Y. (2015). Neural machine translation by jointly learning to align and translate. International Conference on Learning Representations.
3. Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M. and Hennig, P. (2021). Laplace redux: effortless Bayesian deep learning. Advances in Neural Information Processing Systems, 34.
4. Egami, N., Fong, C. J., Grimmer, J., Roberts, M. E. and Stewart, B. M. (2022). How to make causal inferences using texts. Science Advances, 8(42), eabg2652.
5. Farquhar, S., Kossen, J., Kuhn, L. and Gal, Y. (2024). Detecting hallucinations in large language models using semantic entropy. Nature, 630, 625–630.
6. Gal, Y. and Ghahramani, Z. (2016). Dropout as a Bayesian approximation: representing model uncertainty in deep learning. Proceedings of the 33rd International Conference on Machine Learning, 1050–1059.
7. Geiger, A., Lu, H., Icard, T. and Potts, C. (2021). Causal abstractions of neural networks. Advances in Neural Information Processing Systems, 34.
8. Geirhos, R., Jacobsen, J.-H., Michaelis, C., Zemel, R., Brendel, W., Bethge, M. and Wichmann, F. A. (2020). Shortcut learning in deep neural networks. Nature Machine Intelligence, 2, 665–673.
9. Guo, B., Zhang, X., Wang, Z., Jiang, M., Nie, J., Ding, Y., Yue, J. and Wu, Y. (2023). How close is ChatGPT to human experts? Comparison corpus, evaluation, and detection. arXiv:2301.07597.
10. Hill, J. L. (2011). Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20(1), 217–240.
11. Hollmann, N., Müller, S., Purucker, L., Krishnakumar, A., Körfer, M., Hoo, S. B., Schirrmeister, R. T. and Hutter, F. (2025). Accurate predictions on small data with a tabular foundation model. Nature, 637, 319–326.
12. Jin, Z., Chen, Y., Leeb, F., Gresele, L., Kamal, O., Lyu, Z., Blin, K., Gonzalez Adauto, F., Kleiman-Weiner, M., Sachan, M. and Schölkopf, B. (2023). CLadder: assessing causal reasoning in language models. Advances in Neural Information Processing Systems, 36.
13. Kaushik, D., Hovy, E. and Lipton, Z. C. (2020). Learning the difference that makes a difference with counterfactually-augmented data. International Conference on Learning Representations.
14. Kıcıman, E., Ness, R., Sharma, A. and Tan, C. (2023). Causal reasoning and large language models: opening a new frontier for causality. arXiv:2305.00050 (revised 2024; presented at the International Conference on Learning Representations, 2025).
15. Lakshminarayanan, B., Pritzel, A. and Blundell, C. (2017). Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in Neural Information Processing Systems, 30.
16. Louizos, C., Shalit, U., Mooij, J., Sontag, D., Zemel, R. and Welling, M. (2017). Causal effect inference with deep latent-variable models. Advances in Neural Information Processing Systems, 30.
17. MacKay, D. J. C. (1992). A practical Bayesian framework for backpropagation networks. Neural Computation, 4(3), 448–472.
18. Meng, K., Bau, D., Andonian, A. and Belinkov, Y. (2022). Locating and editing factual associations in GPT. Advances in Neural Information Processing Systems, 35.
19. Müller, S., Hollmann, N., Pineda Arango, S., Grabocka, J. and Hutter, F. (2022). Transformers can do Bayesian inference. International Conference on Learning Representations.
20. Nadaraya, E. A. (1964). On estimating regression. Theory of Probability and Its Applications, 9(1), 141–142.
21. Neal, R. M. (1996). Bayesian Learning for Neural Networks. Lecture Notes in Statistics 118. Springer.
22. Peters, J., Bühlmann, P. and Meinshausen, N. (2016). Causal inference by using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Society, Series B, 78(5), 947–1012.
23. Qwen Team (2024). Qwen2.5 technical report. arXiv:2412.15115.
24. Robertson, J., Reuter, A., Guo, S., Hollmann, N., Hutter, F. and Schölkopf, B. (2025). Do-PFN: in-context learning for causal effect estimation. Advances in Neural Information Processing Systems, 38.
25. Rosenfeld, E., Ravikumar, P. and Risteski, A. (2021). The risks of invariant risk minimization. International Conference on Learning Representations.
26. Schölkopf, B., Locatello, F., Bauer, S., Ke, N. R., Kalchbrenner, N., Goyal, A. and Bengio, Y. (2021). Toward causal representation learning. Proceedings of the IEEE, 109(5), 612–634.
27. Shalit, U., Johansson, F. D. and Sontag, D. (2017). Estimating individual treatment effect: generalization bounds and algorithms. Proceedings of the 34th International Conference on Machine Learning, 3076–3085.
28. Shi, C., Blei, D. M. and Veitch, V. (2019). Adapting neural networks for the estimation of treatment effects. Advances in Neural Information Processing Systems, 32.
29. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł. and Polosukhin, I. (2017). Attention is all you need. Advances in Neural Information Processing Systems, 30.
30. Veitch, V., Sridhar, D. and Blei, D. M. (2020). Adapting text embeddings for causal inference. Proceedings of the 36th Conference on Uncertainty in Artificial Intelligence, 919–928.
31. Vig, J., Gehrmann, S., Belinkov, Y., Qian, S., Nevo, D., Singer, Y. and Shieber, S. (2020). Investigating gender bias in language models using causal mediation analysis. Advances in Neural Information Processing Systems, 33.
32. Watson, G. S. (1964). Smooth regression analysis. Sankhyā, Series A, 26(4), 359–372.
33. Xie, S. M., Raghunathan, A., Liang, P. and Ma, T. (2022). An explanation of in-context learning as implicit Bayesian inference. International Conference on Learning Representations.
34. Zečević, M., Willig, M., Dhami, D. S. and Kersting, K. (2023). Causal parrots: large language models may talk causality but are not causal. Transactions on Machine Learning Research.