We would like to understand what trained transformers represent internally. Is it possible to determine what a transformer must represent at specific stages of its architecture? In other words, does the architecture impose constraints on what can and cannot be represented? In this document, I explore this possibility by identifying a constraint that the softmax imposes on the logits to determine next-token probabilities, and pulling this constraint "backward" through each computation (e.g., the unembedding, the RMSNorm, and the MLP). The hope is that the pullback to a prior stage constrains the set of possible objects that the activations at the stage could represent, of which we can then confirm or deny empirically.
Below is the architecture at the final layer and at the "readout stack." The final layer converts resid_pre into resid_mid through attention; afterward, resid_mid is converted into resid_post through the MLP. Resid_post is then converted into resid_final through a final RMSNorm, then into the logits through the unembedding, and finally the next-token probabilities (NTP) through the softmax.
(The text that follows was initially written by Claude and heavily, heavily edited by me.)
After training, the model's NTP (written as $p$) is close to the correct Bayes-optimal one. The NTP are computed by softmaxing the logits:
$$p = \mathrm{softmax}(\ell), \qquad p_i = \frac{e^{\ell_i}}{\sum_j e^{\ell_j}}.$$where $V$ is the vocabulary size, and $i$ and $j$ index $V$.
What does knowing $p$ force about the logits $\ell$? The goal of this section is a decomposition of logit space: one direction that cannot matter for the NTP at all, and a $(V-1)$-dimensional subspace on which the logits must be the centered log-NTP, an object 2.2 defines.
Here $\log Z$ is a single number, shared by all $V$ logits. Knowing $p$ therefore determines what each logit must be up to an unknown constant $c$:
$$\ell = \log p + c\,\mathbf{1}, \qquad c \text{ unknown}, \quad \mathbf{1}=(1,\dots,1).$$Could the logits prefer some particular constant $c$? We can check that no value is preferred by shifting all the logits by the constant $c$:
$$\frac{e^{\ell_i + c}}{\sum_j e^{\ell_j + c}} = \frac{e^{c}\, e^{\ell_i}}{e^{c} \sum_j e^{\ell_j}} = \frac{e^{\ell_i}}{\sum_j e^{\ell_j}} = p_i.$$The factor $e^c$ cancels top and bottom, and NTP is unchanged. So the NTP alone cannot distinguish whether the logits should have one $c$ or another. But perhaps the loss selects a particular $c$ even though the NTP do not. We can check this by computing the loss with and without shifting the logits: $L$ and $L_{\text{shifted}}$.
The loss at one position is the cross-entropy, $L = -\log p_y$, where $y$ is the correct next-token. Using $\log p_y = \ell_y - \log Z$ from 2.1:
$$L = -\ell_y + \log Z.$$Now, we add the same constant $c$ to every logit, substitute the shifted logits into the loss, and simplify:
$$\begin{aligned} L_{\text{shifted}} &= -(\ell_y + c) + \log \sum_j e^{\ell_j + c} && [\text{substitute } \ell_i \to \ell_i + c]\\[2pt] &= -\ell_y - c + \log\Big(e^{c}\sum_j e^{\ell_j}\Big) && [\text{factor } e^c \text{ out of the sum}]\\[2pt] &= -\ell_y - c + c + \log \sum_j e^{\ell_j} && [\log \text{ of a product}]\\[2pt] &= -\ell_y + \log Z \;=\; L. && [\text{the } c\text{'s cancel; definition of } Z] \end{aligned}$$In brief, because the $-c$ from the first term cancels the $+c$ from the second, the loss comes out as the same number for every choice of $c$. A quantity that does not change the loss does not receive a gradient, so training exerts no pressure on $c$. In other words, moving the whole logit vector along $\mathbf 1$ does not change the NTP or the loss.
The claim so far is that the constant-shift part of the logits is not determined, and the constant can differ at every token position. The model is free to represent anything along $\mathbf 1$, as long as it represents log-NTP in the rest (because $\ell_i = \log p_i + \log Z$). The representation that the logits must carry is therefore log-NTP, up to a constant shift.
To see this freedom concretely, I use a Mess3 process, whose vocabulary has $V=3$ tokens so that logit space is three-dimensional and can be drawn directly. Each point below is the logit vector at one token position, one for every history of the process. Each panel makes a different choice of the constants, so the four geometries look different, yet they differ only along $\mathbf 1$ and therefore produce exactly the same NTPs at every position.
How many dimensions does log-NTP have? $p$ has $V$ entries and one constraint, $p_1 + \cdots + p_V = 1$, so it has $V-1$ degrees of freedom. Taking the log is a one-to-one relabeling of each entry, which neither adds nor removes degrees of freedom. So log-NTP is conceptually $(V-1)$-dimensional because the NTP has $V-1$ dimensions. (The orange panel above does show a curved raw log-NTP curving, but this is not an extra dimension.) This conclusion can also be arrived at by recognizing that the subspace that matters for the NTP must be orthogonal to the $\mathbf 1$ direction, which is a 1 dimensional subspace. Therefore, this subspace must be $V-1$. We call this the predictive subspace.
The predictive subspace is the set of vectors orthogonal to $\mathbf 1$. As a linear subspace, its vectors pass through the origin. A basis for it is any $V-1$ independent directions orthogonal to $\mathbf 1$. The directions themselves are not determined, since rotating a basis within the $V-1$ subspace gives an equally valid basis. This is fine! The well-defined objects are the plane itself and the projection onto it, and the projection needs no basis because, critically, removing the $\mathbf 1$-component of a vector is just subtracting its mean, $v \mapsto v - \bar v\,\mathbf 1$.
The predictive subspace does not contain the literal log-NTP. Its entries are all negative and their mean varies across positions, which is the raw log-NTP panel of the figure. To land the content in the plane, we project it by centering it. The result is the centered log-NTP. Its entries sum to zero by construction:
$$\sum_i \big(\log p_i - \overline{\log p}\big) \;=\; \sum_i \log p_i \;-\; V\,\overline{\log p} \;=\; 0.$$So the centered log-NTP is the representative of "log-NTP up to a constant shift" that lives in the predictive subspace, which forms the blue geometry above in 2.1.
The centered log-NTP are defined as:
$$\begin{aligned} \bar\ell &= \overline{\log p} + c && [\text{average both sides}]\\[2pt] \ell_i - \bar\ell &= \log p_i - \overline{\log p}. && [\text{subtract; } c \text{ cancels}] \end{aligned}$$To compute the logits, resid_final $\in\mathbb{R}^{d}$ is passed through the unembedding, which is a linear map and a bias:
$$\ell = \text{resid_final}\;W_U + b_U, \qquad \ell_i = \text{resid_final}\cdot u_i + b_i,$$where $u_i\in\mathbb{R}^{d}$ is the $i$-th column of $W_U$, $b_i$ is the $i$-th entry of the bias, and $d$ is the dimension of the residual stream.
Section 2 forced the content of the logits on the predictive subspace. resid_final sits one step earlier, so we would like the analogous statement one step back: what subspace of resid_final does that content depend on, and what must resid_final equal along them? To answer both questions, we rewrite the conclusion of 2.2 as:
$$\begin{aligned} \ell_i - \bar\ell &= \text{resid_final}\cdot u_i + b_i - \frac{1}{V}\sum_j \big(\text{resid_final}\cdot u_j + b_j\big) && [\text{substitute}]\\[2pt] &= \text{resid_final}\cdot\Big(u_i - \bar u\Big) + \big(b_i - \bar b\big), \qquad \bar u = \frac{1}{V}\sum_j u_j, \quad \bar b = \frac{1}{V}\sum_j b_j. && [\text{dot products are linear}] \end{aligned}$$The simplified equation says that the predictive subspace obtains its content from resid_final by reading it along the centered unembedding directions $u_i - \bar u$, plus a constant. So the content of the predictive subspace is determined by the part of resid_final along the centered unembedding directions. Their span is the subspace of resid_final that determines the prediction, and we could call it the unembedding-predictive subspace.
Any movement of resid_final orthogonal to the unembedding-predictive subspace does not change the dot products in the equation above. Therefore, the NTP cannot change. This is analogous to the $\mathbf 1$ direction from 2.1, except that the free subspace here has $d-(V-1)$ dimensions, where $d$ greatly exceeds $V$.
How large is the unembedding-predictive subspace? The centered unembedding directions sum to 0, $\sum_i (u_i - \bar u) = \sum_i u_i - V\bar u = 0$, so they are linearly dependent and span at most $V-1$ dimensions. The dimension equals $V-1$ exactly when this is their only linear relation, which requires $d \ge V-1$. One difference from the predictive subspace of 2.2 is worth noting. There, no preferred basis existed and none was needed. Here, the architecture hands us specific spanning directions, the $u_i - \bar u$, read directly off the model's own trained weights: no probe is fitted and no new parameters are introduced.
What representation must live in this subspace? For each centered unembedding direction, its dot product with resid_final is forced to equal the $i$-th entry of the centered log-NTP, minus the fixed number $b_i - \bar b$ from the equation above. This bias correction translates the whole picture by one fixed vector and does nothing else, and its entries are of order $10^{-2}$ in our models, so I keep it implicit in what follows. The $V$ centered directions carry one redundancy (they sum to 0), so keep any $V-1$ of them; the dropped one adds nothing, since its reading is minus the sum of the others. If these $V-1$ directions were perpendicular unit axes (an orthonormal basis of the subspace), dot products would be coordinates, and resid_final would contain the blue, centered log-NTP geometry of 2.1 exactly, with its lengths and angles. But they are not necessarily; they can be directions of arbitrary lengths and angles.
However, any independent set of directions is just perpendicular unit axes pushed through one fixed invertible linear map. I will call such a map a "shear" from here on, loosely, and this can include rotations, expansions, and stretches, like the panels of the figure below. Reading the same forced numbers off sheared axes "draws" the same picture, sheared. So the geometry of resid_final in this subspace must be the blue geometry of 2.1 up to a shear, and its representation is a sheared, centered log-NTP.
To compute resid_final, the final layer's resid_post is passed through the final RMSNorm:
$$\text{resid_final} = \gamma \odot \frac{\text{resid_post}}{\text{RMS}(\text{resid_post})}, \qquad \text{RMS}(\text{resid_post}) = \sqrt{\tfrac{1}{d}\textstyle\sum_k \text{resid_post}_k^2} \qquad (\odot \text{ means elementwise product}).$$where $\gamma$ is a learned vector that elementwise scales, and RMS is the L2 norm of a vector divided by $\sqrt{d}$.
We can perform a similar substitution as the previous section, continuing the same chain:
$$\begin{aligned} \ell_i - \bar\ell &= \text{resid_final}\cdot(u_i - \bar u) + \big(b_i - \bar b\big) && [\text{conclusion of Section 3}]\\[2pt] &= \left(\gamma \odot \frac{\text{resid_post}}{\text{RMS}(\text{resid_post})}\right)\cdot(u_i - \bar u) + \big(b_i - \bar b\big) && [\text{substitute the RMSNorm}]\\[2pt] &= \frac{\big(\gamma\odot(u_i-\bar u)\big)\cdot \text{resid_post}}{\text{RMS}(\text{resid_post})} + \big(b_i - \bar b\big). && [\text{regroup: move } \gamma \text{ onto the direction}] \end{aligned}$$The simplified equation says that the predictive content is obtained from resid_post by reading it along the directions $\gamma\odot(u_i-\bar u)$, divided by one per-position number, plus a constant. These directions are the centered unembedding directions bent elementwise by $\gamma$. The bend adds a second shear: in Section 3, reading along the skewed unembedding directions already sheared the picture once, and $\gamma$ bends those directions further. Both distortions are fixed, and they combine into one fixed linear map, so the representation read here is a twice-sheared centered log-NTP, but importantly, scaled at each position by that position's RMS. We could call the span of these directions the gamma-unembedding-predictive subspace. The name is getting long and kind of pretentious...
The RMS-scaling per token position complicates the geometry that we will see. Across token positions, the magnitude of resid_post can differ, so its RMS can differ. In other words, RMS is "idiosyncratic" to token positions, and it cannot be abstracted into a single global shear. This implies that the gamma-unembedding-predictive subspace will not have literally a twice-sheared centered log-NTP, but this geometry whose individual points have varying magnitudes. To be precise, all $V$ dot products $\big(\gamma\odot(u_i-\bar u)\big)\cdot\text{resid_post}$ at a token position (of which $V-1$ are independent, since the directions sum to 0) are divided by an RMS number. Dividing a vector by a number changes its length, not its direction. So at each position, the direction remains the same as the twice-sheared centered log-NTP -- but its length can change. (Dividing each vector in the subspace by its token position's RMS removes the stretches and gives back the twice-sheared centered log-NTP geometry.)
To say precisely which parts of resid_post the constraint touches, split resid_post into its component inside the gamma-unembedding-predictive subspace and its component in the orthogonal complement:
$$\text{resid_post} = p + q, \qquad p \in \text{span}\{\gamma\odot(u_i-\bar u)\}, \qquad q \perp \text{span}\{\gamma\odot(u_i-\bar u)\}.$$Every dot product in the chain reads only $p$. Showing each step, starting from this section's conclusion:
$$\begin{aligned} \ell_i - \bar\ell &= \frac{\big(\gamma\odot(u_i-\bar u)\big)\cdot \text{resid_post}}{\text{RMS}(\text{resid_post})} + \big(b_i - \bar b\big) && [\text{conclusion of this section}]\\[2pt] &= \frac{\big(\gamma\odot(u_i-\bar u)\big)\cdot (p+q)}{\text{RMS}(\text{resid_post})} + \big(b_i - \bar b\big) && [\text{substitute } \text{resid_post} = p+q]\\[2pt] &= \frac{\big(\gamma\odot(u_i-\bar u)\big)\cdot p \,+\, \big(\gamma\odot(u_i-\bar u)\big)\cdot q}{\text{RMS}(\text{resid_post})} + \big(b_i - \bar b\big) && [\text{dot product distributes over the sum}]\\[2pt] &= \frac{\big(\gamma\odot(u_i-\bar u)\big)\cdot p}{\text{RMS}(\text{resid_post})} + \big(b_i - \bar b\big) && [\big(\gamma\odot(u_i-\bar u)\big)\cdot q = 0:\ q \text{ is orthogonal to the subspace}] \end{aligned}$$Next, rewrite the denominator in terms of $p$ and $q$:
$$\begin{aligned} \text{RMS}(\text{resid_post}) &= \frac{\lVert \text{resid_post}\rVert}{\sqrt d} && [\text{definition of RMS}]\\[4pt] \lVert \text{resid_post}\rVert^2 &= \lVert p+q\rVert^2 = \lVert p\rVert^2 + 2\,p\cdot q + \lVert q\rVert^2 && [\text{expand the square}]\\[2pt] &= \lVert p\rVert^2 + \lVert q\rVert^2 && [p\cdot q = 0]\\[4pt] \text{so}\quad \text{RMS}(\text{resid_post}) &= \sqrt{\big(\lVert p\rVert^2 + \lVert q\rVert^2\big)/d}. && \end{aligned}$$Substituting this denominator back into the chain:
$$\ell_i - \bar\ell = \frac{\big(\gamma\odot(u_i-\bar u)\big)\cdot p}{\sqrt{\big(\lVert p\rVert^2 + \lVert q\rVert^2\big)/d}} + \big(b_i-\bar b\big).$$The figure below shows this using the running Mess3 example. Each blue point is one token position's $p$, assuming no twice-shear at all -- so the in-subspace geometry is just the centered log-NTP in this hypothetical. Each red point is the same position's $q$. The constraint never sees $q$'s direction, so the directions drawn are arbitrary; only $q$'s magnitude enters, through the denominator above. At the slider's zero position, the whole resid_post has the same length at every position -- resid_post lies on a shell -- so $q$'s length is $\sqrt{R^2-\lVert p\rVert^2}$, largest exactly where the answer content is smallest. Dragging the slider assigns each position its own random factor, applied to $p$ and $q$ jointly -- each blue point slides along its ray through the origin, and its red point slides in lockstep, so the pinned ratio between them never changes. The directions do not move. In a trained model these common factors are not random, of course -- they are the actual overall scales of resid_post -- but every cloud along the slider is permitted by the constraint. Basically, sliding back to zero is what the final RMSNorm does.
In contrast to Section 3's unembedding-predictive subspace, this is a slight weakening in what must be represented at this subspace, as the magnitude of the centered log-NTP is stored partly in the full resid_post's magnitude. Therefore, whatever writes this content into the subspace, the MLP included, is not required to write the centered log-NTP. It is required to write the RMS-scaled twice-sheared centered log-NTP, which is Section 5's subject. This all means that the model must encode the ingredients to form the centered log-NTP geometry -- the directions and the per-point scales (RMS) to set magnitudes. In addition, it seems like this prediction does not enforce that beliefs must be represented at resid_post! I try to connect this with the constrained belief paper in Section 8.
Some additional notes: Movement of resid_post orthogonal to the gamma-unembedding-predictive subspace leaves every dot product in the equation unchanged. Unlike Section 3, however, orthogonal movement is not perfectly free from affecting NTP, as doing so can change the magnitude of resid_post, which changes $\text{RMS}(\text{resid_post})$.
Additionally, the spanning directions again sum to 0, $\sum_i \gamma\odot(u_i-\bar u) = \gamma\odot\sum_i(u_i-\bar u) = 0$, so the gamma-unembedding-predictive subspace has at most $V-1$ dimensions, and exactly $V-1$ under the same conditions as Section 3, provided $\gamma$ has no zero entries (so that the elementwise bend is invertible). As before, the directions are read directly off the model's trained weights, with no fitted probe.
resid_post is computed as a sum. It is the upper ⊕ in the diagram:
$$\text{resid_post} = \text{resid_mid} + \text{mlp_out}.$$This stage introduces a new kind of freedom, and it is the main idea of this section. Section 4 forced the content of resid_post -- the whole sum. However, forcing the contents of a sum does not force the contents of its summands. resid_mid can carry any amount of the required content, as long as mlp_out supplies the rest. So from this stage of the architecture on, no single module's output is forced.
To keep equations short, let $g_i = \gamma\odot(u_i-\bar u)$ for the spanning directions of the gamma-unembedding-predictive subspace. Starting from $\ell_i - \bar\ell$ and substituting the sum, continuing the chain:
$$\begin{aligned} \ell_i - \bar\ell &= \frac{g_i\cdot \text{resid_post}}{\text{RMS}(\text{resid_post})} + \big(b_i - \bar b\big) && [\text{conclusion of Section 4}]\\[2pt] &= \frac{g_i\cdot \text{resid_mid} + g_i\cdot \text{mlp_out}}{\text{RMS}(\text{resid_post})} + \big(b_i - \bar b\big) && [\text{substitute the sum; dot products distribute}] \end{aligned}$$mlp_out is computed by passing mlp_in through the MLP. mlp_in is an RMSNorm'ed resid_mid. This RMSNorm is different from the RMSNorm of the model's readout stack because it learns a different gamma vector:
$$\text{mlp_in} \;=\; \gamma'\odot \frac{\text{resid_mid}}{\text{RMS}(\text{resid_mid})}, \qquad \text{mlp_out} \;=\; \mathrm{MLP}(\text{mlp_in}) \;=\; W_{\text{out}}\,\mathrm{gelu}\big(W_{\text{in}}\,\text{mlp_in} + b_{\text{in}}\big) + b_{\text{out}},$$where $\gamma'$ is this RMSNorm's gamma.
Starting from $\ell_i - \bar\ell$ and substituting, continuing the chain:
$$\begin{aligned} \ell_i - \bar\ell &= \frac{g_i\cdot \text{resid_mid} \,+\, g_i\cdot \text{mlp_out}}{\text{RMS}(\text{resid_post})} + \big(b_i - \bar b\big) && [\text{conclusion of Section 5}]\\[2pt] &= \frac{g_i\cdot \text{resid_mid} \,+\, g_i\cdot \mathrm{MLP}(\text{mlp_in})}{\text{RMS}(\text{resid_post})} + \big(b_i - \bar b\big) && [\text{substitute the MLP}]\\[2pt] &= \frac{g_i\cdot \text{resid_mid} \,+\, \big(W_{\text{out}}^\top g_i\big)\cdot \mathrm{gelu}\big(W_{\text{in}}\,\text{mlp_in} + b_{\text{in}}\big) \,+\, g_i\cdot b_{\text{out}}}{\text{RMS}(\text{resid_post})} + \big(b_i - \bar b\big). && [\text{expand the MLP; move } W_{\text{out}} \text{ onto the direction}] \end{aligned}$$Every previous section continued by rewriting its newest term as a linear map of the earlier object. Here, the newest term is a fixed direction dotted with $\mathrm{gelu}\big(W_{\text{in}}\,\text{mlp_in}+b_{\text{in}}\big)$, and the gelu is elementwise and nonlinear. These properties introduce a difficulty, which are described next.
Let $z = W_{\text{in}}\,\text{mlp_in} + b_{\text{in}}$ for the vector of pre-activations, right before gelu. This vector has one entry per neuron. gelu is a scalar function, applied to each entry separately: entry $z_j$ comes out as $\mathrm{gelu}(z_j)$, and no entry affects any other. The function is $\mathrm{gelu}(x) = x\,\Phi(x)$, where $\Phi$ is the standard normal CDF. For very negative inputs it outputs approximately $0$; for large positive inputs it outputs approximately $x$; in between, its graph bends smoothly from one behavior to the other:
$$\begin{array}{c|cccccc} x & -3 & -1 & 0 & 1 & 2 & 3\\\hline \mathrm{gelu}(x) & -0.004 & -0.159 & 0 & 0.841 & 1.954 & 2.996\\ \mathrm{gelu}'(x) & -0.01 & -0.08 & 0.50 & 1.08 & 1.09 & 1.01 \end{array}$$The second row of this table shows the slope at each point. A linear function has one slope at every point. gelu's graph is bent, so its slope is different at different points -- essentially $0$ on the far left, essentially $1$ on the far right, and everything in between near the origin.
Why do slopes matter? Suppose mlp_in changes by a very small vector $\delta$. Then $z$ changes by $W_{\text{in}}\,\delta$, and each entry of $\mathrm{gelu}(z)$ changes as: $\mathrm{gelu}\big(z_j + (W_{\text{in}}\,\delta)_j\big) \approx \mathrm{gelu}(z_j) + \mathrm{gelu}'(z_j)\,(W_{\text{in}}\,\delta)_j$. Stacking the entries across j, the change in $\mathrm{gelu}(z)$ is $\mathrm{gelu}'(z)\odot(W_{\text{in}}\,\delta)$, which means that the change in each entry by delta is scaled by that entry's current slope. This reveals that the change in gelu(z) is non-constant, and, specifically, depends on the slope at the particular value of z_j. Therefore, a predictive subspace in mlp_in, if it exists, is not constant across token positions because the pre-image of z depends on what slope gelu(z) takes on, which depends on what z is, and in general z will differ across token positions.
Nevertheless, a predictive subspace can be computed per token position. If we substitute this change of $\mathrm{gelu}(z)$ into the algebra in 6.1, at one fixed token position, then:
$$\begin{aligned} \text{change in } \big(W_{\text{out}}^\top g_i\big)\cdot \mathrm{gelu}(z) &= \big(W_{\text{out}}^\top g_i\big) \cdot \big(\mathrm{gelu}'(z)\odot (W_{\text{in}}\,\delta)\big) && [\text{each entry moves at its own current slope}]\\[2pt] &= \big(\mathrm{gelu}'(z)\odot W_{\text{out}}^\top g_i\big)\cdot \big(W_{\text{in}}\,\delta\big) && [\text{elementwise product moves across the dot}]\\[2pt] &= \Big(W_{\text{in}}^\top\big(\mathrm{gelu}'(z)\odot W_{\text{out}}^\top g_i\big)\Big)\cdot \delta. && [\text{regroup: move } W_{\text{in}} \text{ onto the direction}] \end{aligned}$$In the last line, $W_{\text{in}}$, $W_{\text{out}}$, and $g_i$ are linear maps, and at a fixed token position $z$ is a fixed vector, so $\mathrm{gelu}'(z)$ is a fixed list of numbers that is also linear. So altogether, $W_{\text{in}}^\top\big(\mathrm{gelu}'(z)\odot W_{\text{out}}^\top g_i\big)$ is a fixed direction of mlp_in, and across the tokens i, these directions span a predictive subspace in mlp_in. For the sake of verbosity, again in previous sections, the directions $(u_i - \bar u)$ and $\gamma\odot(u_i-\bar u)$ were linear, so they were the same directions at every token position, which is what allowed a single subspace to host every position's content at once, and is what all the point clouds in this document are pictures of.
To clarify, Section 4 also involved a per-position quantity, the RMS, and absorbed it without losing the subspace -- so a per-position quantity is not by itself fatal. It is worth comparing this with gelu carefully. RMS is a single number at each position, and it divides all the readings at once. Dividing a vector by one number changes its length and nothing else: $(1,1)/2 = (0.5, 0.5)$, and $(1,1)/10 = (0.1, 0.1)$, and both of these lie on the same line through the origin as $(1,1)$. A subspace is a span -- a collection of lines through the origin -- and lines do not record length. So an operation that only changes lengths leaves every line where it was, and the subspace survives untouched. That is why Section 4's per-position RMS cost only a per-position scale on the contents, while the subspace itself remained one fixed object shared by all positions. The slope reweighting is a different kind of operation: not one number for the whole vector, but one number per entry. Multiplying the entries of a vector by different numbers changes which line it lies on. For example, take the direction $(1,1)$ over two entries. At a position where the slopes are $(1, 0)$, it becomes $(1\cdot 1,\ 0\cdot 1) = (1, 0)$: the horizontal line. At a position where the slopes are $(0, 1)$, it becomes $(0, 1)$: the vertical line. These are different lines, so $\mathrm{span}\{(1,0)\}$ and $\mathrm{span}\{(0,1)\}$ are different subspaces. One number times a vector can only stretch it; different numbers per entry can point it somewhere else.
Section 6.2 left us with a definite predictive subspace at every position and no single subspace. Before giving anything up, there is a natural repair to attempt. The set of values a slope can take is known in advance: $\mathrm{gelu}'$ ranges over roughly $[-0.13,\ 1.13]$. So, compute the subspace for every possible slope setting, and take the union of all of them -- the span of every subspace the formula can produce. Whatever direction the write could ever respond to, under any input whatsoever, would be inside this union, and its complement would be safe to ignore. Therefore, it would be an analytic object again: weights and the slope range, no data.
Unfortunately, this fails! Consider a simple case of two neurons, the write direction $(1,1)$, and $W_{\text{in}} = I$ so nothing else is in the way. Here the direction formula becomes $W_{\text{in}}^\top\big((s_1, s_2)\odot(1,1)\big) = (s_1, s_2)$. In words: pick slopes, and those slopes are themselves the direction. Every choice of slopes gives one line. When we sweep across slopes:
$$\begin{array}{ll} \text{slopes } (1,\,0) & \text{the horizontal line}\\ \text{slopes } (1,\,0.25) & \text{a line at } 14^\circ\\ \text{slopes } (1,\,0.5) & \text{a line at } 27^\circ\\ \text{slopes } (1,\,1) & \text{the diagonal}\\ \text{slopes } (0.5,\,1) & \text{a line at } 63^\circ\\ \text{slopes } (0,\,1) & \text{the vertical line} \end{array}$$The lines sweep continuously across the whole quadrant. Their union is the entire plane. Nothing was ruled out.
The full-size version fails the same way. Pick one neuron $j$. Give it slope $1$. Give every other neuron slope $0$. (This can happen: slopes are near $1$ for large pre-activations and near $0$ for very negative ones.) Now plug these slopes into the direction formula. The elementwise product kills every entry except entry $j$: $\mathrm{gelu}'(z)\odot W_{\text{out}}^\top g_i = \big(W_{\text{out}}^\top g_i\big)_j\, e_j$. Then $W_{\text{in}}^\top e_j$ is just row $j$ of $W_{\text{in}}$. So the direction is row $j$ of $W_{\text{in}}$, times a number. Conclusion: row $j$ of $W_{\text{in}}$ is in the union. This works for every neuron $j$ (as long as the write uses that neuron at all, and it essentially always does). So the union contains all $256$ rows of $W_{\text{in}}$. And $256$ vectors in a $64$-dimensional space span everything. The union is all of mlp_in space.
Therefore, no direction of mlp_in is ignorable a priori. The weights permit any direction to matter at some conceivable input. So there is no analytic single subspace at mlp_in -- not by the direct route of 6.2, and not by the union repair. This is where the derivation genuinely ends. Everything after this point is empirical.
It's possible that the context never actually touches all slopes of gelu, and therefore the implied union is not the entire space of mlp_in. We can define an empirical predictive subspace of mlp_in to be the span of the visited gelu slopes' implied subspaces, and call it the MLP-gamma-unembedding-predictive subspace. By construction, movement of mlp_in orthogonal to it does not change mlp_out.
Each visited subspace's directions are $W_{\text{in}}^\top\big(\mathrm{gelu}'(z_t)\odot W_{\text{out}}^\top g_i\big)$, one direction per logit $i$. For one token position $t$, we stack its $V$ directions as the rows of a matrix $J_t$ -- a $V\times d$ matrix. This is one token position's matrix; the subscript says which position, and every visited position has its own. Stacking all of the visited $J_t$ on top of each other gives one tall matrix, which we call $J$ -- a $(T\cdot V)\times d$ matrix, with $T$ being the number of visited positions. The union of the visited subspaces is the span of $J$'s rows. The per-position matrix $J_t$ happens to be the Jacobian of mlp_out with respect to mlp_in at position $t$, which can be shown with the following proof. The Jacobian is the matrix of partial derivatives $\partial\,\Big(\big(W_{\text{out}}^\top g_i\big)\cdot \mathrm{gelu}(z)\Big)\,/\,\partial\,(\text{mlp_in})_k$, evaluated at position $t$'s activation:
$$\begin{aligned} \frac{\partial\,\Big(\big(W_{\text{out}}^\top g_i\big)\cdot \mathrm{gelu}(z)\Big)}{\partial\,(\text{mlp_in})_k} &= \sum_j \big(W_{\text{out}}^\top g_i\big)_j\;\mathrm{gelu}'(z_j)\;\frac{\partial z_j}{\partial\,(\text{mlp_in})_k} && [\text{chain rule; gelu is elementwise, so no cross terms}]\\[2pt] &= \sum_j \big(W_{\text{out}}^\top g_i\big)_j\;\mathrm{gelu}'(z_j)\;(W_{\text{in}})_{jk} && [z = W_{\text{in}}\,\text{mlp_in} + b_{\text{in}} \text{ is linear: } \partial z_j/\partial(\text{mlp_in})_k = (W_{\text{in}})_{jk}]\\[2pt] &= \Big[\,W_{\text{in}}^\top\big(\mathrm{gelu}'(z)\odot W_{\text{out}}^\top g_i\big)\Big]_k. && [\text{the same sum, regrouped}] \end{aligned}$$An obvious next move is to take the simple union: take the span of $J$'s rows. This is slightly unprincipled because each position contributes $V-1$ dimensions. Different positions contribute different ones, because the gelu slopes change. Generic $(V{-}1)$-dimensional subspaces pile up: with $V-1=2$ and $d=64$, about $32$ positions already cover all of mlp_in space -- and there are many more possible token positions in a sequence. So the span of the stack is everything, once again.
The root problem is that a union is all-or-nothing per direction. It asks: was this direction used at least once, at least a little? With many positions, the answer is basically yes for every direction. The question that separates directions is: how much was each direction used? That needs a magnitude per direction, not a yes/no. The magnitude-aware version of the span is the SVD of $J$: the singular directions of $J$ are the union's directions, and each singular value says how much that direction was used. The subspace we want is the span of the top singular directions. In practice, though, we never build the tall matrix $J$. The matrix $M = J^\top J/T$ below has the same eigendirections as $J$'s singular directions, with eigenvalues equal to the squared singular values over $T$, and unlike $J$ it can be accumulated position by position in constant memory: read a batch, add its $J_t^\top J_t$, discard.
The span of the top eigendirections of the SVD of J that capture nearly all of the variance of J are the MLP-gamma-unembedding-predictive subspace.
| site | predictive subspace | dim | constraint type | content |
|---|---|---|---|---|
| logits | complement of $\mathbf 1$ | $V-1$ | content (exact) | centered log-NTP |
| resid_final | span$\{u_i-\bar u\}$ | $V-1$ | content up to shear | sheared, centered log-NTP |
| resid_post | span$\{\gamma\odot(u_i-\bar u)\}$ | $V-1$ | content up to two shears and up to magnitude | the directions of twice-sheared, centered log-NTP (magnitudes are encoded as resid_post's magnitude) |
| mlp_out | same subspace as in resid_post | $V-1$ | no longer forced | |
| mlp_in | top eigendirections of SVD(J) | ~15 | no longer forced |