Mechanistic Interpretability

1. The Residual Stream View

Mechanistic interpretability starts with a conceptual reframing of the transformer architecture. Instead of viewing transformers as a sequence of layers that process data in order, the "residual stream" view treats the residual connection as a central communication channel. Each attention head and MLP layer reads from this stream, performs a computation, and writes its result back to the stream by adding it to the running total.

2. Superposition and Polysemanticity

A central challenge in interpreting neural networks is that individual neurons are often polysemantic: a single neuron may activate for multiple unrelated concepts. For example, a neuron might fire for both "legal terminology" and "the color blue." This happens because of superposition: the model encodes more features than it has neurons by representing features as directions in activation space rather than individual neurons.

Superposition is an efficient compression strategy. If a model has 4096 neurons per layer but needs to track 100,000 features, it can represent each feature as a direction in the 4096-dimensional space. As long as features rarely co-occur, they can share neurons with minimal interference. This means the "true" features of the model are not individual neurons but rather directions in activation space, which are linear combinations of neurons.

3. Sparse Autoencoders (SAEs)

Sparse autoencoders are the primary tool for disentangling superposed representations. An SAE takes the model's activations as input, encodes them into a much higher-dimensional but sparse representation, and then decodes back to the original activation space. The key constraint is sparsity: only a small fraction of the SAE's latent dimensions are active for any given input. Each active dimension corresponds to a single interpretable feature.

# Sparse Autoencoder for Mechanistic Interpretability
import torch
import torch.nn as nn
import torch.nn.functional as F

class SparseAutoencoder(nn.Module):
    """
    Sparse Autoencoder for extracting interpretable features
    from transformer activations.

    Architecture:
    - Encoder: project d_model -> d_sae (expansion)
    - ReLU sparsity
    - Decoder: project d_sae -> d_model (reconstruction)

    The expansion factor (d_sae / d_model) is typically 4x to 64x.
    """

    def __init__(self, d_model, d_sae, l1_coeff=5e-3):
        super().__init__()
        self.d_model = d_model
        self.d_sae = d_sae
        self.l1_coeff = l1_coeff

        # Encoder and decoder
        self.encoder = nn.Linear(d_model, d_sae, bias=True)
        self.decoder = nn.Linear(d_sae, d_model, bias=True)

        # Initialize decoder weights to unit norm
        with torch.no_grad():
            self.decoder.weight.data = F.normalize(
                self.decoder.weight.data, dim=0
            )

    def encode(self, x):
        """Encode activations into sparse feature space."""
        # Subtract decoder bias for centering
        x_centered = x - self.decoder.bias
        # Encode and apply ReLU for sparsity
        z = F.relu(self.encoder(x_centered))
        return z

    def decode(self, z):
        """Reconstruct activations from sparse features."""
        return self.decoder(z)

    def forward(self, x):
        z = self.encode(x)  # sparse features
        x_hat = self.decode(z)  # reconstruction
        return x_hat, z

    def loss(self, x):
        """Combined reconstruction + sparsity loss."""
        x_hat, z = self.forward(x)

        # Reconstruction loss (MSE)
        recon_loss = F.mse_loss(x_hat, x)

        # Sparsity loss (L1 on activations)
        l1_loss = z.abs().mean()

        return recon_loss + self.l1_coeff * l1_loss, {
            "reconstruction": recon_loss.item(),
            "sparsity": l1_loss.item(),
            "alive_features": (z > 0).any(dim=0).sum().item(),
            "avg_active": (z > 0).float().mean().item(),
        }

