Begin with a river, not a pipeline

A pipeline suggests replacement. Stage one emits a representation; stage two consumes it and emits another. A transformer block is easier to understand as a river. The token embeddings enter as the upstream current. Attention adds a tributary whose content depends on relations among token positions. The feed-forward sublayer adds another tributary whose content depends on the current state at each position. The MLP is position-wise, but that state already contains context routed by attention. Empirical analysis of trained transformers has shown feed-forward updates promoting output-token distributions, which is one reason their residual writes deserve separate inspection.[7] The water downstream contains all three contributions, though later currents may mix or cancel them.

One state, three causal views

“Shared state” does not mean that every component sees the same object in the same way. A block exposes three related but distinct views. The read view is the normalised state presented to a sublayer. The write view is the vector that sublayer returns in residual-stream coordinates. The use view is the effect of that changed state on the frozen computation that follows. Treating these as one view is the source of many mistaken explanations.

A direction can persist on the skip path while being invisible to the next normalised read. A component can make a large write that lands in a direction to which the suffix is insensitive. A very small write can matter when the current state sits close to a decision boundary. Conversely, a write that looks decisive under the final unembedding can be cancelled before the model acts on it. Scale, readability and leverage are therefore different properties.

The smallest honest equation

For a common pre-normalised decoder block, let r_l be the matrix of token states entering layer l. It has one row per token position and one column per residual dimension. Let N be a normaliser, A the attention sublayer and M the feed-forward sublayer. The block computes:

Across a stack, the final pre-output state can be written as the embedding contribution plus every attention and feed-forward write observed during that forward pass:

This distinction matters. The sum tells us where to place readouts and interventions. It does not tell us that a capability belongs to one addend. A late MLP may respond to a feature routed by early attention. An attention head may suppress a direction written by another head. A final normaliser can rescale the whole mixture before unembedding. Additivity of state is not additivity of responsibility.

Addition commutes; conditional writing does not

Vector addition itself is commutative. If two already-computed writes are a and m, then r + a + m equals r + m + a. It is tempting to conclude that the attention and feed-forward branches could be reordered without changing the block. The hidden assumption is that the writes were computed independently of the evolving state.

They are not. In the usual sequential block, attention writes a = A(N(r)). The feed-forward network then writes m = M(N(r + a)). If the order is reversed, the MLP receives r, while attention receives a state already altered by the MLP. Both writes can change. Parallel block variants deliberately change this dependency structure; they are more than another drawing of the same sum.

flowchart TD
  A["The answer that appeared<br/>between two additions"] --> B["The representation nobody owns"]
  B --> C["Normalisation decides what a<br/>write can see"]
  C --> D["A visible write is not a<br/>causal contribution"]
  D -. "re-test" .-> A
  classDef first fill:#0a1424,stroke:#5B5CF6,stroke-width:2px,color:#e6f3ff
  classDef step fill:#0a1424,stroke:#3f7da0,stroke-width:1px,color:#e6f3ff
  class A first
  class B,C,D step

Change one variable: move the gate

The original Transformer placed layer normalisation after each residual addition: LayerNorm(x + Sublayer(x)).[1] Many later language-model architectures moved normalisation before the sublayer. The forward functions still contain attention, an MLP, addition and normalisation, yet the optimisation problem changes because the identity route is no longer passed through a normaliser at every step.

What layer normalisation actually removes

For one token vector x with d features, LayerNorm subtracts the feature mean, divides by the feature standard deviation, then applies learned gain and bias. Ba, Kiros and Hinton introduced it as a per-case alternative to batch-dependent normalisation, with the same computation at training and test time.[3]

Normalisation creates equivalence classes

Ignoring the small stabilising constant, LayerNorm maps every vector of the form αx + c1, for positive α, to the same centred and standardised pattern before learned gain and bias. The sublayer cannot distinguish a uniform shift c1, and it treats a positive global rescaling as the same direction. RMSNorm identifies positive rescalings but keeps sensitivity to uniform shifts. These are not philosophical claims about what the model “knows”. They are exact or approximate symmetries of the read operation.

Pre-LN then creates a useful separation. The raw residual state, including components that the normaliser suppresses, continues along the identity route. The branch reads only the equivalence class selected by the normaliser. A fact can therefore be present in the stream yet unavailable to a particular branch under its read map. Whether it later affects output depends on later normalisers, projections and writes. Presence in state and addressability by a sublayer are not the same condition.

The gradient path exposes the same architectural distinction. For a simplified pre-LN update y = x + F(N(x)), the Jacobian is I + J_FJ_N. It contains an explicit identity term. For post-LN, y = N(x + F(x)), the Jacobian is J_N(I + J_F); the direct route is filtered by the normaliser’s Jacobian. Xiong and colleagues used this difference to analyse why pre-LN and post-LN transformers exhibit different gradient behaviour at initialisation.[4] The identity term does not guarantee good optimisation. It explains which route exists before data and training reshape it.

