Residual Stream is Key to Transformer Interpretability

I. A High-Level Overview

A Transformer model processes information in a sequence. It begins with token embedding, where an input token $t$ (represented as a one-hot vector) is mapped to an embedding vector $x_0$ via an embedding matrix $W_E$. This vector then passes through a series of residual blocks. Finally, the output of the last block, $x_{L}$, undergoes token unembedding, where it is mapped to a vector of logits $L(t)$ via an unembedding matrix $W_U$.

In the pre-normalized decoder-only architecture used as our running model, each residual block (or Transformer block) consists of an attention layer followed by an MLP layer. Both sublayers read their input from the residual stream—the central pathway carrying vectors like $x_i, x_{i+1}, \dots$—and subsequently write their results back to it. Suppressing normalization for the moment, each write has the residual form $x_{i+1} = x_i + \operatorname{sublayer}(x_i)$.

II. The Residual Stream as a Communication Channel

If we conceptualize a Transformer as a complex computational device, the residual stream is one of its most critical components (analogous to the cell state in LSTMs or the identity skip path in a residual network). In a pre-normalized Transformer, it is the initial token embedding plus the cumulative sum of the updates written by preceding sublayers: a sublayer reads $\operatorname{Norm}(x_i)$ but still writes its update additively to $x_i$. A post-normalized block, which normalizes after the addition, does not admit this literal cumulative-sum unrolling across blocks.

A defining feature of the residual stream is its additive skip path. This is closely related to a ResNet: in the usual residual block, the skip path is the identity while nonlinearities live in the residual branch, although some variants also apply an activation after the addition. Transformer sublayers likewise contain nonlinear computation, even when their first and last learned maps are linear projections.

In the simplified architecture analyzed by Elhage et al.—where every operation that touches the residual stream is transformed consistently and there is no fixed coordinatewise normalization on that stream—the residual space has no privileged basis ¹. In that setting, we can change the basis of the residual stream and transform the matrices that interact with it without changing the model’s computational behavior. Ordinary LayerNorm, RMSNorm with learned per-coordinate gains, or other coordinate-dependent operations restrict this symmetry, so the exact claim should not be transferred to every Transformer unchanged.

1. Understanding the Privileged Basis

The concept of a “privileged basis” is defined in ¹ as follows:

“A privileged basis occurs when some aspect of a model’s architecture encourages neural network features to align with basis dimensions, for example because of a sparse activation function such as ReLU.”

Let’s dissect this with an example.

Consider a simple MLP trained to classify shapes. Let’s focus on a hidden layer $L$ with $N=4$ neurons, whose activation space is $\mathbb{R}^4$. The standard basis vectors are $e_1=[1,0,0,0], \dots, e_4=[0,0,0,1]$.

• Case 1: No Sparse Activation (e.g., Linear Layer)

  • The feature for “square” might be represented by the dense vector [1.2, -0.9, 0.8, 1.1].
  • The feature for “triangle” might be represented by [-0.8, -1.1, 1.3, -0.9].
  • Here, each feature is a complex combination of all four neurons. No single neuron is the “square detector.” The features are not aligned with the basis vectors. This is a non-privileged basis.

• Case 2: With a Sparse Activation (e.g., ReLU)

  • The pre-activation vector for “square,” [1.2, -0.9, 0.8, 1.1], becomes [1.2, 0, 0.8, 1.1] after passing through ReLU. Because ReLU is applied coordinate by coordinate, rotating the coordinates generally changes the function; the architecture therefore singles out the neuron basis. Training might also produce a sparser representation like [1.5, 0, 0, 0], but ReLU alone does not prove that this will happen or that Neuron 1 is a reliable square detector.
  • Similarly, “triangle” might activate Neuron 3, becoming [0, 0, 1.3, 0].
  • The justified architectural claim is that ReLU privileges the coordinate basis and can produce sparse activations by zeroing negative coordinates. Whether learned features actually align with individual neurons is an empirical question, and even an aligned activation does not by itself establish a causal or exclusive detector.

The bare additive skip path has no coordinatewise activation and therefore does not itself create this pressure. Features can occupy arbitrary directions, although normalization and other operations around the stream may still privilege or restrict particular coordinates.

2. Rotational Invariance

Under the simplified assumptions stated above, we can apply rotations without changing the model’s output. Let’s see why.

