The standard self-attention output is a hop along the normal of an energy surface — a geometric picture that makes the algebra transparent
We start from the standard single-head causal self-attention, stated in a form that will connect directly to the energy landscape.
Consider a single attention head with $d_{\text{model}} = d_k = d$. The input sequence is $\boldsymbol{x}_0, \boldsymbol{x}_1, \ldots, \boldsymbol{x}_t \in \mathbb{R}^{1 \times d}$ (row vectors), stacked into the token matrix
$$X \;=\; \begin{bmatrix} \boldsymbol{x}_0 \\ \boldsymbol{x}_1 \\ \vdots \\ \boldsymbol{x}_t \end{bmatrix} \in \mathbb{R}^{(t+1) \times d}.$$The learnable weight matrices are $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$. Define the combined key-query matrix $M = W_Q\, W_K^\top$ and the transformed token matrix
$$\tilde{Y} \;=\; X\, W_K\, W_Q^\top \;=\; X\, M^\top \in \mathbb{R}^{(t+1) \times d}.$$Row $\tau$ of $\tilde{Y}$ is $\tilde{\boldsymbol{y}}_\tau = \boldsymbol{x}_\tau M^\top$, the basin center for token $\tau$.
In the standard formulation, the query, key, and value vectors are $\boldsymbol{q}_t = \boldsymbol{x}_t W_Q$, $K = X W_K$, $V = X W_V$. The self-attention output at position $t$ is
where $\mathrm{softmax}$ is applied across the row (over the token index $\tau = 0, \ldots, t$). Since $W_Q W_K^\top = M$ and $\tilde{Y} = X M^\top$, the score vector is $\boldsymbol{x}_t\,M\,X^\top = \boldsymbol{x}_t\,\tilde{Y}^\top$. Writing $\beta = 1/\sqrt{d_k}$, the output becomes
$$\boldsymbol{o}_t \;=\; \boldsymbol{a}\,X\,W_V, \qquad a \;=\; \mathrm{softmax}\bigl(\beta\, \boldsymbol{x}_t\,\tilde{Y}^\top\bigr) \;\in\; \mathbb{R}^{1 \times (t+1)}.$$The output $\boldsymbol{o}_t$ is a softmax-weighted average of $\boldsymbol{x}_\tau W_V$, where the weights $a$ depend on the dot product between the query $\boldsymbol{x}_t$ and the rows of $\tilde{Y}$.
This is the formula everyone knows. The question is: what is this computation actually doing, geometrically?
The key observation is that the attention weights $a$ and the weighted combination $\boldsymbol{a}\,\tilde{Y}$ arise naturally as a single step of the Concave-Convex Procedure (CCCP) on a specific energy function. The argument proceeds in three steps: (1) recognize that softmax is the gradient of a concave function, (2) add a convex bowl to get an energy whose density is a Gaussian mixture, (3) observe that the CCCP on this energy reproduces the attention formula.
The starting point is a well-known identity. Define the log-sum-exp (lse) function
$$\mathrm{lse}(\boldsymbol{\xi}) \;=\; \beta^{-1}\log\sum_{\tau \le t} \exp\!\bigl(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau\bigr).$$This is a smooth, convex function of $\boldsymbol{\xi}$ (it is the log-partition function of a categorical distribution). Its gradient is
where $a_\tau = \mathrm{softmax}_\tau(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)$. In other words: the softmax-weighted average of the $\tilde{\boldsymbol{y}}_\tau$ is the gradient of the lse function. This is the link between the attention formula and an energy landscape.
Since $\mathrm{lse}$ is convex, $-\mathrm{lse}$ is concave. Set $E_{\mathrm{cave}}(\boldsymbol{\xi}) = -\mathrm{lse}(\boldsymbol{\xi})$ and add a convex bowl $\frac{1}{2}\|\boldsymbol{\xi}\|^2$ to define the energy
$$E_t(\boldsymbol{\xi}) \;=\; \underbrace{-\mathrm{lse}(\boldsymbol{\xi})}_{\displaystyle E_{\mathrm{cave}}(\boldsymbol{\xi})\;\text{(concave)}} \;+\; \underbrace{\frac{1}{2}\|\boldsymbol{\xi}\|^2}_{\displaystyle\text{bowl (convex)}} \;=\; -\beta^{-1}\log\sum_\tau \exp\!\bigl(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau\bigr) \;+\; \frac{1}{2}\|\boldsymbol{\xi}\|^2.$$By the polarization identity — $\exp(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau) = e^{\frac{\beta}{2}\|\boldsymbol{\xi}\|^2}\, e^{-\frac{\beta}{2}\|\boldsymbol{\xi} - \tilde{\boldsymbol{y}}_\tau\|^2}\, e^{\frac{\beta}{2}\|\tilde{\boldsymbol{y}}_\tau\|^2}$ — the $\frac{1}{2}\|\boldsymbol{\xi}\|^2$ bowl and the $-\frac{1}{2}\|\boldsymbol{\xi}\|^2$ from polarization cancel, and the energy becomes
$$E_t(\boldsymbol{\xi}) \;=\; -\beta^{-1}\log\sum_{\tau} \pi_\tau\,\exp\!\Bigl(-\tfrac{\beta}{2}\|\boldsymbol{\xi} - \tilde{\boldsymbol{y}}_\tau\|^2\Bigr),$$where $\pi_\tau = \exp\!\bigl(\tfrac{\beta}{2}\|\tilde{\boldsymbol{y}}_\tau\|^2\bigr)$. This is the negative log-density of a Gaussian mixture model with one Gaussian per token, centered at the transformed position $\tilde{\boldsymbol{y}}_\tau$. The $\frac{1}{2}\|\boldsymbol{\xi}\|^2$ bowl is essential: it is what turns the lse into a GMM.
The CCCP is a standard procedure for decreasing functions that decompose as convex + concave. For our energy $E_t(\boldsymbol{\xi}) = \underbrace{\frac{1}{2}\|\boldsymbol{\xi}\|^2}_{E_{\mathrm{convex}}} + E_{\mathrm{cave}}(\boldsymbol{\xi})$, the CCCP update rule is
$$\nabla E_{\mathrm{convex}}(\tilde{\boldsymbol{x}}_{t+1}) \;=\; -\nabla E_{\mathrm{cave}}(\boldsymbol{x}_t),$$which evaluates the concave gradient at the current point $\boldsymbol{x}_t$ and matches it to the convex gradient at the next point $\tilde{\boldsymbol{x}}_{t+1}$. This is guaranteed to decrease (or maintain) the energy at each step.
Now the choice of $\frac{1}{2}\|\boldsymbol{\xi}\|^2$ as the convex part becomes essential: its gradient is simply $\nabla E_{\mathrm{convex}}(\boldsymbol{\xi}) = \boldsymbol{\xi}$. So the CCCP rule becomes
$$\tilde{\boldsymbol{x}}_{t+1} \;=\; -\nabla E_{\mathrm{cave}}(\boldsymbol{x}_t).$$The landing point $\tilde{\boldsymbol{x}}_{t+1}$ appears directly on the left-hand side — no extra transformation needed. A different convex term (say $\frac{1}{2}\boldsymbol{\xi}^\top A\boldsymbol{\xi}$) would give $A\tilde{\boldsymbol{x}}_{t+1} = -\nabla E_{\mathrm{cave}}$, requiring an inversion.
Since the full gradient is $\nabla E_t(\boldsymbol{\xi}) = \boldsymbol{\xi} + \nabla E_{\mathrm{cave}}(\boldsymbol{\xi})$, we can also write the update as
$$\tilde{\boldsymbol{x}}_{t+1} \;=\; \boldsymbol{x}_t - \nabla E_t(\boldsymbol{x}_t).$$So in our case, the CCCP update happens to coincide with one full gradient descent step on $E_t$ (with step size 1) — a consequence of the Hessian being $I$.
Now using §2.1 — the fact that $-\nabla E_{\mathrm{cave}} = \nabla\,\mathrm{lse}$ is the softmax-weighted average — we get
$$\tilde{\boldsymbol{x}}_{t+1} \;=\; \nabla\,\mathrm{lse}(\boldsymbol{x}_t) \;=\; \boldsymbol{a}\,\tilde{Y}.$$This is exactly the core computation of self-attention: the softmax-weighted average of the transformed tokens.
To see how the standard attention output relates to the energy landscape, we rewrite the attention formula by substituting $X = \tilde{Y}\,(W_K W_Q^\top)^{-1}$ (valid when $M$ is invertible; the rank-deficient case is handled in §2.5):
$$\boldsymbol{o}_t \;=\; \mathrm{softmax}\!\bigl(\beta\, \boldsymbol{x}_t\,\tilde{Y}^\top\bigr)\;\cdot\; \underbrace{\tilde{Y}\,(W_K W_Q^\top)^{-1}}_{= \;X}\;\cdot\; W_V.$$Now group the first two factors:
This decomposition is the key insight. The attention output factors into two stages:
In words: standard self-attention is one CCCP step on the energy $E_t$, followed by a fixed linear readout.
In practice, $W_Q, W_K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ with $d_k < d_{\text{model}}$, so $M = W_Q W_K^\top$ has rank $\leq d_k$ and the substitution $X = \tilde{Y}\,(W_K W_Q^\top)^{-1}$ used in §2.4 is not valid. Does the geometric picture survive?
Yes. Replace the inverse with the pseudo-inverse $(W_K W_Q^\top)^\dagger$, and the factorization becomes
where $P = M^\top (M^\top)^\dagger$ is the orthogonal projector onto $\mathrm{col}(M^\top)$. The first term is what the hop determines — Stage 1 ($\boldsymbol{a}\,\tilde{Y}$) followed by Stage 2 (readout via pseudo-inverse). The second term is a linear function of $\boldsymbol{a}\,X$ that lives in the null space of $M^\top$ — it does not participate in basin selection and cannot be "seen" by the landscape. When $M$ is full rank, $P = I$ and the residual vanishes, recovering §2.4 exactly.
The CCCP update $\tilde{\boldsymbol{x}}_{t+1} = \boldsymbol{x}_t - \nabla E_t(\boldsymbol{x}_t)$ has a striking geometric meaning that arises from one special property of our energy: the convex part has Hessian equal to the identity (it is $\frac{1}{2}\|\boldsymbol{\xi}\|^2$). This means the CCCP step size is exactly 1 in the "right" sense, which leads to the following construction.
Consider the graph of $E_t$ as a surface $\Gamma = \{(\boldsymbol{\xi}, E_t(\boldsymbol{\xi}))\} \subset \mathbb{R}^{d+1}$. At the point $(\boldsymbol{x}_t, E_t(\boldsymbol{x}_t))$ on this surface, the upward unit normal is
$$\hat{\boldsymbol{n}}(\boldsymbol{x}_t) \;=\; \frac{1}{\sqrt{1 + \|\nabla E_t(\boldsymbol{x}_t)\|^2}}\,\bigl(-\nabla E_t(\boldsymbol{x}_t),\; 1\bigr).$$Now perform the following two-step procedure:
Claim: the horizontal projection of the top of the hop is exactly $\tilde{\boldsymbol{x}}_{t+1} = \boldsymbol{x}_t - \nabla E_t(\boldsymbol{x}_t)$, the CCCP landing point.
The height rise of 1 requires walking a distance $L = \sqrt{1 + \|\nabla E_t\|^2}$ along $\hat{\boldsymbol{n}}$ (since the vertical component of $\hat{\boldsymbol{n}}$ is $1/\sqrt{1 + \|\nabla E_t\|^2}$). The horizontal displacement is then
$$L \cdot \bigl(-\nabla E_t / \sqrt{1 + \|\nabla E_t\|^2}\bigr) \;=\; -\nabla E_t(\boldsymbol{x}_t).$$Adding this to $\boldsymbol{x}_t$ gives $x_t - \nabla E_t(\boldsymbol{x}_t) = \tilde{\boldsymbol{x}}_{t+1}$. The crucial point: the height rise is exactly 1 because the bowl has Hessian $I$. A different convex term would give a different rise, and the construction would not produce the CCCP update.
Note that the vertical drop from the top of the hop back to the surface is not 1 — it depends on the shape of the landscape at the landing point and is typically much larger than 1.
Putting it all together, each autoregressive step of self-attention is:
① Build the energy landscape $E_t(\boldsymbol{\xi})$ from all past tokens $\tilde{\boldsymbol{y}}_0, \ldots, \tilde{\boldsymbol{y}}_t$.
② At the query $\boldsymbol{x}_t$ on the surface, hop along the upward normal until the elevation has increased by 1.
③ Drop vertically to the surface. The horizontal projection of this hop is $\tilde{\boldsymbol{x}}_{t+1}$.
④ Apply the fixed linear readout $(M^\top)^{-1} W_V$ to get $\boldsymbol{o}_t = \boldsymbol{x}_{t+1}$.
The landscape is reconstructed from scratch at every time step (with one more basin as each token is added), and the query hops across it to select the next state. This is why we call the mechanism energy landscape hopping.
In §2 we showed that the standard softmax attention is a CCCP step on an energy landscape. A natural question: what happens if we replace the softmax with a different nonlinearity? We can answer this by reverse-engineering the energy function — starting from the attention equation, then asking: "what energy has this update as its CCCP?"
We follow the same logic as §2. In §2 the key was: softmax is the gradient of lse. Here the parallel question is: tanh is the gradient of what?
Step 1: tanh is the gradient of $\beta^{-1}\log\cosh$. The identity $\frac{d}{du}\log\cosh(u) = \tanh(u)$ gives
$$\frac{\partial}{\partial \boldsymbol{\xi}_j}\,\beta^{-1}\sum_\tau \log\cosh(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau) \;=\; \sum_\tau \tanh(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)\,[\tilde{\boldsymbol{y}}_\tau]_j.$$So $\sum_\tau \tanh(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)\,\tilde{\boldsymbol{y}}_\tau = \tilde{Y}^\top \tanh(\beta\,\tilde{Y}\,\boldsymbol{\xi})$ is the gradient of $\beta^{-1}\sum_\tau \log\cosh(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)$, exactly paralleling the softmax case.
Step 2: set $E_{\mathrm{cave}} = -\beta^{-1}\sum_\tau \log\cosh(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)$ and add the bowl.
Step 3: expand cosh to reveal the modes. Using $\sum_\tau \log\cosh = \log\prod_\tau\cosh$, expand each factor:
$$\cosh(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau) \;=\; \tfrac{1}{2}\bigl(e^{+\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau} + e^{-\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau}\bigr).$$Each $\cosh$ contributes two terms: one for $+\tilde{\boldsymbol{y}}_\tau$ and one for $-\tilde{\boldsymbol{y}}_\tau$. Taking the product over all $n$ tokens, each token independently chooses $+$ or $-$, giving $2^n$ configurations:
$$\prod_\tau \cosh(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau) \;=\; \frac{1}{2^n}\sum_{\varepsilon \in \{-1,+1\}^n} \exp\!\left(\beta\,\boldsymbol{\xi}^\top \sum_\tau \varepsilon_\tau\, \tilde{\boldsymbol{y}}_\tau\right).$$Define the $2^n$ mode centers $\boldsymbol{z}_\varepsilon = \sum_\tau \varepsilon_\tau\,\tilde{\boldsymbol{y}}_\tau$. The $\beta^{-1}\log$ of this product is $\beta^{-1}\log\sum_\varepsilon \exp(\beta\,\boldsymbol{\xi}^\top \boldsymbol{z}_\varepsilon) - n\log 2/\beta$, which is the lse function evaluated at the $2^n$ mode centers (plus a constant). By the same polarization as §2.2:
This is exactly the same GMM form as the softmax energy, but summing over $2^n$ mode centers $\boldsymbol{z}_\varepsilon$ instead of $n$ centers $\tilde{\boldsymbol{y}}_\tau$. The $\beta$ appears in the same places: $\beta^{-1}$ outside the log, $\beta/2$ inside each Gaussian. For $n = 2$ tokens: $2^2 = 4$ modes at $\tilde{\boldsymbol{y}}_0 + \tilde{\boldsymbol{y}}_1$, $\tilde{\boldsymbol{y}}_0 - \tilde{\boldsymbol{y}}_1$, $-\tilde{\boldsymbol{y}}_0 + \tilde{\boldsymbol{y}}_1$, $-\tilde{\boldsymbol{y}}_0 - \tilde{\boldsymbol{y}}_1$ — the vertices of a parallelogram.
Step 4: the CCCP. $\tilde{\boldsymbol{x}}_{t+1} = -\nabla E_{\mathrm{cave}}(\boldsymbol{x}_t) = \tilde{Y}^\top \tanh(\beta\,\tilde{Y}\,\boldsymbol{x}_t) = \sum_\tau \tanh(\beta\,\boldsymbol{x}_t^\top \tilde{\boldsymbol{y}}_\tau)\,\tilde{\boldsymbol{y}}_\tau$.
The same logic one more time. If we remove the softmax, the raw scores $\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau$ are used directly as weights: $\tilde{\boldsymbol{x}}_{t+1} = \sum_\tau \beta(\boldsymbol{x}_t^\top \tilde{\boldsymbol{y}}_\tau)\,\tilde{\boldsymbol{y}}_\tau = \beta\,\tilde{Y}^\top \tilde{Y}\,\boldsymbol{x}_t$. This linear map is the gradient of what?
$$\nabla_{\boldsymbol{\xi}} \;\tfrac{\beta}{2}\sum_\tau (\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)^2 \;=\; \beta\sum_\tau (\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)\,\tilde{\boldsymbol{y}}_\tau \;=\; \beta\,\tilde{Y}^\top \tilde{Y}\,\boldsymbol{\xi}.$$So $E_{\mathrm{cave}}(\boldsymbol{\xi}) = -\frac{\beta}{2}\sum_\tau (\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)^2$ and the energy is
This is a pure quadratic form — not a mixture. The energy has no basins and no basin selection: the CCCP landing $\beta\,\tilde{Y}^\top \tilde{Y}\,\boldsymbol{x}_t$ is a fixed linear transformation of the query. When $\beta\,\tilde{Y}^\top \tilde{Y}$ has eigenvalues exceeding 1, the energy becomes a saddle (unbounded below) — reflecting the instability of unnormalized linear attention.
Note on concavity: $E_{\mathrm{cave}}(\boldsymbol{\xi}) = -\frac{\beta}{2}\boldsymbol{\xi}^\top \tilde{Y}^\top \tilde{Y}\,\boldsymbol{\xi}$ has Hessian $-\beta\,\tilde{Y}^\top\tilde{Y} \preceq 0$ (negative semidefinite), so it is concave. The CCCP decomposition and the hop construction from §3 are both valid — the bowl still has Hessian $I$, so the hop along the surface normal with height rise 1 gives the correct CCCP landing. The only difference is that the full energy may be a saddle, so repeated hops can diverge — but each individual hop is well-defined and still decreases the energy.
| Property | Softmax | Tanh | Linear (no softmax) |
|---|---|---|---|
| Per-token factor | $\exp(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)$ | $\cosh(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)$ | $\beta(\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)^2$ |
| $E_{\mathrm{cave}}(\boldsymbol{\xi})$ | $-\beta^{-1}\mathrm{lse}$ | $-\beta^{-1}\sum_\tau \log\cosh(\beta\,\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)$ | $-\frac{\beta}{2}\sum_\tau (\boldsymbol{\xi}^\top \tilde{\boldsymbol{y}}_\tau)^2$ |
| CCCP landing | $\boldsymbol{a}\,\tilde{Y}$ | $\tilde{Y}^\top\tanh(\beta\,\tilde{Y}\,\boldsymbol{x}_t)$ | $\beta\,\tilde{Y}^\top \tilde{Y}\,\boldsymbol{x}_t$ |
| Energy form | $-\beta^{-1}\log\sum \pi_\tau e^{-\frac{\beta}{2}\|\cdot\|^2}$ | $-\beta^{-1}\log\sum \pi_\varepsilon e^{-\frac{\beta}{2}\|\cdot\|^2} + c$ | $\frac{1}{2}\boldsymbol{\xi}^\top(I - \beta\,\tilde{Y}^\top\tilde{Y})\boldsymbol{\xi}$ |
| Landscape | GMM, $n$ modes | GMM, $2^n$ modes | Quadratic (possibly saddle) |
Verify the geometric interpretation in $d = 2$ with two learnable matrices. Generate tokens autoregressively and watch the landscape grow.
Basic workflow: Start with the rotation preset. Click "Generate x₁" — the landscape has one basin, and the hop lands there. Click again for x₂ — now two basins appear, and the query selects the closer one. Keep generating to watch the token trajectory form a circle (2D panel, right).
Verify the geometric claim:
Compare the three kernels:
Explore the parameters: