Mechanistic Interpretability: Understanding AI System Internals for Safety

In March 2023 a small Bengaluru team fine-tuned an open-weight language model to answer questions about India's space program for a public outreach chatbot. Two weeks after launch, a schoolteacher reported that the bot confidently stated Chandrayaan-3 landed on the Moon in "October 2023" — fluent, well-formatted, wrong by two months. The obvious fix is to retrain on corrected data, but retraining a multi-billion-parameter model to change one fact is like repaving a highway to fix a single pothole: expensive, slow, and it risks damaging everything nearby. The team wanted a scalpel, not a bulldozer — a way to find the exact place inside the network where "Chandrayaan-3 landing date" is stored, and edit only that. This is the question mechanistic interpretability answers: not "does the model behave correctly on this input" (which SHAP, LIME, and attention-map inspection can already tell you) but "which specific components, in which specific layer, causally produce this output — and can I intervene on exactly those components." That shift, from describing behavior to establishing causal internal structure, is what makes mechanistic interpretability a safety tool rather than a debugging convenience: the same machinery that locates a stale fact can locate a deceptive or harmful capability and suppress it at its source.

From Behavior to Mechanism

Post-hoc explanation methods treat the network as a function to be probed from outside: perturb the input, watch the output move, attribute credit to input features. This tells you what the model is sensitive to, not how it computes the answer. Mechanistic interpretability instead opens the box and asks about the actual computational graph — residual stream, attention heads, MLP blocks — as engineered structure, and demands causal evidence: if I surgically change this internal quantity and nothing else, does the output change in the way my theory predicts? Two ideas dominate this chapter, both central to current safety research at labs building frontier LLMs: first, that individual neurons are usually not clean, human-interpretable units of meaning, because of a phenomenon called superposition; second, that once you have a way to recover interpretable units, you can trace them causally through the network's layers to locate where a specific piece of knowledge or behavior lives, and edit it there. This is the actual research lineage behind our Chandrayaan example: Elhage et al. (2022), "Toy Models of Superposition" (Anthropic); Bricken et al. (2023), "Towards Monosemanticity" (Anthropic); Templeton et al. (2024), "Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 Sonnet" (Anthropic); Meng, Bau, Andonian, and Belinkov (2022), "Locating and Editing Factual Associations in GPT" (NeurIPS), the paper that introduced the ROME method and the causal-tracing procedure this chapter works through by hand.

Superposition: Why a Neuron Rarely Means One Thing

A transformer's hidden layers are narrow relative to the number of distinct concepts a large model needs to represent — think millions of facts, styles, and abstractions squeezed through a residual stream of a few thousand dimensions. If each concept needed its own dedicated neuron, the model would run out of room almost immediately. Elhage et al. (2022) showed that networks resolve this by packing more features than they have dimensions into the same space, tolerating a controlled amount of cross-talk between features as the price of extra capacity. This works well precisely when features are sparse — rarely active at the same time — because a ReLU nonlinearity can clip away small negative interference before it corrupts the output.

Here is the mechanism made concrete, in two neurons. Suppose three features must be packed into a 2-dimensional hidden layer. Give each feature its own unit-length direction in that 2D space, spaced 120° apart so their pairwise cosine similarity is exactly −0.5:

import numpy as np

# Three feature directions in a 2-neuron hidden space, 120 degrees apart
W = np.array([
    [ 1.0,        0.0       ],   # feature 1
    [-0.5,        0.8660254 ],   # feature 2
    [-0.5,       -0.8660254 ],   # feature 3
])

def encode_decode(x):
    h = x @ W                      # compress 3 features into 2 "neurons"
    x_hat = np.maximum(h @ W.T, 0.0)   # decode, then ReLU
    return h, x_hat

# Case A: sparse — only feature 1 fires
x_a = np.array([0.8, 0.0, 0.0])
h_a, xhat_a = encode_decode(x_a)
print(h_a, xhat_a)

Trace it by hand. h_a = x_a @ W = 0.8*[1,0] = [0.8, 0.0]. Decoding multiplies h_a by each feature direction: x̂₁ = h_a·W[0] = 0.8·1 + 0·0 = 0.8; x̂₂ = h_a·W[1] = 0.8·(−0.5) + 0·0.866 = −0.4; x̂₃ = h_a·W[2] = 0.8·(−0.5) + 0·(−0.866) = −0.4. After ReLU, the negative entries vanish and x̂ = [0.8, 0.0, 0.0] — an exact reconstruction of a feature vector using only two neurons for three features. The interference is real (raw dot products of −0.4 for the silent features) but the nonlinearity absorbs it completely, because only one feature was active.

Now break the sparsity assumption — fire two features at once, feature 1 at 0.8 and feature 2 at 0.6:

