---
date: '2025-09-12'
description: linear algebra applications for transformers
id: notes
modified: 2026-06-07 01:29:37 GMT-04:00
seealso:
- '[[thoughts/Transformers|Transformers]]'
- '[[thoughts/Attention|Attention]]'
- '[[thoughts/LLMs|LLMs]]'
- '[[lectures/412/tools|Tools]]'
slides: true
tags:
- workshop
- math/linalg
title: supplement to 0[dot]412
transclude:
title: false
created: '2025-09-12'
published: '2025-09-12'
pageLayout: default
slug: lectures/412/notes
permalink: https://aarnphm.xyz/lectures/412/notes.md
generator:
quartz: v4.6.0
hostedProvider: Cloudflare
baseUrl: aarnphm.xyz
full: https://aarnphm.xyz/llms-full.txt
---
## introduction
- structure: embed tokens, inject positions, run masked multi-head attention, apply position-wise MLPs, and close with the unembedding that converts residual coordinates into logits.
> \[!abstract\] readings
>
> residual-stream linearization and MLP interpretation. \[@elhage2021mathematical; @geva2021transformerfeedforwardlayerskeyvalue\]
>
> - residual stream: Anthropic’s decomposition of attention heads and unembedding matrices into analyzable paths to motivate each algebraic step. \[@elhage2021mathematical\]
> - MLP as KV memory: highlight how the ffns cache features, setting up the later mixture-of-experts discussion. \[@geva2021transformerfeedforwardlayerskeyvalue\]
## decoder-only stack as linear maps
Let $d_{model}$ denote the residual stream width and $T$ the sequence length. Represent the token sequence as a matrix $X_0\in\mathbb{R}^{T\times d_{model}}$ where each row is the residual vector at a position.
### embeddings and residual coordinates
- token embedding: $X_0 = S E$ where $S\in\mathbb{R}^{T\times |\mathcal{V}|}$ is a one-hot indicator matrix and $E\in\mathbb{R}^{|\mathcal{V}|\times d_{model}}$ is the learned embedding.
- learned absolute and sinusoidal encodings add a position vector to the residual stream. RoPE instead rotates query and key coordinates inside attention, so it is not an additive matrix $P$ in $X_0$. \[@vaswani2023attentionneed\]
- layer norm: treat pre-normalization as an affine map $\mathcal{N}(x)=D(x)(x-\mu(x))$ with diagonal scaling $D(x)$. In practice we approximate $D$ as invertible on the support of interest to reason about nearby linearizations.
- residual stream: every block writes back to $X$ through additive updates, so we model the stack as
$$
X_{\ell+1} = X_{\ell} + F_{\ell}(X_{\ell}), \quad \ell = 0,\dots,L-1
$$
where $F_{\ell}$ sums attention and mlp contributions.
**Definition (token embedding map).** The map $\varepsilon: \{1,\dots,|\mathcal{V}|\}\to\mathbb{R}^{d_{model}}$ satisfies $\varepsilon(i)=e_i^\top E$ with $e_i$ the $i$th standard basis vector. For a one-hot token matrix $S$, the initial residual matrix is $X_0 = S E$.
**Lemma (injectivity for additive position vectors).** The map $\widetilde\varepsilon(a,i)=\varepsilon(a)+p_i$ is injective exactly when
$$
\varepsilon(a)-\varepsilon(b)\ne p_j-p_i
$$
for every distinct pair $(a,i)\ne(b,j)$.
_Proof._ Equality of two combined embeddings is equivalent to $\varepsilon(a)-\varepsilon(b)=p_j-p_i$. excluding that equality for every distinct pair is precisely injectivity. distinct token embeddings and distinct position vectors alone are insufficient because their differences can cancel. [^embed-proof]
**Lemma (layer norm linearization).** Let $z=x-\mu\mathbf1$ and $s=\sqrt{d^{-1}\lVert z\rVert_2^2+\epsilon}$. the Jacobian of $\mathcal N(x)=\gamma\odot z/s+\beta$ is
$$
J_{\mathcal N}(x)
=\frac{\operatorname{diag}(\gamma)}{s}
\left(I-\frac{\mathbf1\mathbf1^T}{d}-\frac{zz^T}{ds^2}\right).
$$
_Proof._ Differentiate $\mu$, $z$, and $s$, then collect the three terms. the bracketed matrix annihilates $\mathbf1$, so layer norm is not locally invertible. when $\epsilon>0$ it is not a projection: its eigenvalue along $z$ is $\epsilon/s^2$. [^layernorm-proof]
### single-head attention as a tensor product
Fix a head $h$ inside layer $\ell$. Following the transformer-circuits factorization we write the attention output as a Kronecker-style product[^Kronecker] acting on $X_{\ell}$:\[@elhage2021mathematical\]
$$
H^{(h)}(X_{\ell}) = A^{(h)}(X_{\ell})X_{\ell}W_{OV}^{(h)}
$$
[^Kronecker]: In machine learning, we often deal with matrices and tensors where we want to multiply “one side.”
In transformers specifically, our activations are often 2D arrays representing vectors at different context indices, and we often want to multiply **per position** or **across positions**. Frequently, we want to do both! Tensor (or equivalently, Kronecker) products are a really clean way to denote this. We denote them with the $\otimes$ symbol.
Very pragmatically:
- A product like $\mathrm{Id} \otimes W$ (with identity on the left) represents multiplying each position in our context by a matrix $W$.
- A product like $A \otimes \mathrm{Id}$ (with identity on the right) represents multiplying **across positions**.
- A product like $A \otimes W$ multiplies the vector at each position by $W$ _and_ across positions with $A$. It doesn’t matter which order you do this in.
- The products obey the **mixed-product property**:
$$
(A \otimes B) \cdot (C \otimes D) = (A C) \otimes (B D).
$$
There are several completely equivalent ways to interpret these products. If the symbol is unfamiliar, you can pick whichever you feel most comfortable with:
- **Left-right multiplying**: Multiplying $x$ by a tensor product $A \otimes W$ is equivalent to simultaneously left and right multiplying:
$$
(A \otimes W)\, x = A\, x\, W^\top.
$$
When we sum them, it is equivalent to adding the results of this multiplication:
$$
(A_1 \otimes W_1 + A_2 \otimes W_2)\, x = A_1 x W_1^\top + A_2 x W_2^\top.
$$
- **Kronecker product**: The operations we want to perform are linear transformations on a flattened (“vectorized”) version of the activation matrix $x$. But flattening gives us a huge vector, and we need to map our matrices to a much larger block matrix that performs operations like “multiply the elements which previously corresponded to a vector by this matrix.” The correct operation to do this is the Kronecker product. So we can interpret $\otimes$ as a Kronecker product acting on the vectorization of $x$, and everything works out equivalently.
- **Tensor product**: $A \otimes W$ can be interpreted as a tensor product turning the matrices $A$ and $W$ into a 4D tensor. In NumPy notation, it is equivalent to
$$
A[:, :, \mathrm{None}, \mathrm{None}] * W[\mathrm{None}, \mathrm{None}, :, :]
$$
(though one wouldn’t computationally represent them in that form). More formally, $A$ and $W$ are “type $(1,1)$” tensors (matrices mapping vectors to vectors), and $A \otimes W$ is a “type $(2,2)$” tensor (which can map matrices to matrices).
with $A^{(h)}(X_{\ell})\in\mathbb{R}^{T\times T}$ encoding token-to-token routing and $W_{OV}^{(h)} = W_V^{(h)}W_O^{(h)}\in\mathbb{R}^{d_{model}\times d_{model}}$ a low-rank map ($\mathrm{rank} \le d_{head}$).
- Because the only nonlinearity resides in the softmax defining $A^{(h)}$, freezing $A^{(h)}$ linearizes the head completely.
- the idea of [[thoughts/mathematical framework transformers circuits#attention heads as information movement|information movement]] comes from the fact attention compute value vectors for each tokens from residual stream, and linearly combine each of those vectors based on attention pattern.
- i.e linear map $[n_\text{context}, d_{\text{model}}] \to [n_{\text{context}}, d_{\text{model}}]$
- This enables path analysis:
- expand logits as sums of products of $W_{OV}$ and $W_{QK}$ matrices and map each product to a circuit acting on a basis vector in the residual stream. \[@elhage2021mathematical\]
- An attention head applies two linear operations, $A$ across positions and $W_VW_O$ across residual coordinates.
- $A$ governs which token’s information to move from and to
- $W_VW_O$ governs which information is read from the source token and how it is written to the destination token.
> \[!tip\] intuition
>
> think of $W_{OV}$ as choosing a subspace to copy from the source token and another to write into at the destination. singular vectors of $W_{OV}$ highlight reusable features; tracking them across layers surfaces induction heads and copy circuits.
- causal mask: define $M\in\mathbb{R}^{T\times T}$ with $M_{ij}=0$ for $j\le i$ and $M_{ij}=-\infty$ otherwise. then $A^{(h)} = \operatorname{softmax}\!\left(\frac{Q^{(h)}K^{(h)\top}}{\sqrt{d_k}} + M\right)$ is lower triangular and can have a nonzero diagonal because a token may attend to itself.
- stochasticity: each row of $A^{(h)}$ sums to $1$. causal attention is generally not column-stochastic, so it does not lie in the Birkhoff polytope.[^birkhoff]
**Proposition (causal support is closed under composition).** If $A,B\in\mathbb R^{T\times T}$ are lower triangular attention matrices, then $AB$ is lower triangular.
_Proof._ For $j>i$,
$$
(AB)_{ij}=\sum_k A_{ik}B_{kj}=0,
$$
because a nonzero summand would require both $k\le i$ and $j\le k$, which is impossible. the product remains causal. it is not generally nilpotent because its diagonal can be nonzero.
#### observations on single-head attention
- $A$ is the **only non-linear part** of this equation (being computed from a softmax). This means that if we fix the attention pattern, attention heads perform a linear operation. This also means that, **without fixing** $A$, attention heads are “half-linear” in some sense, since the per-token linear operation is constant.
- $W_Q$ and $W_K$ always operate together. They’re never independent. Similarly, $W_O$ and $W_V$ always operate together as well.
- Although they’re parameterized as separate matrices, $W_V W_O$ and $W_Q W_K^T$ can be treated as individual low-rank maps under the row-activation convention.
- Keys, queries and value vectors are, in some sense, superficial. They’re intermediary by-products of computing these low-rank matrices. One could easily reparameterize both factors of the low-rank matrices to create different vectors, but still function identically.
- define $W_{OV}=W_VW_O$ and $W_{QK}=W_QW_K^T$ to keep the row-activation convention explicit.
- Products of attention heads behave much like attention heads themselves. By the distributive property,
$$
H^{(h_2)}(H^{(h_1)}(X))
=(A^{h_2}A^{h_1})X(W_{OV}^{h_1}W_{OV}^{h_2}).
$$
the composite has sequence operator $A^{h_2}A^{h_1}$ and residual operator $W_{OV}^{h_1}W_{OV}^{h_2}$. this is the virtual-head bookkeeping device.
### block-level composition
multi-head aggregation: a layer with heads $h=1,\dots,H$ yields
$$
M_{\ell}(X_{\ell}) = \sum_{h=1}^{H} H^{(h)}(X_{\ell})
$$
and the residual update becomes $X_{\ell+1}=X_{\ell}+M_{\ell}(X_{\ell})+G_{\ell}(X_{\ell})$ where $G_{\ell}$ is the mlp contribution.
- mlp as rank-expansion: with row activations, $G_\ell(x)=\phi(xW_1)W_2$, where $W_1\in\mathbb R^{d_{model}\times d_{ff}}$ and $W_2\in\mathbb R^{d_{ff}\times d_{model}}$. freezing the diagonal activation Jacobian gives the local residual-coordinate map $W_1\operatorname{diag}(\phi'(xW_1))W_2$.
- key-value memory view: geva et al. show that feed-forward neurons implement key-value associative memories; $W_1$ encodes keys, the activation acts as a selector, and $W_2$ writes cached values back into the residual stream, matching the residual-circuit picture.\[@geva2021transformerfeedforwardlayerskeyvalue\]
- virtual heads: products like $H^{(h_2)}\circ H^{(h_1)}$ have attention map $A^{(h_2)}A^{(h_1)}$ and row-convention OV map $W_{OV}^{(h_1)}W_{OV}^{(h_2)}$.\[@elhage2021mathematical\]
**Proposition (linearized residual path expansion).** Freeze the attention patterns, layer-norm Jacobians, and MLP activation Jacobians at one input. each layer is then a linear operator $I+B_\ell$ on $\operatorname{vec}(X)$, where $B_\ell$ contains both sequence-routing and residual-space factors. the unembedded output is
$$
\operatorname{vec}(\operatorname{logits})
=(W_U^T\otimes I_T)
\prod_{\ell=0}^{L-1}(I+B_\ell)
\operatorname{vec}(X_0).
$$
_Proof._ Expanding the ordered product chooses either the identity or one update from each layer. every term is therefore an ordered residual path. attention terms retain both $A^{(\ell,h)}$ and $W_{OV}^{(\ell,h)}$; MLP terms retain their local Jacobians. the unembedding maps the final residual coordinates to logits.\[@elhage2021mathematical\][^path-proof]
### mlp superposition (geva et al.)
- Factor the feed-forward block as a key-value memory: $W_1$ supplies keys, $W_2$ provides values, and the non-linearity selects which keys fire.\[@geva2021transformerfeedforwardlayerskeyvalue\]
- Superposition appears because different features reuse the same neuron; activations overlap in the residual basis until you rotate into a privileged coordinate system (see the toy-model section below).
**Lemma (key-value decomposition).** Let $x\in\mathbb{R}^{d_{model}}$, $K=W_1$, $V=W_2^\top$, and $a=\phi(Kx)$ with $\phi$ applied elementwise. Then the MLP output can be written as
$$
G_{\ell}(x) = V^\top a = \sum_{j=1}^{d_{ff}} a_j\, v_j,
$$
where $v_j$ is the $j$th column of $V^\top$. Each neuron contributes a rank-one write $a_j v_j$; overlapping keys $K_j$ mean several semantic features can share the same neuron and therefore superpose inside $a$.[^geva-proof]
> \[!note\] spectral diagnostics
>
> singular values of $W_{OV}^{(h)}$ act like feature strength indicators.
>
> - The singular values $\sigma_i$ quantify how much a head stretches each principal direction; sharp decay (large $\sigma_1$, tiny tail) signals a nearly one-dimensional copier, while flat spectra imply superposed features with comparable gain.[^sigma-footnote]
> - Effective rank $r_{\text{eff}} = \exp\!\big(-\sum_i p_i \log p_i\big)$ with $p_i = \sigma_i^2 / \sum_j \sigma_j^2$ summarises the decay; it stays small for copy heads and grows when heads mix many directions.
> - The cumulative energy $E_k = \sum_{i\le k}\sigma_i^2 / \sum_i \sigma_i^2$ tells you how many singular directions you need to capture a target fraction of variance.
> - The diagnostic script [[lectures/412/tools#scripts]] computes these curves, logs spectra, and now measures effective rank so you can separate low-rank copiers from high-rank mixers in practice.
> - Example below: layer 12 head 8 on the prompt “The true meaning of life is absurdity, and suffering” shows both spectrum and attention routing for the same head.
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>
>