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$.

Each residual block (or Transformer block) consists of an attention layer followed by an MLP layer. Both layers 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. This write operation is performed via a residual connection: $x_{i+1} = x_i + \text{layer}(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 skip connections in ResNets). At its core, the residual stream is simply the cumulative sum of the outputs from all preceding layers, added to the initial token embedding.

A defining feature of the Transformer’s residual stream is its linear and additive structure. This is a key difference from architectures like ResNet, where non-linear activations (e.g., ReLU) are often applied within the skip connection path. Each layer in a Transformer block reads its input from the stream via linear transformations. Similarly, it writes its output back to the stream, typically after another linear transformation.

One of the most profound consequences of the stream’s linear, additive nature is that it lacks a privileged basis ¹. This implies that we can rotate the vector space of the residual stream—and correspondingly rotate the matrices that interact with it—without changing the model’s computational behavior.

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. If the network further learns to isolate features, this might evolve into a sparser representation like [1.5, 0, 0, 0]. Now, this vector is perfectly aligned with the first basis vector, $e_1$. We can confidently say that Neuron 1 has learned to detect squares.
  • Similarly, “triangle” might activate Neuron 3, becoming [0, 0, 1.3, 0].
  • It can be observed that, due to the presence of the ReLU activation (or more generally, as originally defined, “occurs when some aspect of a model’s architecture encourages neural network features to align with basis dimensions”), features tend to align with individual neurons. Consequently, one can “confidently” ascribe responsibility to specific neurons — for example, designating a given neuron as detecting the presence of a square. This alignment makes the representation more interpretable, and such a coordinate system is referred to as a privileged basis.
  • In addition, ReLU contributes both non-linearity and sparsity to the activation vectors (since negative activations are zeroed out), which further reinforces this interpretability structure.

The residual stream, being purely linear, does not have this architectural pressure towards sparsity. Features can exist in any arbitrary direction.

2. Rotational Invariance

Because the residual stream is basis-free, we can apply such 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 remain invariant after applying these compensatory rotations to the weight matrices, we say the residual stream is rotationally invariant or basis-free.

III. Virtual Weights

The linearity of the residual stream has another powerful implication: we can analyze the interaction between any two layers by composing their weight matrices into a single “virtual weight.”

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 across multiple layers. The input to component $i$ is influenced by the output of component $j$ (where $j<i$) via the virtual weight $W_I^i W_O^j$. This allows us to think of the Transformer as a network where every layer directly communicates with every subsequent layer, mediated by these virtual weights.

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 high dimensionality allows different layers and attention heads to operate on distinct, often non-overlapping, subspaces.

In a multi-head attention layer, each head has a relatively small output dimension ($d_{head} = d_{model} / n_{heads}$, often 64). When these outputs are projected back into the residual stream, they are likely to occupy different subspaces. It’s possible for these subspaces to be nearly orthogonal (disjoint), allowing heads to write information without interfering with each other ².

Once information is added to the residual stream, it persists until it is explicitly modified or overwritten by a subsequent layer. From this perspective, the dimensionality of the residual stream, $d_{model}$, acts as the model’s communication bandwidth or working memory. Increasing $d_{model}$ theoretically increases the capacity for components to store and share information.

Furthermore, studies suggest that the token embedding and unembedding matrices ($W_E, W_U$) often interact with only a small fraction of the available dimensions ³. This leaves a large number of “free” dimensions in the residual stream for intermediate layers to use for computation.

However, this bandwidth is in very high demand. It is the sole channel for communication between all components. The computational dimensions of the components often far exceed the residual stream’s dimension. For instance, the MLP hidden layer dimension is typically $4 \times d_{model}$. This mismatch creates computational bottlenecks.

• 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.

  • By construction, each value vector has dimensionality $d_{model}/h$, where $h$ denotes the number of attention heads. Thus, its dimensionality is much smaller than that of the residual stream.
  • 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, the information in the residual stream $x_s$ is compressed into $v_s$ and subsequently transferred to the residual stream at position $t$. Thus, between two residual streams, the value vector functions as a bottleneck activation. Importantly, the value vector is the only mechanism by which information can be transmitted from one token to another.

Because of the high bandwidth demand imposed on the residual stream, certain MLP neurons or attention heads can be interpreted as performing memory management operations. For instance, they may clear specific residual dimensions allocated by earlier layers by reading out information and then writing back the negative of that information. This resembles the behavior of a memory-cleaning process: the act of writing the negation cancels out the previous signal, thereby freeing up representational capacity in the residual stream..

V. Attention Heads Operate as an Ensemble of Independent Operations

A key design principle of multi-head attention is that the heads, $h \in H$, operate in parallel and independently. The output of the attention layer is the sum of the outputs of individual heads.

Recall Attention Mechanism:

Let $r^{h_{k}}$ be the result vector from head $k$ (with dimension $d_{head}$). In the original Transformer paper, these results are concatenated and then projected by a single output matrix $W_O$ (with dimension $[n_{heads} \cdot d_{head} \times d_{model}]$). We can decompose this operation. Let $W_O$ be composed of sub-matrices $W_O^k$ (each of size $[d_{head} \times d_{model}]$), one for each head. The concatenation and projection is equivalent to:

This decomposition shows that the total output is simply the sum of each head’s output projected independently into the residual stream. Each head can be thought of as contributing its own update vector, and these are all added together.

VI. Attention Heads as Information Movers

The fundamental operation of an attention head is to move information between token positions. It reads information from the residual stream at one set of positions and writes that information to the residual stream at another position.

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

• Let $x$ be the matrix of input vectors from the residual stream (shape $[N \times d_{model}]$, where $N$ is sequence length).
• The head computes value vectors $V = x W_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, the output written to the stream is $H = R W_O = (AV)W_O$.

This sequence of operations—per-token projection, across-token mixing, per-token projection—can be elegantly expressed using the Kronecker product ($\otimes$).

VII. One-Layer Attention-Only Transformers

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