1 Attention functions and transformers
Throughout, \(C\) is a finite-dimensional real vector space. Smooth means \(C^\infty \), and \(\Delta _{k-1}\subset \mathbb {R}^k\) is the standard simplex.
A smooth map \(a\colon C^k\to C\) is an attention function if there are a smooth \(\omega \colon C^k\to \Delta _{k-1}\), a smooth convex \(\phi \colon \mathbb {R}^k\to \mathbb {R}\) and a polynomial map \(\sigma \colon C^k\to \mathbb {R}^k\) with
The data \(\omega ,\phi ,\sigma \) are carried along with \(a\). A polynomial map is one whose components are polynomials in the coordinates of \(x_1,\ldots ,x_k\).
An attention function whose scoring function has the form \(\sigma (x)=(P_{k-1}(x_1,x_k),\ldots ,P_0(x_k,x_k))\) for polynomials \(P_d\colon C^2\to \mathbb {R}\) of bidegree at most \((1,1)\) is biaffine; it is bilinear if every \(P_d\) is a bilinear form.
Every polynomial \(P\colon C\times C\to \mathbb {R}\) of bidegree at most \((1,1)\) has a unique decomposition \(P(u,v)=B(u,v)+L(u)+R(v)+D\) with \(B\) bilinear, \(L,R\) linear and \(D\in \mathbb {R}\), namely
The paper calls this elementary and gives no proof. Each monomial of bidegree at most \((1,1)\) is a product of an affine function of \(u\) and an affine function of \(v\); expanding the product shows that the four expressions are bilinear, linear, linear and constant, and they add up to \(P\). Evaluating any decomposition at \((0,0)\), \((u,0)\) and \((0,v)\) gives the formulas, hence uniqueness.
Let a bilinear attention function on a head space \(H\) have forms \(B_{k-1},\ldots ,B_0\), let \(L_j\colon H\to H^\vee \), \(L_j(v)(u)=B_j(u,v)\), assume \(B_0\) nondegenerate, and set \(T_d=L_0^{-1}L_d\). The function is RoPE-like if \(T_{d+e}=T_dT_e\) whenever \(d,e\ge 0\) and \(d+e{\lt}k\).
With \(B_0\) nondegenerate, the function is RoPE-like if and only if there is \(S\in \operatorname {End}(H)\) with \(T_d=S^d\) for all \(0\le d{\lt}k\). When \(S\) exists it is \(T_1=L_0^{-1}L_1\).
The RoPE-like condition gives \(L_{d+e}=L_dL_0^{-1}L_e\), in particular \(L_{d+1}=L_dS\) with \(S=T_1\). By induction \(L_d=L_0S^d\), that is \(T_d=S^d\). Conversely \(S^{d+e}=S^dS^e\).
For attention functions \(a_k\colon C^k\to C\), \(k=2,\ldots ,\ell \), the causal attention head is \(A(x_1,\ldots ,x_\ell )=(x_1,a_2(x_1,x_2),\ldots ,a_\ell (x_1,\ldots ,x_\ell ))\).
For a head \(A\), a linear value map \(V\colon C\to C\) and a map \(M\colon C\to C\), the single-head transformer is \(T=T_M\circ T_A\) with \(T_A(x)_i=x_i+V(A(x)_i)\) and \(T_M(y)_i=y_i+M(y_i)\).
Fix a norm on \(C\) and put \(\lVert x\rVert _\infty =\max _k\lVert x_k\rVert \). If \(\lVert M(y)\rVert \le L\lVert y\rVert \) for all \(y\) and \(V=\sum _i\lVert V_i\rVert \) (operator norms), then \(\lVert T(x)-x\rVert _\infty \le (L+V+LV)\lVert x\rVert _\infty \).
Each component of a head is a convex combination of its inputs, so its norm is at most \(\lVert x\rVert _\infty \); the attention block is then bounded by \(V\lVert x\rVert _\infty \). The \(k\)-th component of \(T(x)-x\) is \(A(x)_k+M(x_k+A(x)_k)\), whose norm is at most \(V\lVert x\rVert _\infty +L(1+V)\lVert x\rVert _\infty \). The Lean statement assumes only that each head returns convex combinations of its inputs, which is what the proof uses.
Let \(\widehat C=C\oplus \mathbb {R}e\) and \(\iota (x)=x+e\). For every single-head QKV transformer \(T\colon C^\ell \to C^\ell \), with scores \(\langle Q(x_k+p_k),K(x_j+p_j)\rangle /\sqrt n\), there is a single-head bilinear transformer \(\widehat T\) on \(\widehat C\) with the same potentials such that \(\widehat T(\iota x_1,\ldots ,\iota x_\ell )=\iota ^\ell (T(x_1,\ldots ,x_\ell ))\).
Define \(\widehat B(u+\alpha e,v+\beta e)=\langle Q(v+\beta p_k),K(u+\alpha p_j)\rangle /\sqrt n\), which expands to the paper’s four-term formula. It agrees with the QKV score on the chart \(C+e\), so the weights agree there and \(\sum _j\omega _j(x_j+e)=a_k(x)+e\). Extend \(V\) and \(M\) by ignoring the new coordinate. The homogenized weights are simplex-valued on all of \(\widehat C^k\) because \(\widehat\sigma (u+\alpha e,\ldots )=\sigma (x)\) with \(x_j=u_j+(\alpha _j-1)p_j\); this also gives their smoothness. The Lean statement allows any potentials, softmax included.