Two-head fixed-width Transformers implement finite-grid lookups

Put $u=(1,{\delta }^{-1},\dots ,{\delta }^{-d+1})$. For a grid column $x=(a_1\delta ,\dots ,a_d\delta )$, where $a_i\in \{0,\dots ,m-1\}$, its scalar code is $\mathrm{\ell }(x)=u^{\mathsf{T}}x,\mathrm{\ell }(x)\mathrm{/}\delta =a_1+a_{2}m+\cdots +a_d{m}^{d-1}\mathrm{.}$ Thus the valid column codes are distinct multiples of $\delta$ in $[0,B]$, where $B=\delta (q-1)$. If a column has a sentinel coordinate $-{\delta }^{-nd}$, the sum of all its other coordinate contributions is at most ${\delta }^{-d+1}-1$, so its code is at most $-{\delta }^{-nd}+{\delta }^{-d+1}-1<0$.

For a scalar threshold $b$, the single-head update $\psi (Z;b)=e_1(u^{\mathsf{T}}Z)\mathrm{hmax}\text{\hspace{0.17em}⁣}((u^{\mathsf{T}}Z{)}^{\mathsf{T}}(u^{\mathsf{T}}Z-b{\mathbf{1}}_n^{\mathsf{T}}))$ is an allowed attention update contribution: its key and value projections are $u$, its query projection is $u$ with bias $-b$, and its output vector is $e_1$. In query column $j$, if ${\mathrm{\ell }}_j>b$ it returns the largest current scalar code, and if ${\mathrm{\ell }}_j<b$ it returns the smallest. If ${\mathrm{\ell }}_j=b$, all query scores are zero and equal-weight hardmax returns the mean $\bar{\mathrm{\ell }}=n^{-1}{\sum }_k{\mathrm{\ell }}_k$. Thus, writing ${\mathrm{\ell }}_{\max }={\max }_k{\mathrm{\ell }}_k$ and ${\mathrm{\ell }}_{\min }={\min }_k{\mathrm{\ell }}_k$, the contribution $\psi (Z;b)-\psi (Z;b^{\mathrm{\prime }})$ for $b<b^{\mathrm{\prime }}$ is ${\mathrm{\ell }}_{\max }-{\mathrm{\ell }}_{\min }$ in the first coordinate when $b<{\mathrm{\ell }}_j<b^{\mathrm{\prime }}$, $\bar{\mathrm{\ell }}-{\mathrm{\ell }}_{\min }$ when ${\mathrm{\ell }}_j=b$, ${\mathrm{\ell }}_{\max }-\bar{\mathrm{\ell }}$ when ${\mathrm{\ell }}_j=b^{\mathrm{\prime }}$, and zero when ${\mathrm{\ell }}_j<b$ or ${\mathrm{\ell }}_j>b^{\mathrm{\prime }}$. In particular, if no current code equals either threshold, it shifts exactly the columns strictly between them by ${\mathrm{\ell }}_{\max }-{\mathrm{\ell }}_{\min }$.

Sweep the valid scalar codes $0,\delta ,\dots ,B$ in increasing order, using the two-head update $Z\mapsto Z+q(\psi (Z;\mathrm{\ell }-\delta \mathrm{/}2)-\psi (Z;\mathrm{\ell }+\delta \mathrm{/}2))$ at code $\mathrm{\ell }$. All initial codes, including sentinel codes, are multiples of $\delta$: every grid coordinate lies in $\delta \mathbb{Z}$ and every entry of $u$ is an integer. Since $q=m^d$ is an integer, each strict selective shift preserves this lattice; inductively no current code can equal any half-grid threshold $\mathrm{\ell }\pm \delta \mathrm{/}2$. This also applies to duplicate-column inputs and to the position-tagged lattice. For $L\in {\widehat{G}}_{\delta }$, order its distinct column codes as ${\mathrm{\ell }}_1<\cdots <{\mathrm{\ell }}_n$. The first sweep shifts only ${\mathrm{\ell }}_1$ to ${\tilde{\mathrm{\ell }}}_1={\mathrm{\ell }}_1+q({\mathrm{\ell }}_n-{\mathrm{\ell }}_1)>B$, since ${\mathrm{\ell }}_n-{\mathrm{\ell }}_1\ge \delta$ and $q\delta =B+\delta$. Each later sweep shifts the next unmodified code and gives ${\tilde{\mathrm{\ell }}}_j={\mathrm{\ell }}_j+q({\tilde{\mathrm{\ell }}}_{j-1}-{\mathrm{\ell }}_j)>{\tilde{\mathrm{\ell }}}_{j-1}(2\le j\le n)\mathrm{.}$ Thus shifted columns leave all later threshold intervals, and the final scalar codes satisfy $0<{\tilde{\mathrm{\ell }}}_1<\cdots <{\tilde{\mathrm{\ell }}}_n$.