# Training an SAE on GPT-2 MLP activations
def train_sae_on_model(
    model_name="gpt2",
    layer_idx=6,
    component="mlp",
    expansion_factor=8,
    num_tokens=10_000_000,
    batch_size=4096,
    lr=3e-4,
):
    """Train a sparse autoencoder on transformer activations."""
    from transformers import AutoModelForCausalLM, AutoTokenizer

    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    d_model = model.config.n_embd
    d_sae = d_model * expansion_factor

    sae = SparseAutoencoder(d_model, d_sae)
    optimizer = torch.optim.Adam(sae.parameters(), lr=lr)

    # Collect activations using hooks
    activations_buffer = []

    def hook_fn(module, input, output):
        activations_buffer.append(output.detach())

    # Register hook on the target layer
    if component == "mlp":
        handle = model.transformer.h[layer_idx].mlp.register_forward_hook(hook_fn)
    else:
        handle = model.transformer.h[layer_idx].attn.register_forward_hook(hook_fn)

    # Training loop (simplified)
    for step in range(num_tokens // batch_size):
        # Get a batch of activations
        batch_acts = get_activation_batch(
            model, tokenizer, batch_size
        )

        loss, metrics = sae.loss(batch_acts)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Normalize decoder weights after each step
        with torch.no_grad():
            sae.decoder.weight.data = F.normalize(
                sae.decoder.weight.data, dim=0
            )

        if step % 100 == 0:
            print(
                f"Step {step}: recon={metrics['reconstruction']:.4f}, "
                f"l1={metrics['sparsity']:.4f}, "
                f"alive={metrics['alive_features']}/{d_sae}"
            )

    handle.remove()
    return sae

4. Activation Patching

Activation patching (also called causal tracing or interchange intervention) is a technique for identifying which components of a model are responsible for a specific behavior. The idea is to run the model on two inputs (a "clean" input that triggers the behavior and a "corrupted" input that does not), and then selectively replace activations from one run with activations from the other. If replacing a specific component's activation restores the original behavior, that component is causally important.

# Activation Patching with TransformerLens
import transformer_lens
from transformer_lens import HookedTransformer, utils
import torch

# Load model with TransformerLens
model = HookedTransformer.from_pretrained("gpt2-small")

# Example: Which components know that "The Eiffel Tower is in" -> "Paris"?
clean_prompt = "The Eiffel Tower is in"
corrupted_prompt = "The Colosseum is in"  # different answer: "Rome"

# Get clean logits for the target token
clean_logits, clean_cache = model.run_with_cache(clean_prompt)
target_token = model.to_single_token(" Paris")
clean_logit_diff = (
    clean_logits[0, -1, target_token]
    - clean_logits[0, -1, model.to_single_token(" Rome")]
).item()

# Get corrupted cache
_, corrupted_cache = model.run_with_cache(corrupted_prompt)

# Patch each attention head and measure the effect
results = torch.zeros(model.cfg.n_layers, model.cfg.n_heads)

for layer in range(model.cfg.n_layers):
    for head in range(model.cfg.n_heads):
        # Define hook that patches this head's output
        def patch_hook(activation, hook, layer=layer, head=head):
            # Replace corrupted activation with clean activation
            activation[:, :, head, :] = clean_cache[
                hook.name
            ][:, :, head, :]
            return activation

        # Run corrupted input with this one head patched
        hook_name = f"blocks.{layer}.attn.hook_z"
        patched_logits = model.run_with_hooks(
            corrupted_prompt,
            fwd_hooks=[(hook_name, patch_hook)],
        )

        # Measure recovery of clean behavior
        patched_diff = (
            patched_logits[0, -1, target_token]
            - patched_logits[0, -1, model.to_single_token(" Rome")]
        ).item()

        results[layer, head] = patched_diff

print("Top 5 most important heads for 'Eiffel Tower -> Paris':")
flat = results.flatten()
top_indices = flat.argsort(descending=True)[:5]
for idx in top_indices:
    layer = idx // model.cfg.n_heads
    head = idx % model.cfg.n_heads
    print(f"  L{layer}H{head}: logit diff = {flat[idx]:.3f}")

5. TransformerLens and nnsight

Two primary libraries support mechanistic interpretability research on transformers. TransformerLens provides a re-implementation of common models with full hook access at every computational step. nnsight offers a different approach, allowing intervention on any PyTorch model without re-implementation.

# nnsight: Intervening on any PyTorch model
from nnsight import LanguageModel

# Wrap any HuggingFace model
model = LanguageModel("gpt2", device_map="auto")

# Use the tracing context to inspect and modify activations
with model.trace("The cat sat on the") as tracer:
    # Access any module's output
    layer_5_output = model.transformer.h[5].output[0]

    # Save it for inspection
    layer_5_output.save()

    # You can also modify activations in-place
    # model.transformer.h[5].mlp.output[:] *= 0  # ablate MLP

# Access saved activations after the trace
print(f"Layer 5 output shape: {layer_5_output.value.shape}")

# Activation patching with nnsight
clean_text = "The Eiffel Tower is in"
corrupt_text = "The Colosseum is in"

# Get clean activations
with model.trace(clean_text) as tracer:
    clean_resid = model.transformer.h[8].output[0].save()

# Patch corrupted run with clean activations at layer 8
with model.trace(corrupt_text) as tracer:
    model.transformer.h[8].output[0][:] = clean_resid.value
    patched_logits = model.lm_head.output.save()

print(f"Patched prediction: {patched_logits.value[0, -1].argmax()}")