The Flow of Attention

Author(s): GSO1 Originally published on Towards AI. The Flow of Attention Introduction Picture an input prompt to a large language model as a cloud of points in a high-dimensional vector space E, one point for each token. As the model processes the prompt layer by layer, the cloud reconfigures itself — points shift position, some drift together, others pull apart — so that each token’s location better reflects how it relates to the others in context. By the final layer, the cloud has settled into a configuration that encodes the contextualized meaning of the prompt. This final configuration is not a random scatter: semantically related tokens have drifted into clusters, a structure we will see emerge as the natural endpoint of the dynamics described in this article. The next token is then read off from a single point in this configuration: the last token’s final vector determines a set of alignments with the rows of the unembedding matrix and softmax over these alignments produces a probability distribution over the vocabulary, from which the next token is sampled. This article is about that reconfiguration — what drives it, what shape it takes, and what it reveals about transformer attention. Some notation and a working definition of attention Let X = { x(1), x(2), …, x(N) } be a sequence of N token embeddings in the d-dimensional space E, and let q = x(i) be the embedding of the i-th token, which we call the query. Attention seeks a vector Δ(q, X) in E such that q* = q + Δ(q, X) is a contextualized version of q, that is, the i-th token’s representation updated by content drawn from X. The vector Δ(q, X) is called the contextual update of q and is computed by an attention head. By letting i run from 1 to N in q = x(i), the attention head generates a contextual update for every token of X. Arranging these updates as the columns of a d × N matrix gives the sequence Δ(X), which transforms the original cloud X into a new cloud X* by X* = X + Δ(X) where each point of X is displaced by its contextual update. This is one layer’s worth of reconfiguration. Two operators in one space — a geometric view of attention The two previous articles in this series developed the geometry underlying how a single attention head computes Δ(q, X), step by step and at a more leisurely pace: Article 1 recast standard attention as a problem of estimation, and Article 2 reduced its machinery to two operators acting in a single space. What follows is the compressed version. Note both articles focused on the cross-token interaction that defines attention and set aside other layer components such as layer normalization and feed-forward networks, which do not facilitate information transfer between tokens. Positional encoding does shape the underlying geometry of attention, but how relevance scores are determined from that geometry is unchanged; a short note at the end of this article explains why. Article 1 showed that Δ(q, X) is computed in two steps. First, attention weights are inferred from relevance scores given by scaled dot products of learned query (Q) and key (K) projections into low-dimensional subspaces of E. The inference is a softmax normalization of these scores — a computationally efficient, closed-form solution to a free-energy minimization problem that fits the relevance signal while preventing the weights from collapsing onto a single token. Second, the contextual update Δ(q, X) is computed as an attention-weighted average of N learned value (V) projections, the closed-form solution to a classical weighted least squares problem. This two-step formulation, built around four projection matrices and three intermediate spaces, is the standard QKV model of attention. Article 2 showed that the QKV machinery — multiple projections, dot products, softmax normalizations, weighted sums — reduces to two operators acting in the single space E: • A bilinear form M that assigns a relevance score to an ordered pair of embeddings from E. • A linear operator F that extracts content from a token embedding. For a fixed query q, the function M(q, t) is a linear functional on E (a linear map from E to the real numbers) that scores the relevance of any token t to q, that is, the degree to which q should attend to t. Geometrically, it stratifies E into parallel hyperplanes of constant relevance, all normal to the relevance gradient r_q = Mᵀq, where M also denotes the matrix of the bilinear form such that M(x₁, x₂) = x₁ᵀMx₂. Mathematicians call this stack of hyperplanes a foliation; the individual hyperplanes are its leaves. Every token x sits on exactly one leaf, with score M(q, x) = ⟨r_q, x⟩ — the projection of x onto r_q. Tokens on the same leaf score identically, and since r_q points toward increasing relevance, the leaves are ordered by score along r_q. Change q and r_q changes, reorienting the foliation and rescoring every token in X. Relevance, then, is determined by leaf membership. Solving the free-energy minimization problem on these scores yields the attention weights via softmax, and the contextual update Δ(q, X) is the attention-weighted average of the content vectors F(xⱼ). This update is a vector in the range of F — the query-agnostic content subspace shared by all queries — and is identical to the update produced by the QKV model. This single-space formulation is the EMF model of attention. The EMF model is what we use in the rest of this article. Its advantage for what follows is geometric: because everything happens in E, the contextual updates Δ(q, X) are vectors in the same space as the tokens themselves. The reconfiguration X → X + Δ(X) therefore represents a motion of the cloud through E, governed by two operators per layer — M shaping where each query looks, F shaping what it finds. That motion, layer after layer, is what the next sections describe. The residual stream A transformer isn’t just […]