Let $R$ be an arbitrary orthogonal rotation matrix, meaning $R^{-1} = R^T$ and $R^T R = I$. Suppose we rotate a vector $x_i$ on the residual stream to get $x'_i = R x_i$. For the model’s behavior to remain unchanged, every component that interacts with the stream must adapt.

Consider an attention component:

• Read Operation: The component reads from the stream using matrices $W_Q, W_K, W_V$. To preserve the computation, we need new matrices $W'_Q, W'_K, W'_V$ such that:

Substituting $x'_i = Rx_i$, we get $W_Q x_i = W'_Q R x_i$. For this to hold for all $x_i$, we must have $W_Q = W'_Q R$, which implies $W'_Q = W_Q R^{-1} = W_Q R^T$. Thus, the new weight matrices simply “un-rotate” the input before applying the original transformation: $W'_Q x'_i = (W_Q R^T)(R x_i) = W_Q x_i$. The underlying logic is unchanged.

• Write Operation: The layer writes its output back to the rotated stream:

Since all vectors on the stream must be consistently rotated, $x'_{i+1} = R x_{i+1}$ and $x'_i = R x_i$. Substituting these into the original update rule $x_{i+1} = x_i + W_O \cdot \text{head\_output}$ gives:

This requires $W'_O = R W_O$. The output projection is simply rotated along with the rest of the space.

Since the internal calculations of each component in this simplified model remain invariant after applying these compensatory rotations to the weight matrices, we call its residual stream rotationally invariant or basis-free. A normalized production architecture needs a separate symmetry analysis; the calculation above is not a proof for it.

III. Virtual Weights

The additive structure of the residual stream has another useful implication: for a direct write from one component followed by a linear read from a later component, we can compose their projection matrices into a “virtual weight.” This describes that direct path, not the complete nonlinear computation between the components.

Let $C_j$ be the computation of component $j$ (e.g., an attention head or MLP), with input weights $W_I^j$ and output weights $W_O^j$. The update rule at step $j$ is:

Now, consider how the next component, $j+1$, reads from the stream:

The term $W_I^{j+1} W_O^j$ is a virtual weight matrix. It directly maps the output of component $j$ to the input of component $j+1$. This shows that information written by layer $j$ is read by layer $j+1$ through this composite matrix.

We can extend this bookkeeping across multiple layers. In the normalization-free model, the input projection of component $i$ receives the direct contribution of component $j$ (where $j<i$) through $W_I^i W_O^j$. This supports the useful picture of later components reading earlier writes through the shared stream, while nonlinearities, attention patterns, normalization, and intervening writes still determine the actual computation.

IV. Subspaces and Bandwidth of the Residual Stream

The residual stream is a high-dimensional vector space (e.g., $d_{model} = 768$ for BERT-base, $d_{model} = 2304$ for Gemma-2B). This dimensionality makes it possible for different layers and attention heads to use partly distinct subspaces, but the architecture does not require those subspaces to be disjoint or orthogonal.

In a multi-head attention layer, each head has a relatively small value dimension (often, but not always, $d_{head}=d_{model}/n_{heads}$). After its output projection, head $k$ can write only into the image of $W_O^k$, a subspace of dimension at most $d_{head}$. Different heads may learn partly separated or nearly orthogonal write subspaces, but this is an empirical property of a trained model rather than a guarantee of multi-head attention. Overlap can also be useful because it lets heads reinforce or cancel one another.

Because later updates are added, an earlier write remains as an algebraic summand of the residual stream. That does not guarantee that the information remains decodable: later writes, normalization, or cancellation can obscure or erase it. With this caveat, $d_{model}$ is a useful analogy for communication bandwidth or working memory, and a larger stream offers more representational degrees of freedom without giving a simple capacity formula.

The framework paper further notes that embedding and unembedding maps can leave directions that are weakly constrained by direct token input or logit output ². Calling these directions “free” is an intuition, not a guarantee that trained intermediate layers reserve or independently allocate them; the relevant ranks and usage must be measured in the model being analyzed.

This shared bandwidth can nevertheless constrain communication between sublayers. The computational dimensions of components can exceed the residual stream’s dimension; for instance, many MLP hidden layers are wider than $d_{model}$. This makes the residual stream a bottleneck in the dimensional sense defined below, although dimensionality alone does not show that useful information is lost.

