Causal Mediation Analysis in Language Models

Description

Causal mediation analysis traces which internal components of a language model carry the effect of a change in the input to the output. The analyst makes a controlled change to the input, such as swapping 'nurse' for 'man' in a prompt, and measures how the model's next-word probabilities shift. To find the indirect effect through one component (a single neuron, or an attention head, which is the part of the model that chooses which earlier words to focus on), they feed the model the original input but give that component the value it took on the changed input. To find the direct effect, they feed the changed input but hold that component at its original value. Repeating this across many components shows where in the network a behaviour is carried.

Example Use Cases

Safety

A safety team finds that a language model gives unsafe medical advice when a request is phrased as an instruction rather than a question. They compare outputs for the two phrasings and use mediation analysis to find which attention heads carry most of the shift in the probability of the unsafe reply, so a fix can target those heads and the analysis can be rerun to check it worked.

Explainability

A research team studying why a language model sometimes gets subject-verb agreement wrong in sentences with an intervening clause changes the number of the subject noun and measures the shift in the verb's probability. Mediation analysis then shows which neurons carry the number information, and whether the same neurons are used across sentence structures.

Limitations

Needs white-box access and knowledge of the architecture. The analyst must read and overwrite activations at chosen points inside the model, so it cannot be used on a model offered only through an API, and a poorly chosen intervention point can give misleading conclusions.
Conclusions hold for the contrast tested. Swapping 'nurse' for 'man' and tracking one probability ratio tests a single contrast, and a different prompt set or outcome measure can point to different components. Effects also differ between models: larger GPT-2 models showed larger effects, and masked language models matched less well.
Every component and every prompt needs its own forward passes, so testing all the neurons or heads of a large model on many prompts is expensive. Analysts often test a subset or groups of components, which can miss effects elsewhere in the network.
Single-component effects do not simply add up. Vig et al. found that intervening on all neurons together produced an effect far larger than the sum of the single-neuron effects, so ranking components one at a time can understate or misplace where a behaviour is carried.
The split into direct and indirect effects is not guaranteed to be exact. In a non-linear network the total effect need not equal their sum, and Vig et al. report only an approximate match for GPT-2, so the split should be checked rather than assumed.

Resources

Research Papers

Investigating Gender Bias in Language Models Using Causal Mediation Analysis

Jesse Vig et al.•Apr 26, 2020

Common methods for interpreting neural models in natural language processing typically examine either their structure or their behavior, but not both. We propose a methodology grounded in the theory of causal mediation analysis for interpreting which parts of a model are causally implicated in its behavior. It enables us to analyze the mechanisms by which information flows from input to output through various model components, known as mediators. We apply this methodology to analyze gender bias in pre-trained Transformer language models. We study the role of individual neurons and attention heads in mediating gender bias across three datasets designed to gauge a model's sensitivity to gender bias. Our mediation analysis reveals that gender bias effects are (i) sparse, concentrated in a small part of the network; (ii) synergistic, amplified or repressed by different components; and (iii) decomposable into effects flowing directly from the input and indirectly through the mediators.

Software Packages

pyvene: A Library for Understanding and Improving PyTorch Models via Interventions

Oct 5, 2026

Stanford NLP Python library for understanding and improving PyTorch models via interventions

CausalMediationAnalysis (code for Vig et al. 2020)

Jul 15, 2026

Code for the paper "Causal Mediation Analysis for Interpreting Neural NLP: The Case of Gender Bias"

Tutorials

Causal Tracing (ROME) — pyvene tutorial

Stanford NLP Python library for understanding and improving PyTorch models via interventions - stanfordnlp/pyvene

Documentations

pyvene documentation

Related Techniques

Name	Description	Assurance Goals
Neuron Activation Analysis	Neuron activation analysis examines when individual neurons in a neural network respond strongly. The analyst runs a set of inputs through the model, finds the inputs that activate a chosen neuron most and least, and looks for the concept that separates them, such as buildings in an image classifier or French text in a language model. Some methods test the match automatically, for example by scoring how closely a neuron's activation lines up with labelled examples of a concept. The result is a list of neurons with candidate meanings and the evidence behind each. These meanings are correlations between inputs and activations, so they show what a neuron responds to, not what it does to the model's output.	Explainability
Sparse Autoencoder Feature Analysis	Sparse autoencoder feature analysis looks for human-interpretable units inside a neural network, usually a language model, by learning from the activations at one layer a large dictionary of directions, called features, such that each activation vector can be rebuilt from only a few of them. Unlike neuron activation analysis it does not take the network's own neurons as the units, since single neurons often respond to several unrelated things, and unlike concept activation vectors it needs no labelled examples of a concept chosen in advance, because the features are found without supervision. The analyst records activations at the chosen layer over a large body of ordinary text and trains an autoencoder, a network with one hidden layer usually several times wider than the layer studied, to reconstruct them while a sparsity penalty keeps most hidden units at zero; the decoder weights of each hidden unit then give one feature direction. Each feature is labelled from the inputs on which it fires most strongly, by a person or by a language model that writes a description, which is then scored on how well it predicts the feature's activations on other text, and the label is tested by removing or amplifying the feature inside the running model and checking the change in output. The output is a catalogue of labelled features with example activations, interpretability scores and measured effects, plus a reconstruction score saying how much of the layer's behaviour the dictionary accounts for; the features can then serve as units for causal mediation analysis, which traces pathways, while this method finds and checks the units themselves.	Explainability
Attention Weight Analysis	Attention weight analysis reads the weights an attention-based neural network assigns to the parts of its input, to see which positions it draws on at each step. The weights come from an ordinary forward pass and are plotted as heatmaps, one row per output step or token. In recurrent encoder-decoder models, such as early neural machine translation or image captioning, each output step has a row over the source words or image regions, read as a soft alignment in which weight is shared across several source positions. In transformers, each attention head (one of several parallel attention units in a layer) has its own matrix at every layer; the analyst views these per head and layer, or averages them over many inputs to profile a head, for example by how much attention it gives each kind of token, or how often its strongest attention lands on a labelled relation such as a verb and its object. The heatmaps show where each step looks; on their own they do not measure how much each input contributed to the prediction, because the weights leave out what is done with the attended values and information mixes between layers.	Explainability