Write $g_j={\mathrm{\ell }}_{j+1}-{\mathrm{\ell }}_j\ge \delta$. Unrolling the recurrence in step 2.1 gives $t:={\tilde{\mathrm{\ell }}}_n={\mathrm{\ell }}_n+{\sum }_{j=1}^{n-1}(q^n-q^{n-j})g_j\mathrm{.}$ Consequently $\delta q^{n-1}(q-1)\le t\le \delta q^n(q-1)-\delta (q-1{)}^2\mathrm{.}$ The lower bound follows from the $j=1$ term; the upper bound follows from ${\mathrm{\ell }}_n\le B$ and ${\sum }_j{g}_j\le B$. The map $({\mathrm{\ell }}_1,\dots ,{\mathrm{\ell }}_n)\mapsto t$ is injective: if two ordered tuples give the same $t$, put $a_j=({\mathrm{\ell }}_j-{\mathrm{\ell }}_j^{\mathrm{\prime }})\mathrm{/}\delta \in \mathbb{Z}$, so $\mathrm{\mid }{a}_j\mathrm{\mid }\le q-1$. The difference equation is $a_n+q(a_{n-1}-a_n)+\cdots +q^{n-1}(a_1-a_2)+q^n(a_n-a_1)=0.$ Reduction modulo $q$ first gives $a_n=0$; division by $q$ and repeated reduction modulo $q$ then gives $a_{n-1}=0,\dots ,a_1=0$. Hence distinct unordered column multisets have distinct $t$ values, separated by at least $\delta$.

Apply one final single-head update $Z\mapsto Z+K\psi (Z;0)$, with a positive integer $K$. For valid $L$, all ${\tilde{\mathrm{\ell }}}_j>0$, so it adds $Kt$ to every first coordinate and changes the scalar codes to $c_{L,j}=Kt+{\tilde{\mathrm{\ell }}}_j\mathrm{.}$ For different column multisets, their $t$ values differ by at least $\delta$, while $0<{\tilde{\mathrm{\ell }}}_j\le t$. Choose $K$ so that $K\delta >\max t$ over the finite valid grid. The intervals $[Kt,Kt+t]$ for different multisets are disjoint, and codes within one multiset are distinct by step 2.1. All resulting codes are multiples of $\delta$, hence distinct codes are at least $\delta$ apart.

If $L$ has duplicate columns but no sentinel, its scalar codes have only $r<n$ distinct levels. The same sweep shifts these $r$ levels, and its largest value is at most $\delta q^r(q-1)-\delta (q-1{)}^2\le \delta q^{n-1}(q-1)-\delta (q-1{)}^2$, strictly below the lower bound for valid $t$ in step 3.1. A sentinel column has a negative code by step 1.1 and remains the minimum throughout the sweep. If all codes are negative, the final update in step 4.1 shifts every code further negative. If some codes are nonnegative, the first nonnegative code is shifted by at least $q(q^n-B)>\delta q^{n+1}$; the later positive codes are shifted by still larger amounts. This exceeds every valid $t$ in step 3.1. Thus, after the final update, each invalid scalar code has one of three forms: below every valid code, above every valid code, or negative. Since there are finitely many grid inputs and all codes are multiples of $\delta$, choose $K$ large enough that every invalid code lies at least $\delta$ outside the interval from the least to the greatest valid code.

The preceding construction gives a finite scalar code for each token of each valid context; these codes are all distinct, and column permutations carry each code and its token to the corresponding permuted token by [F2]. Let ${\lambda }_{-}$ and ${\lambda }_{+}$ be the least and greatest valid codes. Because the table $A$ and grid are finite, enlarge $K$ so that $\mathrm{\mid }{u}^{\mathsf{T}}A(L{)}_{:,j}\mathrm{\mid }<{\lambda }_{-}-\delta \mathrm{/}3$ for every valid $L,j$. A three-piece constant activation that is zero on $[{\lambda }_{-}-\delta \mathrm{/}3,{\lambda }_{+}+\delta \mathrm{/}3]$ and one outside this interval detects exactly the invalid codes. Subtracting a constant larger than the maximum absolute coordinate of every encoded grid state makes every invalid state strictly negative; for each coordinate $r$, the allowed residual update $Z\mapsto Z+e_r\mathrm{ReLU}(-e_r^{\mathsf{T}}Z)$ replaces that coordinate by $z_r+\max (0,-z_r)=\max (0,z_r)$. These $d$ updates send every invalid state to zero and leave every valid, coordinatewise nonnegative state unchanged. For each valid code $c_{L,j}$, a three-piece pulse equal to one on $[-\delta \mathrm{/}3,\delta \mathrm{/}3]$ and zero outside, applied to $u^{\mathsf{T}}Z-c_{L,j}$, selects exactly its token. Add the vector $A(L{)}_{:,j}-Z_{:,j}$ at that token. The updated token has scalar value $u^{\mathsf{T}}A(L{)}_{:,j}$ below every code, so later pulses cannot select it; the remaining tokens retain their codes until selected. Equivariance of $A$ makes each update correct for every permutation of $L$. This yields $A(L)$ on ${\widehat{G}}_{\delta }$ and zero on every other point of $G_{\delta }^{+}$.

For the position-tagged grid, the scalar code of column $j$ has the form ${\mathrm{\ell }}_j=(j-1)S+r_j,S={\sum }_{k=0}^{d-1}{\delta }^{-k},0\le r_j\le B,r_j\in \delta \mathbb{Z}\mathrm{.}$ Since $S-B\ge \delta$, codes in a matrix satisfy ${\mathrm{\ell }}_1<\cdots <{\mathrm{\ell }}_n$ and each individual code identifies its column vector. Choose an integer $Q\ge q$ with $Q\delta >(n-1)S+B$. Sweep all possible position-tagged scalar codes in increasing order, replacing the shift multiplier $q$ in step 2.1 by $Q$. All codes remain in $\delta \mathbb{Z}$, so none equals a half-grid threshold. Since distinct initial codes differ by at least $\delta$, the first shifted code exceeds the largest possible initial code, and each subsequent shifted code exceeds its predecessor. Thus the recurrence of step 2.1 holds with $Q$ in place of $q$, and different matrices give different ordered tuples $({\mathrm{\ell }}_1,\dots ,{\mathrm{\ell }}_n)$. The divisibility argument in step 3.1 applies modulo $Q$: for each fixed position, the offsets $(j-1)S$ cancel and the difference of two $r_j$ values divided by $\delta$ has magnitude at most $q-1<Q$. Thus the final maxima $t$ distinguish all position-tagged matrices, and the codes $Kt+{\tilde{\mathrm{\ell }}}_j$ distinguish every token in every matrix when $K\delta >\max t$. Choose $K$ still larger so these positive codes exceed all finitely many target values $\mathrm{\mid }{u}^{\mathsf{T}}A(L{)}_{:,j}\mathrm{\mid }+\delta \mathrm{/}3$. The pulse layers of step 6.1 then implement the arbitrary table $A$; no permutation condition is needed because the position offsets distinguish columns.

Each update used above is either a two-head attention update, a single-head attention update with the second head set to zero, or a tokenwise scalar piecewise-affine update; an unused attention or feedforward sublayer is made the identity by setting its output weights to zero. Therefore the finite constructions are compositions of the modified Yun blocks specified in the Statement, proving both lookup assertions.