A small experiment makes the path visible

To expose the mechanism without claiming a new benchmark, we initialised matched 18-block synthetic networks. Each block contained two small GELU feed-forward operators standing in for attention and the MLP. One stack used pre-LN; the other used post-LN. We measured the median parameter-gradient norm at each depth across 30 random seeds under the same mean-squared-error objective.

The logit lens is a question, not an answer

A decoder-only language model usually applies a final normaliser and an unembedding matrix to the last residual state. If we reuse that output machinery at an earlier layer, we obtain an intermediate vocabulary distribution. This “logit lens” asks which tokens are readable from the current state under the final decoder. A tuned lens learns a separate affine translator at each layer and has been shown to improve reliability over the raw logit lens in tested models.[8]

The readout is useful because it turns an opaque vector into a prediction trajectory. Its assumption is strong: the intermediate representation is sufficiently aligned with the final decoder. Earlier layers may use a different basis or encode information for later nonlinear processing. Final normalisation also means that raw component projections do not always compose as a simple sum unless the normalisation effect is handled explicitly.

Direct projection is local accounting

Let d_p,q be the difference between the unembedding directions for tokens p and q. For a component write w, the quantity d_p,q^Tw is its direct contribution to the p-versus-q logit difference under a linear readout. This is a clean accounting identity when no final normaliser intervenes. With a final LayerNorm or RMSNorm, the exact output change depends on the whole residual state, so analysts either apply the full readout to counterfactual states or use a stated local approximation.

Even when computed exactly, direct projection answers only what the write points towards at that site. It does not include how the write changes a later attention pattern, opens an MLP feature, alters normalisation statistics or triggers a compensating write. For that reason, three numbers must remain separate:

No one scalar is “the contribution”. Each is an answer to a different counterfactual. Agreement strengthens a mechanism claim. Disagreement reveals interaction structure that a component leaderboard would hide.

Separate three questions

The central practical contribution of this paper is a Residual Contribution Audit. It separates three properties that are often collapsed:

Write asks whether a component changed the residual state in a specified direction or subspace. Survival asks whether that change remains available after later layers, rather than being rotated, overwritten or cancelled. Use asks whether changing that write changes the tested behaviour through the model’s own downstream computation.

A fourth diagnostic, recovery, asks whether restoring the clean write into a corrupted run recovers the clean result. Recovery is especially informative because it tests a counterfactual carrier: can this state difference transmit the relevant information when placed back into the original computation?

Intervene on matched runs

Causal tracing and activation patching compare a clean run with a matched corrupted run. An internal activation from one run is inserted into the other, and the effect on a defined metric is measured. Meng and colleagues used a causal intervention of this kind to identify states that mediated factual predictions in the GPT models they studied.[9]

Ablation and patching answer different questions. Ablation asks whether the current route is necessary under a chosen replacement, often zero or a resampled activation. Clean-to-corrupt patching asks whether a state is sufficient to restore some behaviour in the corrupted context. Neither establishes unique responsibility. Redundant routes can hide necessity. A large patch can carry collateral information. Off-manifold replacements can create damage that the endogenous system never encounters.

Build the rival generator before trusting the result

A causal-looking pattern can be produced by several mechanisms. The practical defence is to name the strongest rival before running the intervention, then add the control that would separate it from the preferred explanation.

This discipline changes a negative result as well. A failed patch does not mean the information was absent. It may mean the patch site was too broad, too narrow, mistimed or incompatible with the corrupted context. The result lowers one specified mechanism claim, not every possible account of the behaviour.

stateDiagram-v2
  [*] --> Claim
  state "Begin with a river, not a pipeline" as Claim
  state "Computation as a dye stream" as Test
  state "One state, three causal views" as Revision
  Claim --> Test: expose mechanism
  Test --> Decision: evidence survives
  Test --> Revision: rival survives
  Revision --> Claim: narrow or rebuild
  Decision --> [*]

Five hook points are enough to begin

For each block, record the incoming residual state, the attention write, the post-attention state, the feed-forward write and the outgoing state. Record normalised reads when the question concerns what a sublayer detected. Bind every trace to the exact model weights, input tokens, mask, normalisation configuration and decoding metric.

Do not log only attention weights. They show how value vectors were mixed, not the full vector written after value and output projections. Do not log only hidden states. Their differences conflate all writes between two sites. A trace becomes causally useful when it preserves both the state and the delta that changed it.

The smallest valid audit packet

A convincing residual experiment needs more than tensors. It needs a contract that makes the counterfactual reproducible. The smallest useful packet contains four linked records.

