II · THE IDEA · ARTIFICIAL INTELLIGENCE
Causal Mediation Analysis in Neural Networks
▶ Listen · narrated
When a network gets something right, the useful question is often not only what it used, but which internal route carried the effect from input to output.
At a glance
- Aim
- Separate a direct path from a path through a mediator
- Mediator
- An internal quantity whose role is under test
- Contrast
- Clean versus corrupted (or alternative) inputs
- Output
- A decomposition of effect, not a full causal graph
Think of a message sent across a city. You change the starting instruction and watch whether the final delivery changes. Then you ask how much of that change travelled through one particular junction—a bridge you can temporarily freeze or restore—and how much found other streets.
Causal mediation analysis on a neural network is that kind of bookkeeping. You pick an internal part of the model as the junction. You run a normal input and a changed input. You swap the junction’s internal state between those runs and see how much of the answer gap comes back. The share that comes back through the junction is the indirect path. The share that remains even when the junction is held fixed is the direct path.
The method does not prove that the junction is the only place the work happens. It only measures how much of this particular change, on this task contrast, appears to travel through the part you opened.
Causal mediation in neural networks estimates how much of an input-induced change in an outcome is transmitted through a selected internal mediator m. Let X be an input intervention (often a clean versus corrupted prompt pair), Y an output functional such as a logit difference, and m = m(x) an activation or subspace computed on the forward pass. The total effect contrasts Y under x_clean and x_corr. An indirect effect is obtained by interventions that set m to its value from the other run while leaving the rest of the computation to proceed; a direct effect holds m fixed while the input contrast is applied.
In implementation this is almost always activation patching or path patching rather than weight-level surgery: cache activations on both runs, splice m at a chosen layer or component, and re-measure Y. Natural indirect effect and natural direct effect style estimands from statistics are approximated by these splices, but only under the assumption that the patch isolates the mediator of interest and that the corruption defines a meaningful counterfactual on X.
Limitations are practical and structural. Mediators that are too wide absorb unrelated computation; mediators that are too narrow miss distributed pathways. Correlated components make ‘direct’ a residual label, not a claim of zero elsewhere. Nonlinearities and residual connections mean a single splice may not correspond to a unique graphical path. The decomposition is therefore design-relative: it is evidence about the share of an effect recoverable through m under a named intervention, not automatic discovery of a complete circuit.
Look closer
Two routes, one measured outcome
The analysis treats a model run as a small causal system: an input intervention, a chosen internal mediator, and a downstream output. The indirect path is the portion of the output change that appears to travel through that mediator. The direct path is the portion that remains when the mediator is not allowed to carry the difference. What you see in practice is not a story about every weight, but a contrast between runs that differ only where the experimenter intervenes.
The mediator is a choice, not a discovery
Nothing in the method names the ‘true’ internal cause on its own. An experimenter picks a layer, a neuron set, an attention head, or another quantity and asks how much of an effect is mediated by that choice. A large mediated share means the selected part is sufficient to carry much of the measured difference under the intervention design. A small share means the effect largely bypasses it—or that the mediator was cut too coarsely to show its role.
Interventions must be paired
Typical designs compare a clean input with a corrupted or counterfactual one, then splice internal states between those runs. The visible object is a set of paired forward passes: restore the mediator from the clean run into the corrupted run, or vice versa, and watch how much of the output gap closes. The numbers only mean what the splice allows; they do not, by themselves, prove that the mediator is the only locus of the computation.
The story
Causal mediation analysis, used on neural networks, is a way of splitting a measured change in behaviour into a part that runs through a chosen internal quantity and a part that does not. The motivation is simple enough: knowing that a model relies on a feature or a representation is weaker than knowing whether that reliance is concentrated in a particular internal locus or widely distributed.
The basic setup has three roles. There is an input factor the experimenter can change. There is a mediator—an activation, a residual stream vector, a head output, or another internal signal selected in advance. And there is an outcome, usually a logit, a probability, or a task score. An effect of the input on the outcome is then asked to divide into an indirect effect, travelling via the mediator, and a direct effect, travelling by other internal routes.
In network practice this almost always means interventions on activations rather than training-time surgery. One common pattern is to run the model on a clean input and on a corrupted or alternative input that removes or alters the cue of interest. Internal states from one run are then inserted into the other at the mediator’s location. If restoring the clean mediator largely restores the clean behaviour, the mediator is carrying a substantial indirect effect under that design. If behaviour barely moves, the measured effect is mostly direct relative to that mediator—or the mediator was defined at the wrong grain.
The language of direct and indirect paths is borrowed from causal mediation in statistics, but the network version is only as strong as its intervention. The method does not automatically yield a full causal graph of the model. It answers a narrower question: of the output gap produced by this input contrast, how much is blocked or restored when we hold this particular internal object fixed. Change the contrast, the mediator, or the splicing point, and the split can change with them.
That narrowness is also the method’s discipline. It forces a claim about ‘where the computation lives’ to be stated as a measurable share of an effect, under named interventions, rather than as a general impression from saliency or probing alone. Probes can show that information is present; mediation-style interventions ask whether that information is on the route the model actually uses for the behaviour under study.
Results are therefore easiest to trust when the task contrast is sharp, the mediator is modular enough to splice, and the outcome metric is stable. They are hardest to trust when many entangled parts each carry a little of the effect, when the corruption changes more than the intended cue, or when the mediator is a large block that mixes several computations. In those regimes the decomposition still produces numbers, but the numbers under-determine the mechanistic story.
Used carefully, the approach sits between purely observational interpretability and full causal discovery. It does not replace either. It adds a bookkeeping tool: a way to say, of a known behavioural effect, how much appears to travel through a part we chose to open, and how much appears to go around it.
Why it mattered then
As interpretability moved from asking what a representation contains to asking what a model uses, simple correlations became insufficient. A probe can light up on a feature that is present but idle for the decision at hand. Mediation-style decompositions mattered because they tied internal parts to behavioural contrasts with explicit interventions. In that moment the field needed language for partial, path-specific claims: not ‘this neuron is the dog detector’, but ‘this much of this output gap runs through this component under this intervention’. That shift let arguments about mechanism be checked against spliced runs rather than against intuition alone.
Why it matters now
Large models still offer only limited access to training data and still invite over-confident stories about single circuits. Direct and indirect effect estimates remain useful because they scale to concrete behavioural questions—factual recall, refusal, syntactic agreement, and similar contrasts—without requiring a complete schematic of the network. They also make disagreement sharper: two groups can name the same mediator and the same input contrast and still report different mediated shares if their corruptions or metrics differ. That friction is productive. It keeps interpretability claims attached to designs that others can rerun on open-weight models, and it pushes explanations toward quantities that survive intervention rather than toward narratives that only fit observational traces.
The surprising detail
The ‘direct’ effect is not a picture of a single clean wire from input to output. It is only the leftover: whatever portion of the effect is not carried by the mediator you chose. Enlarge the mediator, and some of what was called direct can reappear as indirect. Shrink it, and the reverse happens. The labels look absolute in a diagram; in the method they are relative to the cut you made.
What is disputed
The supplied source list does not include primary papers on causal mediation in neural networks; it names unrelated work on a Landau–Zener system and on concept maps for software-product-line onboarding. The lesson therefore stays at the level of method structure and interpretation and does not attribute specific numerical findings, authors, or experiment dates to those sources.
Remember this
Causal mediation splits a behavioural effect into a share through a chosen internal mediator and a share that bypasses it—under stated interventions, not as a full map of the network.
Test yourself
A team reports that an attention head mediates 80 percent of a factual-answer effect. They corrupted the subject name in the prompt and restored that head’s clean activations. What two design facts must you know before treating the 80 percent as evidence that the head is the main locus of the fact?
First, whether the corruption altered only the intended cue or also changed other routes the model might use; a broad corruption can inflate the apparent role of whatever you restore. Second, how the mediator was bounded: if the ‘head’ splice includes nearby residual or layer-norm state, part of the mediated share may belong to those companions. The percentage is a property of the intervention design, not a free-standing fact about the head alone.
Go deeper
- [2004.03863] The effect of the second nearest neighbor interaction on the population transfer in a four-particle Landau-Zener system · arxiv.org
- [2103.03829] Onboarding in Software Product Lines: ConceptMaps as Welcome Guides · arxiv.org
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.