• For example, the residual stream can be regarded as a form of bottleneck activation. MLP layers at different depths (which typically have activations of higher dimensionality than the residual stream) must communicate with one another through the residual stream. Consequently, the residual stream acts as an intermediary between two MLP layers whose hidden activations may have much larger dimensionality. Moreover, the residual stream is the only pathway through which any given MLP layer can communicate with subsequent layers. It must also carry forward information originating from other MLP layers along the path toward the extreme bottleneck.

• Similarly, a value vector (in the $Q, K, V$ decomposition of an attention head) also constitutes a bottleneck activation.

  • In the common equal-width parameterization, each value vector has dimensionality $d_v=d_{model}/h$, where $h$ denotes the number of attention heads. More generally, $d_v$ is an architectural choice.
  • Let $x_s$ denote the residual stream at token position $s$. The corresponding value vector is $v_s = x_s W_V$. This value vector $v_s$ is then used to update the residual stream at another position $t$: $$x_t = x_t + (\text{attention\_score} \times v_s) W_O .$$
  • In this way, information from residual state $x_s$ is projected into $v_s$ and contributes to the update at position $t$. When $d_v<d_{model}$, that projection is a dimensional bottleneck. Within the standard Transformer block covered here, attention’s weighted value path is the mechanism for cross-token communication; token-wise MLPs and normalizations do not themselves mix positions. Architectures with convolutions, recurrence, routing, or other mixing layers need not satisfy this claim.

In a particular trained model, an MLP neuron or attention head that writes the negative of an earlier component’s contribution can be interpreted as performing a memory-management-like operation. The algebra makes cancellation possible, but it does not imply that components generally allocate named dimensions or deliberately “free memory”; that interpretation requires causal evidence from the model.

V. Attention Heads Contribute Parallel Updates

A multi-head attention layer computes each head from the same input in parallel. Conditional on that input, one head’s forward calculation does not consume another head’s result, and the projected head outputs are summed. The heads are not statistically or functionally independent in general: they are trained jointly, share the residual input, and can interact through their summed write and through later layers.

Using row vectors for tokens, let $X\in\mathbb R^{N\times d}$ and let $X_i$ denote the row at token position $i$. For head $k$, define

Here $A^{(k)}_{ij}$ is the weight with which query position $i$ reads value position $j$; masked entries with $j>i$ are zero. Using $\ell$ for layer depth avoids overloading the token index $i$.

Let $r^{(k)}$ be the result row vector from head $k$, with dimension $d_v$. In the original Transformer paper, the $h$ head results are concatenated and then projected by $W_O\in\mathbb R^{hd_v\times d}$. Partition $W_O$ into row blocks $W_O^{(k)}\in\mathbb R^{d_v\times d}$, one for each head. Then

This decomposition shows that the total output is the sum of the separately projected head results. Each head contributes an update vector, but, as noted above, this algebraic decomposition does not imply that the learned head functions are independent.

VI. Attention Heads as Information Movers

Within the standard Transformer architecture covered here, an attention head is the component that moves information between token positions. It reads value information from residual states at source positions and contributes a weighted update at each query position.

To formalize this, let’s analyze the computation for a single head.

• Let $X$ be the matrix of input row vectors from the residual stream, with shape $N\times d_{model}$, where $N$ is sequence length.
• The head computes value vectors $V = X W_V\in\mathbb R^{N\times d_v}$. This is a per-token operation.
• It computes an attention matrix $A$ (shape $[N \times N]$), where $A_{ij}$ is the softmax score from query $i$ to key $j$.
• The result vectors are computed by mixing values: $R = A V$. This is an across-token operation, where the result for token $i$, $r_i$, is $\sum_j A_{ij} v_j$.
• Finally, with $W_O\in\mathbb R^{d_v\times d_{model}}$, the output written to the stream is $H = R W_O = AXW_VW_O\in\mathbb R^{N\times d_{model}}$.

This sequence of operations—per-token projection, across-token mixing, per-token projection—can be expressed using the Kronecker product ($\otimes$), following the notation in the framework paper ³.

We will stop at this general decomposition. To go further, we would need to study how $A(X)$ is formed in a one-layer attention-only Transformer; that is a question for another note.

Citation

1
@misc{ln2025residual,
2
    author={Nguyen Le},
3
    title={Residual Stream is Key to Transformer Interpretability},
4
    year={2025},
5
    url={https://lenguyen.vercel.app/note/math-transformers}
6
}

References