Two invariants make the packet auditable. The clean and corrupted runs must use the same frozen parameters and execution settings. The intervention must alter only the declared tensor sites. Hashes or deterministic checks can verify both. A positive result then permits a narrow statement about causal leverage for that metric and context. A negative result lowers that statement, but cannot establish global absence or lack of representation.

This packet is intentionally smaller than a full interpretability platform. It can be implemented with forward hooks, saved activations and a metric function. Its value lies in forcing the analyst to expose assumptions that prose otherwise hides.

A realistic configured example

The executable lab trains a one-block, one-task decoder on synthetic sequences. The last position contains a fixed query token; the target is the first token. Causal attention permits the query position to read earlier positions. The model is small enough to train on a CPU in seconds, yet it contains token and positional embeddings, pre-LN multi-head attention, a GELU MLP, residual additions, a final normaliser and tied unembedding.

On the fixed seed supplied below, the model reaches perfect batch accuracy on the synthetic task. For the held-out sequence [7, 3, 9, 5, 0], the target probability is near zero at the embedding readout, 0.86 after the attention write, and approximately 1 after the MLP. Attention ablation reduces target probability to about 0.007. MLP ablation leaves the correct top token with probability 0.86. Patching the clean attention write into a corrupted sequence restores the clean answer.

Those numbers establish only what the configured experiment tests: attention carries the first-token information to the query position in this trained model, and the MLP sharpens the readout. They do not establish a general role for attention or MLPs in language models.

Run the full Python experiment

Save the code as residual_stream_lab.py, install PyTorch, and run it with Python 3.10 or later. It trains only on generated integer sequences. No external dataset or model download is required.

from __future__ import annotations

import statistics
import torch
from torch import nn
import torch.nn.functional as F

# Reproducible, synthetic experiment.
torch.manual_seed(7)
torch.set_num_threads(1)

VOCAB = 12
SEQ = 5
D_MODEL = 32
HEADS = 4


class Block(nn.Module):
    def __init__(self, d_model: int = D_MODEL) -> None:
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(
            d_model, HEADS, batch_first=True
        )
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )

    def forward(
        self,
        x: torch.Tensor,
        *,
        ablate_attn: bool = False,
        ablate_mlp: bool = False,
        patch_attn: torch.Tensor | None = None,
        trace: dict[str, torch.Tensor] | None = None,
    ) -> torch.Tensor:
        r0 = x
        z = self.ln1(r0)
        causal_mask = torch.triu(
            torch.ones(SEQ, SEQ, dtype=torch.bool), diagonal=1
        )
        attn_write, _ = self.attn(
            z, z, z, attn_mask=causal_mask, need_weights=False
        )

        if ablate_attn:
            attn_write = torch.zeros_like(attn_write)
        if patch_attn is not None:
            attn_write = attn_write.clone()
            attn_write[:, -1, :] = patch_attn

        r1 = r0 + attn_write
        mlp_write = self.mlp(self.ln2(r1))
        if ablate_mlp:
            mlp_write = torch.zeros_like(mlp_write)
        r2 = r1 + mlp_write

        if trace is not None:
            trace.update(
                r0=r0.detach(),
                attn_write=attn_write.detach(),
                r1=r1.detach(),
                mlp_write=mlp_write.detach(),
                r2=r2.detach(),
            )
        return r2


class TinyCopyModel(nn.Module):
    """At the final query position, predict the first token."""

    def __init__(self) -> None:
        super().__init__()
        self.token_embedding = nn.Embedding(VOCAB, D_MODEL)
        self.position = nn.Parameter(torch.randn(SEQ, D_MODEL) * 0.02)
        self.block = Block()
        self.final_norm = nn.LayerNorm(D_MODEL)
        self.unembed = nn.Linear(D_MODEL, VOCAB, bias=False)
        self.unembed.weight = self.token_embedding.weight

    def forward(self, tokens: torch.Tensor, **block_kwargs) -> torch.Tensor:
        x = self.token_embedding(tokens) + self.position
        x = self.block(x, **block_kwargs)
        return self.unembed(self.final_norm(x[:, -1, :]))

    def readout(self, residual: torch.Tensor) -> torch.Tensor:
        return self.unembed(self.final_norm(residual[:, -1, :]))


def synthetic_batch(size: int = 128) -> tuple[torch.Tensor, torch.Tensor]:
    tokens = torch.randint(0, VOCAB, (size, SEQ))
    tokens[:, -1] = 0  # fixed query marker
    target = tokens[:, 0].clone()
    return tokens, target