x_b = np.array([0.8, 0.6, 0.0])
h_b, xhat_b = encode_decode(x_b)
print(h_b, xhat_b)

h_b = 0.8·[1,0] + 0.6·[−0.5, 0.8660254] = [0.5, 0.51962]. Now decode: x̂₁ = 0.5·1 + 0.51962·0 = 0.5 (true value 0.8 — a −0.3 error). x̂₂ = 0.5·(−0.5) + 0.51962·0.8660254 = −0.25 + 0.450 = 0.200 (true value 0.6 — a −0.4 error). x̂₃ = 0.5·(−0.5) + 0.51962·(−0.8660254) = −0.25 − 0.450 = −0.700, clipped by ReLU to 0 (true value 0, correctly suppressed). Compressing two simultaneously active features into two neurons is no longer free: reconstruction degrades to [0.5, 0.2, 0.0] against a true [0.8, 0.6, 0.0]. This is the superposition trade-off in miniature — a network happily represents more features than it has neurons, but pays an interference cost that grows as sparsity breaks down. It is also why a single neuron's top-activating examples are a poor guide to "what it means": in this toy model, neuron 1 (the h_x coordinate) fires positively for feature 1 and negatively for features 2 and 3 simultaneously — it is not "about" any one concept, it is a coordinate in a shared code.

Sparse Autoencoders: Un-mixing Superposed Activations

If real hidden layers are superposed mixtures like this, reading individual neurons will mislead you. Bricken et al. (2023) and, at production scale, Templeton et al. (2024) address this with a sparse autoencoder (SAE): a small network trained to take the model's actual (superposed) activation vector h and re-expand it into a much larger, overcomplete set of "dictionary" features, most of which are forced to be exactly zero for any given input. The forward pass is f = ReLU(W_enc·h + b_enc), ĥ = W_dec·f, trained to minimize reconstruction error ‖h − ĥ‖² plus an L1 sparsity penalty λ·Σ|f_i| on the feature activations. The L1 term is not optional decoration — without it, the SAE has no pressure to keep few features active per input, and it will simply relearn a re-entangled, equally uninterpretable code. Templeton et al. (2024) trained SAEs with dictionaries far larger than Claude 3 Sonnet's residual-stream width — into the millions of learned directions — specifically because superposition means you need more "slots" than dimensions to pull the mixture apart, and among the recovered directions they found individually interpretable, safety-relevant features, including ones that responded selectively to concepts like sycophancy, deception, and insecure code.

Our toy example already contains a working SAE, hiding in plain sight. If the encoder weights equal the true feature directions (W_enc = W, b_enc = 0) and the decoder is its transpose, then f = ReLU(W·h) is exactly the x̂ we computed above by hand — case A recovers f = [0.8, 0, 0] perfectly, a clean, monosemantic, single-feature code, because the input was sparse. A real SAE never gets handed W; it has to discover a good approximation of it purely from the statistics of many activation vectors, which is why it needs a large amount of data and an overcomplete dictionary: enough directions on offer, and enough sparsity pressure, that it converges on approximately the network's true (if implicit) feature basis rather than some other, tangled decomposition that fits the same reconstruction loss.

Causal Tracing: Finding Where a Fact Lives

Knowing that a fact is represented somewhere as a sparse, superposed feature is not the same as knowing where in the network's depth it is computed and stored. Meng et al. (2022) solved this with causal tracing, the technique behind ROME (Rank-One Model Editing), and it is the tool the Bengaluru team actually needs. The procedure has three runs:

1. Clean run: feed the true prompt ("The Chandrayaan-3 lander touched down on the Moon in ___"), record the hidden state at every layer and token position, and note that the model assigns high probability to the correct continuation.
2. Corrupted run: add noise to the embeddings of the subject tokens ("Chandrayaan-3"), which destroys the model's grip on which entity is being discussed; the correct continuation's probability collapses.
3. Patched run: for a chosen layer and token position, run the corrupted prompt again but splice in — "patch" — the clean hidden state at just that one (layer, position). Measure how much probability of the correct token comes back.

Repeating step 3 across every layer and position produces a map of causal responsibility: the (layer, position) pairs where patching recovers the most probability are the ones doing the causal work of storing that association. Geva et al. (2021), "Transformer Feed-Forward Layers Are Key-Value Memories" (EMNLP), gives the mechanistic reason this tends to concentrate in mid-network MLP blocks at the subject's final token: an MLP's first weight matrix acts like a bank of "keys" that detect when a particular entity or concept is present in the residual stream, and its second weight matrix writes a corresponding "value" back in — literally a key-value lookup implemented in feed-forward weights, sitting exactly at the token position where the model has just finished resolving which entity the sentence is about.

