Residual Stream as a Dynamical System State

I recently had the chance to read Elhage et al.’s A Mathematical Framework for Transformer Circuits, the paper that gets cited in every mechanistic interpretability paper. As a control systems student, I found it interesting how much of a transformer layer maps directly onto control systems ideas, and inspired me to start this series, A Transformer as a Dynamical System.

The concept of residual stream is often used as a tool for mechanistic interpretability, where you try to reverse engineer a black box architecture. Here, I’m going to use this as an internal state in a dynamical system.

The residual stream as the object

Traditionally, transformer architecture is understood as a multi layer system, where an input goes into each block one after another until we get the output we need. The paper shows that the blocks do not have to be nested like this. The output of a transformer is a sum: the embedded vector, plus every attention head output, plus every MLP output. This running sum is what the paper calls the ‘residual stream’.

We can think of the difference like this:

$$y=f_3(f_2(f_1(x)))\text{vs.}y=x+g_1(x)+g_2(x+g_1(x))+g_3(\cdot )$$

These are the same function family. Any $f$ can be written as $x+g(x)$, so the residual form does not make anything new possible, but it can pull the identity out of every layer. Passing information through a layer now costs nothing by default, and each layer’s contribution is its own term, so you can measure it. In the nested version those terms exist too, but they are stuck inside the next layer’s input and you cannot separate them out.

Every layer reads from the residual stream and writes back into it:

$$X^{(\ell +1)}=X^{(\ell )}+{\mathrm{Attn}}_{\ell }(X^{(\ell )})+{\mathrm{MLP}}_{\ell }\left(X^{(\ell )}+{\mathrm{Attn}}_{\ell }(X^{(\ell )})\right)$$

Here $x_i$ is the vector for token position $i$, and $X^{(\ell )}\in {\mathbb{R}}^{N\times d_{\text{model}}}$ is those vectors stacked as rows. The rows only interact through attention. The MLP works on each row by itself. So the state is the whole matrix, not a single vector.

This has the form of a forward-Euler step $x_{k+1}=x_k+h{F}_k(x_k)$ with $h=1$ and layer-dependent $F_k$. It could be seen as a time varying, finite horizon discrete system, not yet an ODE. Depth is the time axis here. Sequence position is a separate axis, which is what a recurrent network would have called time, and I am not touching it in this series.

What this sets up

If the residual stream is a state, then the standard control questions apply.

Reachability: if I add a vector to $X^{(\ell )}$, which final states $X^{(L)}$ can I actually get to?

Input class: every current method adds a constant vector. Would a state-dependent input do better, and is the constant vector just a special case of that?

Which layer: steering works best in the middle blocks, empirically. In this framing the injection layer is where the control horizon starts, so ‘why the middle’ becomes a question about the dynamics from $\ell$ to $L$.

I hope to answer these questions as I continue on in this series.