Learning Instruments You Don’t Have: Synthesizing Instrumental Variables with Graph Attention

Ask a causal question, does this drug lower risk? does this policy raise wages?, and the first enemy you meet is confounding. Something you didn’t measure nudges both the treatment and the outcome, and a naïve regression happily reports that nudge as if it were the effect of the treatment itself.

The textbook cure is the instrumental variable (IV): a source of variation that moves the treatment but touches the outcome only through the treatment. The trouble is practical, good instruments are rare, and the ones people reach for are often only barely valid. This post is about a question I find genuinely fun: if you can’t find an instrument, can you learn one? It’s the idea behind my synthetic-instrumental-variables project.

The setup, briefly

Write the treatment as \(T\), the outcome as \(Y\), the covariates we observe as \(\mathbf{X}\), and the confounder we don’t as \(\mathbf{U}\). A simple structural model:

\[Y = \tau\,T + \phi(\mathbf{U}) + \varepsilon, \qquad T = \psi(\mathbf{U}, \mathbf{X}) + \eta .\]

We want \(\tau\), the causal effect of \(T\) on \(Y\). Because \(\mathbf{U}\) sits in both equations, regressing \(Y\) on \(T\) gives an estimate that is biased, the regression can’t tell \(\tau\) apart from the shared influence of \(\mathbf{U}\).

An instrument \(Z\) rescues us if it satisfies two conditions:

Relevance, it actually moves the treatment: \(\operatorname{Cov}(Z, T) \neq 0\).
Exogeneity (the exclusion restriction), it affects the outcome only through the treatment: \(\operatorname{Cov}(Z, \varepsilon) = 0\).

Given such a \(Z\), two-stage least squares (2SLS) recovers the effect. In the single-instrument case it’s the wonderfully compact ratio

\[\hat\tau_{\text{2SLS}} = \frac{\operatorname{Cov}(Z, Y)}{\operatorname{Cov}(Z, T)} .\]

The synthetic instrument Ẑ is built from covariates X and must be relevant to T, while having no direct path to Y. The unobserved confounder U (red) is what makes naïve regression biased.

The idea: build the instrument you wish you had

Instead of hunting for a pre-existing instrument, we ask a model to construct one from the covariates we already have. Not every covariate is a good ingredient, some are themselves confounded, some are irrelevant, so the model has to be selective, and it has to learn interactions, not just weights.

That’s a natural fit for a graph-attention network (GAT). Represent the covariates as nodes on a fully-connected graph and let attention decide which covariates (and which combinations) form a strong, well-behaved instrument. The network outputs a single scalar per unit, our synthetic instrument \(Z = f_\theta(\mathbf{X})\).

The whole trick lives in the loss. We want \(Z\) to be strongly relevant to the treatment, while staying uncorrelated with the part of the outcome the treatment doesn’t explain, the second-stage residual \(\hat\varepsilon\), which we use as a tractable proxy for exogeneity:

\[\mathcal{L}(\theta) = \underbrace{-\,\big|\widehat{\operatorname{Corr}}(Z, T)\big|}_{\text{be relevant}} \;+\; \lambda \underbrace{\big|\widehat{\operatorname{Corr}}(Z, \hat\varepsilon)\big|}_{\text{stay exogenous}} .\]

The two terms pull against each other, and \(\lambda\) sets the trade-off. Minimising this is, in spirit, asking the network for the most useful instrument that still behaves itself.

In code

The core is small, a two-layer GAT that emits one number per unit, and a loss that trades relevance against exogeneity:

import torch
import torch.nn as nn
from torch_geometric.nn import GATConv


class InstrumentGenerator(nn.Module):
    """Turn covariates X into a single synthetic instrument Z = f(X)."""
    def __init__(self, d_in, hidden=64, heads=4):
        super().__init__()
        self.g1 = GATConv(d_in, hidden, heads=heads, concat=True)
        self.g2 = GATConv(hidden * heads, hidden, heads=1, concat=False)
        self.head = nn.Linear(hidden, 1)

    def forward(self, x, edge_index):
        h = torch.relu(self.g1(x, edge_index))
        h = torch.relu(self.g2(h, edge_index))
        return self.head(h).squeeze(-1)          # one instrument value per unit


def pearson(a, b, eps=1e-8):
    a, b = a - a.mean(), b - b.mean()
    return (a * b).mean() / (a.std() * b.std() + eps)


def iv_loss(z, t, resid, lam=1.0):
    """Maximise relevance to T; penalise correlation with the 2SLS residual."""
    relevance  = pearson(z, t).abs()             # want this large
    exogeneity = pearson(z, resid).abs()         # want this near zero
    return -relevance + lam * exogeneity

Training alternates between fitting 2SLS with the current instrument (to get the residual \(\hat\varepsilon\)) and nudging the generator to lower iv_loss. The edge_index wires up the covariate graph; in the simplest version it’s fully-connected and attention does the pruning.

Does it help?

I tested it on the LaLonde (NSW) benchmark, the classic stress-test where a naïve comparison is badly confounded and the experimental ground truth is known. Pairing the learned instrument with 2SLS pulled the treatment-effect estimate substantially back toward the experimental benchmark, cutting bias relative to ordinary least squares (about 18% on my synthetic-confounding experiments, \(n \approx 10\text{k}\)).

It is emphatically not magic, and it’s worth being honest about why:

Exogeneity can’t be conjured from nothing. If a covariate is itself confounded, any instrument built from it inherits that flaw. The penalty term is a proxy, it discourages the symptom, not the disease.
It needs overlap and signal. With weak covariates or thin support, the “instrument’’ is weak too, and weak instruments have their own well-known pathologies.
Validity remains an assumption. No loss function can prove the exclusion restriction. I treat synthetic IVs as a tool for exploration and sensitivity analysis, one lens among several, rather than a black box that stamps a causal number.

Why I like this direction

It sits exactly where my research lives: taking a clean idea from causal inference and asking what a modern, differentiable, representation-learning view adds, without letting the flexibility of deep learning quietly launder away the assumptions that make causal claims meaningful. Learned instruments won’t replace careful design, but as a way to search the space of candidate instruments and to reason about robustness, I think they’re a genuinely useful addition to the toolbox.

The code, including the GAT generator, the custom loss, and the 2SLS evaluation against biased OLS, is on GitHub, feedback and issues are very welcome.