Understanding the Causal Transformer: Estimating Counterfactual Outcomes Over Time

I'm applying to PhD programs in AI for medicine, and most of my work so far has been in medical imaging — vision transformers, attention mechanisms, representation learning for mammography. Causal inference over time is a different tradition, but a closely related one: both are ultimately about learning representations that generalize past the biases of the data they were trained on. This post is my attempt to understand Melnychuk, Frauen & Feuerriegel's Causal Transformer (ICML 2022) properly — the theory, the architecture, and the training mechanism — rather than just skimming the abstract.

I did not fully retrain their model on MIMIC-III in the time I had. What I did do: work through the math from first principles, understand every architectural choice well enough to explain why it exists, and build a small, self-contained experiment that isolates and tests the paper's central mechanism — adversarial balanced representations — on synthetic data I generated myself. The full paper's transformer backbone is really about scaling this mechanism to long, complex patient histories; the mechanism itself is what I wanted to actually understand and verify.

1. The problem: why you can't just compare treated vs. untreated patients

Start with the simplest version, no time at all. A patient can either receive a drug or not. Formally, there are two potential outcomes — $Y[1]$ if treated, $Y[0]$ if not — and the treatment effect is $Y[1] - Y[0]$. The catch, sometimes called the fundamental problem of causal inference (Rubin, Neyman): for any real patient, you observe exactly one of these two outcomes, never both. You can't rerun the same patient's life twice.

The obvious fix — compare average outcomes between treated and untreated patients — only works if treatment was assigned independently of anything that also affects the outcome. In an RCT, it is, by design. In observational data (electronic health records), it usually isn't: sicker patients are more likely to receive treatment. That shared cause is a confounder, and if you ignore it, your effect estimate isn't measuring the drug — it's measuring the drug tangled up with "who was sick to begin with."

Now add time. A patient isn't one decision; they have a trajectory: covariates $X_t$, a treatment decision $A_t$, and an outcome $Y_t$ at every step. The confounder isn't a fixed trait anymore — today's treatment depends on yesterday's outcome and covariates, and today's covariates are shaped by yesterday's treatment. This feedback loop is time-varying confounding, and it's why the static-case tricks (matching, fixed propensity scores) don't transfer.

The paper's formal target, for a patient history $\bar{H}_t = \{\bar{X}_t, \bar{A}_{t-1}, \bar{Y}_t, V\}$ and a hypothetical future treatment plan $\bar{a}_{t:t+\tau-1}$:

$$\mathbb{E}\left[Y_{t+\tau}[\bar{a}_{t:t+\tau-1}] \mid \bar{H}_t\right]$$

In words: given everything known about this patient so far, what would their outcome be $\tau$ steps ahead, if this specific treatment sequence were applied? Three assumptions make this estimable from observational data at all (Appendix A): consistency (a patient's potential outcome under the treatment they actually received equals what you observed — the bridge connecting "potential" to "observed"), sequential overlap (nobody was treated or untreated with probability exactly 0 or 1, or there's no data to learn the counterfactual from), and sequential ignorability (no unobserved confounders — everything relevant to both treatment and outcome is captured in the recorded history). That third one is untestable from data alone and is the framework's real weak point — worth being upfront about rather than treating the method as a free lunch.

2. Why a transformer, and why three of them

Standard sequence models (LSTMs, as used in the prior state of the art — RMSNs, CRN, G-Net) process one sequence. The Causal Transformer's architectural bet is that a patient trajectory is really three intertwined sequences — covariates, treatments, outcomes — that causally influence each other, so it runs three parallel transformer subnetworks connected by cross-attention.

Self-attention itself, stripped to basics: for each time step, you compute a query (what am I looking for), and compare it against every other step's key (what do you have) via a dot product, turned into weights via softmax, which then combine the values (the actual content):

$$\text{Attn}(Q,K,V)_i = \sum_j \alpha_{ij} V_j, \qquad \alpha_{ij} = \text{softmax}_j\left(\frac{Q_i^\top K_j}{\sqrt{d_{qkv}}}\right)$$

For a causal problem, position $i$ must never attend to $j > i$ — the pre-softmax score is set to $-\infty$ for future positions so they get exactly zero weight after softmax. Get this wrong and your "results" are just future information leaking backward.

Cross-attention is identical math, except the query comes from one subnetwork while the key and value come from another. This is how, e.g., the treatment subnetwork's representation at time $t$ absorbs relevant information from the covariate subnetwork — capturing exactly the entanglement that produces confounding in the first place. The paper's ablation (Table 3) confirms this three-subnetwork structure matters: a single merged-subnetwork variant (EDCT) underperforms.

One more deliberate choice worth flagging: relative rather than absolute positional encoding. What matters for a clinical sequence is "this event was 3 steps before that one," not "this happened on day 47" — and relative encoding also generalizes to trajectory lengths never seen during training, since it only ever encodes distance, not calendar position.

3. The actual mechanism: adversarial balanced representations

This is the part I wanted to verify hands-on, because it's the part actually doing the causal-debiasing work — the transformer backbone is capacity, this loss is what makes the capacity trustworthy.

