Transformer Math

Method 1: Activation Patching (exact, slow)

Run the model on a clean input and a corrupted input. For each component , swap its activation from the clean run into the corrupted run and measure the change in output metric :

Large effect = component causally matters. Requires one forward pass per component — O(N) total.

Method 2: Attribution Patching (approximate, fast)

First-order Taylor approximation. The attribution of feature on output is:

Gradient tells you sensitivity; activation difference tells you magnitude. One forward + one backward pass covers all components. Accuracy is within ~90% of full patching on most circuits.

Method 3: Circuit Tracing (SAE + attribution combined)

Replace MLP and attention outputs with their SAE decompositions. Now run attribution patching over the feature graph instead of neurons. This gives a directed graph where nodes are named features and edges are causal attributions:

Prune edges below a threshold to get the sparse subgraph — the “circuit”. Each node is interpretable (SAE feature), each edge is a measured causal strength.

Cross-Layer Transcoders (CLTs) — the 2025 upgrade

Standard SAEs are trained per-layer. Ameisen et al. (2025) instead train Cross-Layer Transcoders: each CLT reads from the residual stream at its own layer but outputs contributions to all subsequent MLP layers. This lets a single feature represent a concept that persists and propagates across depth, rather than needing separate features at each layer.

Edge weight between source feature and target feature is:

where is the source feature's activation, is the Jacobian through attention and residual connections, and the sum runs over the decoder-to-encoder paths. Loss uses JumpReLU (not standard ReLU) plus a tanh-based sparsity penalty to reduce dead features.

PyTorch: Basic Activation Patching

import torch
from contextlib import contextmanager

@contextmanager
def patch_activation(model, layer_name: str, patch_value: torch.Tensor):
    """Context manager to swap one layer's output mid-forward-pass."""
    hooks = []
    def hook_fn(module, input, output):
        return patch_value   # replace with clean-run activation
    handle = dict(model.named_modules())[layer_name].register_forward_hook(hook_fn)
    hooks.append(handle)
    try:
        yield
    finally:
        for h in hooks:
            h.remove()

def activation_patching_score(model, clean_tokens, corrupt_tokens, layer_name,
                               clean_cache, metric_fn):
    """
    Measure how much layer_name causally matters for the metric.
    metric_fn(logits) -> scalar (e.g., logit diff between two tokens)
    """
    # Baseline: corrupted run
    with torch.no_grad():
        corrupt_logits = model(corrupt_tokens)
    baseline = metric_fn(corrupt_logits)

    # Patched: corrupted run but swap in the clean activation
    clean_act = clean_cache[layer_name]
    with torch.no_grad():
        with patch_activation(model, layer_name, clean_act):
            patched_logits = model(corrupt_tokens)
    patched = metric_fn(patched_logits)

    return (patched - baseline).item()  # positive = component helps