The loop, structurally, looks like this (the two acts that produce and read hidden states are standard forward-hook instrumentation, assumed here rather than reimplemented):

# run_with_patch(prompt_ids, layer, pos, patch_vector) -> P(correct_token)
# clean_acts[layer][pos] holds hidden-state vectors from a prior clean
# forward pass, captured via forward hooks.
# (both assumed helpers, not shown)

n_layers, n_positions = 8, 6
recovery = [[0.0] * n_positions for _ in range(n_layers)]

for layer in range(n_layers):
    for pos in range(n_positions):
        patch_vector = clean_acts[layer][pos]
        recovery[layer][pos] = run_with_patch(
            corrupted_prompt_ids, layer, pos, patch_vector
        )

The diagram below shows what a run of this loop looks like at the position of the subject's last token, across a small illustrative 8-layer stand-in network (the numbers are constructed for teaching, not measurements from a real model — Meng et al.'s actual GPT-2/GPT-J results show the same qualitative shape: near-zero recovery at early layers, a sharp peak in a narrow mid-network band, decay afterward). Patching layer 0 barely moves the needle (probability stays near the corrupted baseline of 0.03) because too little computation has happened yet for that layer's state to matter downstream. Patching layer 4 recovers the most — probability jumps to 0.62 — identifying that as the layer where the model has both resolved "which entity" and written the associated fact into the stream. Later layers recover less, because by then the (still-corrupted) information has already been used to route attention and build downstream representations that a single later patch cannot undo.

Once causal tracing has localized the association to layer 4's MLP at the subject's last token, ROME edits it directly: it treats that MLP's second weight matrix as storing key→value associations and computes a targeted, rank-one update that changes the value written for the "Chandrayaan-3 → landing date" key, leaving every other stored association untouched (to first order). The Bengaluru team's fix, then, is not "retrain the model" but "locate the layer, verify the localization with the recovery-probability curve, apply a rank-one edit at that layer." The same machinery generalizes past facts: if a model has learned a specific unsafe behavior — say, a narrow circuit that produces compliant-sounding text for a class of harmful requests — causal tracing can, in principle, localize which layer's components are causally responsible for that behavior, which is what makes this a safety technique and not merely a debugging one. Templeton et al. (2024) is the evidence this generalizes past toy examples: at production LLM scale, SAE-recovered features included ones that tracked concepts directly relevant to safety monitoring, not just factual recall.

A Common Misconception

Students who have just learned that "neurons fire for things" often reach for a specific bad habit: pick a neuron, feed the model many inputs, read off the ones that make it fire hardest, and declare that the neuron "represents" that concept. This is exactly the mistake superposition warns against. In our two-neuron toy model, neuron 1 (the h_x coordinate) fires positively when feature 1 is present and negatively when features 2 or 3 are present — inspecting only its top-activating inputs would show you feature 1's examples and miss that the same neuron is doing double duty as negative evidence for two other, unrelated concepts. A neuron's activation is a coordinate in a shared, compressed code, not a label. The correct unit of analysis is a direction in activation space — potentially a combination of many neurons — recovered by an SAE or isolated by causal intervention, and even then, "this direction correlates with concept X" is a hypothesis to test causally (does patching it change the output the way the hypothesis predicts?), not a conclusion to read off a scatter plot of top activations.

Active Recall

Attempt each question before reading its answer.

1. What does it mean for a network to represent features "in superposition," and why does it happen at all instead of the network just using one neuron per feature?
2. Using the three-feature/two-neuron toy model (feature directions at 0°, 120°, 240°), compute the reconstructed x̂ if only feature 3 fires, with x = [0, 0, 1.0]. Show the intermediate h.
3. Take worked example "Case B" (features 1 and 2 both active, x = [0.8, 0.6, 0.0], giving x̂ = [0.5, 0.2, 0.0]). Suppose feature 2's activation magnitude rises from 0.6 to 0.9. Recompute h and the full reconstructed x̂, and describe how the errors on all three features change, not just feature 2's.
4. Why must a sparse autoencoder's dictionary (hidden) layer be larger than the model's residual-stream width, and why does removing the L1 penalty break the method even if reconstruction loss stays low?
5. In causal tracing, why does the strongest causal effect for a factual association typically show up at mid-network MLP layers at the subject's last token, rather than at attention layers or at the final token of the sentence? What would you conclude if patching every layer produced almost no recovery?
6. Beyond correcting a wrong fact, name one way that layer-localization (causal tracing plus SAE features) could be used for AI safety specifically, and one limitation of relying on it.

Answers.