After the stacked blocks, the three subnetworks' hidden states are averaged into a single vector $\Phi_t$ per time step. The training objective wants $\Phi_t$ to satisfy two properties simultaneously:

(a) predictive of the next outcome — fit via a network $G_Y$ minimizing $\mathcal{L}_{G_Y} = \| Y_{t+1} - G_Y(\Phi_t, A_t) \|^2$

(b) non-predictive of the current treatment — a classifier $G_A$ is trained to predict $A_t$ from $\Phi_t$ as well as it can, and then $\Phi_t$ itself is pushed toward confusing that classifier, via a cross-entropy loss against a uniform distribution over treatments rather than the classic gradient-reversal trick used by the prior method (CRN).

Property (b) is the whole point: if $\Phi_t$ can't reveal what treatment was actually given, then whatever $G_Y$ predicts from $\Phi_t$ can't be riding on the "who tends to get treated" shortcut — which is precisely the shortcut that produces confounding bias. Theorem 4.1 in the paper formalizes this: the optimal balanced representation is exactly the one where the treatment-conditional distributions of $\Phi_t$ are identical across all treatment groups — provably, via a reversed-KL-divergence argument, not just empirically.

4. A small experiment: does this actually work?

To test this mechanism directly rather than take it on faith, I built the simplest possible version: one confounder, one binary treatment, no time dimension, no transformer — just enough structure to isolate whether adversarial balancing reduces confounding bias the way the theory says it should.

Setup. A confounder $Z \sim \mathcal{N}(0,1)$ (think: patient severity), a confounded treatment assignment $A \sim \text{Bernoulli}(\sigma(\gamma Z))$ with $\gamma = 2.5$ (so sicker patients are much more likely to be "treated" — strong confounding, $\text{corr}(A,Z) = 0.65$), and a true outcome $Y = \beta_A A + \beta_Z Z + \varepsilon$ with $\beta_A = 2$ (the true treatment effect I want to recover) and $\beta_Z = 3$.

I compared three estimators of the average treatment effect:

Method	Estimated ATE	Bias vs. true (2.0)
Naive regression, $Y \sim A$ only	5.886	+3.886
Oracle regression, $Y \sim A + Z$	2.010	+0.010
Adversarially balanced $\Phi(Z)$, trained via domain confusion	2.200	+0.200

The naive estimate is badly wrong — nearly 3x the true effect — entirely because it can't tell "the drug helped" apart from "sicker patients got the drug." The oracle (which gets to see $Z$ directly, hand-included in the regression) is essentially unbiased, as expected — this is the best case, and only works here because $Z$ is a single clean scalar I happen to already know about.

The adversarially balanced representation, trained with no direct access to hand-crafted deconfounding — only the outcome-prediction and domain-confusion losses — recovered an ATE of 2.2, a small fraction of the naive method's bias. That's the mechanism from Section 4.3 working, at toy scale.

The honest caveat, which turned out to be the most interesting part. I checked how well a treatment classifier could still predict $A$ from the balanced $\Phi$ after training: 80.9% accuracy, versus 50% for a truly fully-balanced representation. The balancing is real but incomplete. I tried strengthening the domain-confusion coefficient ($\alpha$) to push harder toward full balance — and it made the bias worse, not better (bias grew to +0.75 at $\alpha=2$, +1.26 at $\alpha=4$), without meaningfully improving classifier accuracy. This is not a bug in my toy code — it's the exact instability the paper names explicitly in Section 4.3: adversarial objectives of this kind are prone to oscillation, and the paper's actual fix is an exponential moving average (EMA) over model parameters during training, which my minimal demo deliberately left out to keep the mechanism isolated. Seeing the instability appear on cue, for the reason the paper says it will, was more convincing to me than a clean result would have been.

# full toy code: naive vs oracle vs adversarially-balanced ATE estimation
# (abbreviated here — full script in the repo linked below)

phi = Repr()                      # small representation network Phi(Z)
G_Y = nn.Linear(17, 1)             # outcome head: [Phi, A] -> Y
G_A = nn.Linear(16, 1)             # treatment classifier: Phi -> A

for step in range(3000):
    # G_A tries its best to predict treatment from the representation
    loss_A = bce(G_A(phi(Z).detach()), A)
    # Phi + G_Y: predict Y well, but confuse G_A toward a uniform guess
    loss_Y = mse(G_Y(cat[phi(Z), A]), Y)
    loss_conf = bce(G_A(phi(Z)), uniform_target)
    total = loss_Y + alpha * loss_conf

5. What I'd want to explore next

Three directions, roughly in order of what I'd tackle first with more time: (1) add the EMA stabilization from the paper and check whether it actually closes the gap between 80.9% and 50% classifier accuracy without the bias blowing up — direct test of their stated fix; (2) extend the toy setup to an actual short time series with feedback between $A_t$ and $Z_{t+1}$, which is the genuinely harder time-varying case the full paper targets; (3) run the real Causal Transformer on their released synthetic tumor-growth simulator, which is lighter-weight than MIMIC-III and would let me sanity-check the full architecture against my understanding of the mechanism above.