def train_copy_model(steps: int = 700) -> TinyCopyModel:
    model = TinyCopyModel()
    optimiser = torch.optim.AdamW(model.parameters(), lr=2e-3)

    for step in range(steps):
        tokens, target = synthetic_batch()
        logits = model(tokens)
        loss = F.cross_entropy(logits, target)
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()

        if step in {0, 99, 299, steps - 1}:
            accuracy = (logits.argmax(-1) == target).float().mean().item()
            print(
                f"step={step + 1:3d} loss={loss.item():.4f} "
                f"batch_accuracy={accuracy:.3f}"
            )
    return model.eval()


def audit(model: TinyCopyModel) -> None:
    clean = torch.tensor([[7, 3, 9, 5, 0]])
    corrupt = torch.tensor([[2, 3, 9, 5, 0]])
    trace: dict[str, torch.Tensor] = {}

    with torch.no_grad():
        clean_logits = model(clean, trace=trace)
        corrupt_logits = model(corrupt)
        no_attention = model(clean, ablate_attn=True)
        no_mlp = model(clean, ablate_mlp=True)
        patched = model(
            corrupt,
            patch_attn=trace["attn_write"][:, -1, :],
        )

    print("\nLayer-wise readouts for target token 7")
    stages = {
        "embedding": trace["r0"],
        "after attention": trace["r1"],
        "after MLP": trace["r2"],
    }
    for name, residual in stages.items():
        logits = model.readout(residual)
        probability = logits.softmax(-1)[0, 7].item()
        print(
            f"{name:16s} p(target)={probability:.3f} "
            f"top={logits.argmax(-1).item()}"
        )

    print("\nInterventions")
    cases = {
        "clean": clean_logits,
        "attention ablated": no_attention,
        "MLP ablated": no_mlp,
        "corrupt": corrupt_logits,
        "corrupt + clean attention patch": patched,
    }
    for name, logits in cases.items():
        probability = logits.softmax(-1)[0, 7].item()
        print(
            f"{name:34s} p(target)={probability:.3f} "
            f"top={logits.argmax(-1).item()}"
        )

    print("\nWrite norms at the query position")
    print("attention", trace["attn_write"][0, -1].norm().item())
    print("MLP      ", trace["mlp_write"][0, -1].norm().item())


# Optional: reproduce Figure 5 with matched synthetic stacks.
class FFOperator(nn.Module):
    def __init__(self, d_model: int) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class PreLNBlock(nn.Module):
    def __init__(self, d_model: int) -> None:
        super().__init__()
        self.n1, self.n2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
        self.a, self.m = FFOperator(d_model), FFOperator(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.a(self.n1(x))
        return x + self.m(self.n2(x))


class PostLNBlock(nn.Module):
    def __init__(self, d_model: int) -> None:
        super().__init__()
        self.n1, self.n2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
        self.a, self.m = FFOperator(d_model), FFOperator(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.n1(x + self.a(x))
        return self.n2(x + self.m(x))


def gradient_profile(seeds: int = 30, depth: int = 18, d_model: int = 48):
    medians: dict[str, list[float]] = {}
    for name, block_type in {"pre": PreLNBlock, "post": PostLNBlock}.items():
        runs: list[list[float]] = []
        for seed in range(seeds):
            torch.manual_seed(seed)
            blocks = nn.ModuleList([block_type(d_model) for _ in range(depth)])
            output = nn.Linear(d_model, d_model)
            x = torch.randn(8, 6, d_model)
            target = torch.randn_like(x)
            y = x
            for block in blocks:
                y = block(y)
            loss = F.mse_loss(output(y), target)
            loss.backward()
            runs.append([
                statistics.mean(
                    p.grad.detach().norm().item()
                    for p in block.parameters()
                    if p.grad is not None
                )
                for block in blocks
            ])
        medians[name] = [
            statistics.median(run[layer] for run in runs)
            for layer in range(depth)
        ]
    print("\nMedian gradient profiles")
    for name, values in medians.items():
        print(name, [round(value, 5) for value in values])


if __name__ == "__main__":
    trained_model = train_copy_model()
    audit(trained_model)
    gradient_profile()

What it tests: layer-wise readout, attention and MLP ablation, clean-to-corrupt patch recovery, component write magnitude, and a matched pre-LN versus post-LN gradient illustration. What it cannot establish: a general decomposition of language-model cognition, semantic localisation, or unique causal responsibility.

The strongest alternative explanation

The residual-stream view can itself become too coarse. A component output is a vector that may superpose many features. One block may contain several circuits with opposing effects. Layer boundaries are engineering units, not guaranteed cognitive units. Gated residuals, mixture-of-experts MLPs, parallel attention-MLP layouts and normalisation-free variants change the exact contract.

The alternative is not to abandon the residual view. It is to treat it as a measurement coordinate. It tells us where state persists and where updates enter. Feature-level explanations still require finer decompositions, controlled interventions and stability across inputs. The residual stream is a causal workspace in the architectural sense, not a certificate that its axes correspond to human concepts.