1. A network has far fewer neurons in a layer than the number of distinct features useful for its task, so it packs multiple features into the same dimensions rather than dedicating one neuron per feature. This is safe when features are sparse (rarely co-active), because a ReLU-style nonlinearity clips away the resulting small negative interference; the network trades a controlled reconstruction error for extra representational capacity, and Elhage et al. (2022) show this trade is favorable whenever the loss reduction from representing more features outweighs the interference cost.

2. h = x₃ · row2 = 1.0 · [−0.5, −0.8660254] = [−0.5, −0.8660254]. Decoding: x̂₁ = h·row0 = (−0.5)(1) + (−0.8660254)(0) = −0.5 → ReLU → 0 (matches true 0). x̂₂ = h·row1 = (−0.5)(−0.5) + (−0.8660254)(0.8660254) = 0.25 − 0.75 = −0.5 → ReLU → 0 (matches true 0). x̂₃ = h·row2 = (−0.5)(−0.5) + (−0.8660254)(−0.8660254) = 0.25 + 0.75 = 1.0 (matches true 1.0). Reconstruction is exact — by the 120°-symmetric construction, any single active feature reconstructs perfectly, regardless of which one it is.

3. New input x = [0.8, 0.9, 0.0]. h = 0.8·[1,0] + 0.9·[−0.5, 0.8660254] = [0.8 − 0.45, 0 + 0.779423] = [0.35, 0.779423]. Decoding: x̂₁ = 0.35·1 + 0.779423·0 = 0.35 (true 0.8, error −0.45 — worse than the −0.30 error in the original case B). x̂₂ = 0.35·(−0.5) + 0.779423·0.8660254 = −0.175 + 0.675 = 0.50 (true 0.9, error −0.40 — about the same magnitude as before). x̂₃ = 0.35·(−0.5) + 0.779423·(−0.8660254) = −0.175 − 0.675 = −0.850 → ReLU → 0 (still correctly suppressed, true value 0). The ripple: raising feature 2's magnitude does not just worsen feature 2's own reconstruction — it worsens feature 1's reconstruction more than proportionally (error grew from −0.30 to −0.45) because the shared code point h moved further from feature 1's clean direction, while feature 3 stays safely zeroed regardless, since ReLU only ever needs the projection onto its direction to stay negative. This is the general shape of superposition interference: it is asymmetric and depends on geometry, not just on how much any one feature changed.

4. Superposition packs more features into a layer than it has dimensions, so undoing the compression needs more "slots" than the original width — an overcomplete dictionary. Templeton et al. (2024) used dictionaries with millions of learned directions, far exceeding Claude 3 Sonnet's residual width, precisely to have enough room to separate fine-grained features. The L1 penalty is what forces the code to be sparse per input; without it, the SAE can hit low reconstruction loss by activating many overlapping dictionary atoms for every input (effectively re-creating a superposed, entangled code), which reconstructs well but produces features that are just as polysemantic and uninterpretable as the original neurons — defeating the entire purpose.

5. Geva et al. (2021) show that a feed-forward (MLP) block's first matrix acts as a bank of key-detectors and its second matrix writes an associated value into the residual stream — a key-value lookup implemented in weights. That lookup fires when the relevant entity has just been resolved, which is why the causal effect concentrates at the subject's last token (the point where the model has finished encoding which entity is being discussed) and at mid-network layers (early enough that later layers can still use the corrected value, late enough that entity resolution has already happened). If patching every layer recovered almost nothing, the association is not cleanly localized to a single-point intervention — it might be diffusely spread across many layers and positions, or later computation might depend on other, unpatched tokens (e.g., attention pulling context from elsewhere) that this single-point patch cannot fix. Either way, it signals that a ROME-style single-layer edit will not reliably work for that association, and a broader intervention (multi-layer editing or fine-tuning) is needed instead.

6. Safety use: once a specific harmful or deceptive capability is causally localized to a layer or a small set of SAE-recovered feature directions, you can suppress or edit it directly — ablate the direction, or apply a rank-one weight edit — without retraining the whole model, and Templeton et al. (2024) found SAE features in a production model that tracked exactly this kind of safety-relevant concept (e.g., features responsive to deception or insecure code), which could plausibly seed an activation-based monitor that flags unsafe generation in real time. Limitation: localization found via causal tracing on a curated set of prompts may not generalize — the same underlying capability could be represented differently for different phrasings, or redundantly spread across layers, so a single-point edit or ablation found on one prompt distribution can be incomplete, and because of interference (the same problem as in superposition), suppressing one direction can degrade unrelated associations that happen to share components of that direction.

Think About It

Think about this: How would you explain mechanistic interpretability: understanding ai system internals for safety to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where mechanistic interpretability: understanding ai system internals for safety is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting mechanistic interpretability: understanding ai system internals for safety to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind mechanistic interpretability: understanding ai system internals for safety, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.