2D Gated Linear Attention with Per-Head Kronecker Manhattan Gate

Block context in the full architecture

At spatial position $(i,j)$, let the incoming feature be $x_{ij}\in\mathbb{R}^{d_{\mathrm{model}}}$. For each head $h\in\{1,\dots,N_h\}$:

with $q^{(h)}_{ij},k^{(h)}_{ij},v^{(h)}_{ij}\in\mathbb{R}^{d_h}$ and typically $d_{\mathrm{model}}=N_h d_h$. After computing head outputs $o^{(h)}_{ij}$, concatenate heads, apply an output projection, and use residual + normalization in the surrounding transformer block.

Grid and features

Let the input be a 2D grid of size $H\times W$ (for us $H=W=50{,}000$). For each position $(i,j)$ and each head $h\in\{1,\dots,N_h\}$, define

Define the rank-1 update matrix (per head)

We will construct a position-dependent state $S^{(h)}_{ij}\in\mathbb{R}^{d_h\times d_h}$ and output

Per-head learned Manhattan gate (Kronecker separable)

For each head $h$, learn two scalars

e.g. $\rho^{(h)}_x=\sigma(\theta^{(h)}_x)$ and $\rho^{(h)}_y=\sigma(\theta^{(h)}_y)$.

Define 1D distance matrices

The induced 2D gate over the flattened grid is the Kronecker product

Equivalently, on a grid of matrices,

meaning entrywise

Finally,

Thus each $o^{(h)}_{ij}$ depends on all $HW$ sites through the state $S^{(h)}_{ij}$ (global receptive field).

Directed (triangular) decomposition for implementation

Define directed (triangular) 1D matrices using indicator functions (no piecewise blocks):

Then the absolute-distance matrices satisfy

Therefore the target 2D operator is

This defines the exact kernel $(\rho_x^{(h)})^{|i-k|}(\rho_y^{(h)})^{|j-\ell|}$.

Scan (RNN) realization of directed operators

Fix a head $h$.

Forward $x$ scan (left-to-right) on each column $j$:

Backward $x$ scan (right-to-left) on each column $j$:

Forward $y$ scan (top-to-bottom) on each row $i$:

Backward $y$ scan (bottom-to-top) on each row $i$:

These scans implement the action of $\Gamma^{(h)\rightarrow}_x$, $\Gamma^{(h)\leftarrow}_x$, $\Gamma^{(h)\downarrow}_y$, $\Gamma^{(h)\uparrow}_y$ on the grid of updates.

Exact separable composition (no ad-hoc mixing)

To realize the Kronecker-separable operator $S^{(h)}=\Gamma^{(h)}_x U^{(h)} \Gamma^{(h)\top}_y$, we apply the $y$-operator and $x$-operator by composition:

In practice, each application of $\Gamma^{(h)}_x$ or $\Gamma^{(h)}_y$ uses the bidirectional decomposition:

Thus the full absolute-distance kernel is implemented exactly by bidirectional scans (two per axis) plus the identity correction.

Scan count and parallelism accounting

The composed bidirectional implementation uses exactly four directed scans per head:

This is the same directed-scan count as the old four-corner realization. The key difference is execution order: pass 2 depends on the materialized/intermediate $M^{(h)}$, so the two passes are sequential with respect to each other. Parallelism is internal to each pass (all scan lines and heads can run in parallel), not across the two passes.

Chunked implementation of a directed scan (GLA-style)

We describe one directed scan on a length-$N$ line with constant $\rho\in(0,1)$:

where $U_t\in\mathbb{R}^{d_h\times d_h}$.

Choose chunk size $C$ and let $m=N/C$. Chunk $b$ contains indices $t=(b-1)C+1,\dots,bC$. Define the chunk decay

and the chunk effective update

Then the state at chunk boundaries satisfies the exact chunk recurrence

which is a scan over $m$ chunks (small).

For any $t$ inside chunk $b$ with $t_0=(b-1)C$, the exact decomposition is

The first term uses the incoming chunk-boundary state $S_{t_0}$; the second term is a local triangular computation of size $C$.

This is the basis for GPU efficiency:

• compute $B_b$ for all chunks in parallel across many lines (rows/cols) and heads,
• scan across chunks to obtain boundary states $S_{bC}$,
• compute within-chunk outputs using the above local formula.

Tiling a $50k\times 50k$ slide (exact global coupling)

We process the slide in spatial tiles of size $T\times T$ for memory reasons, but we do not reset scan states at tile boundaries.

For a directed scan (e.g. $x\rightarrow$), the sufficient statistic to continue the exact same computation in the adjacent tile is the boundary state on the shared edge. For example, for each global row index $j$, the state at the right edge of a tile is

Passing this boundary state to the neighboring tile on the right yields exactly the same result as scanning the full length-$W$ row in one shot (up to floating-point associativity). Analogous boundary states exist for $x\leftarrow$, $y\downarrow$, and $y\uparrow$. Hence the computation is streamable over tiles and distributable across GPUs while preserving the full-slide receptive field defined by the Kronecker Manhattan gate.

For the composed pipeline, tiling is a two-pass stream over the slide:

• pass 1 ($y$ filtering): stream tiles in column-major order, propagate $y$-scan boundary states, and materialize tilewise $M^{(h)}$,
• pass 2 ($x$ filtering): stream tiles in row-major order over $M^{(h)}$, propagate $x$-scan boundary states, and produce $S^{(h)}$.

So there are two full tile-streaming passes. Boundary state for one directed scan is one $d_h\times d_h$ matrix per scan line per head. Therefore, boundary storage scales as $O(H\,N_h\,d_h^2)$ for one $y$ direction and $O(W\,N_h\,d_h^2)$ for one $x$ direction. At $H=W=50{,}000$, $N_h=8$, $d_h=64$, this is

which is about $6.55$ GB in float32 (about $3.28$ GB in float16/bfloat16) per stored direction.

Normalization and stability at very large spatial extent

Without normalization, the state magnitude can grow large when $\rho_x^{(h)},\rho_y^{(h)}\to 1$. For large grids, total gate mass is approximately

so for $\rho_x^{(h)}=\rho_y^{(h)}=0.9999$, this scale is about $10^8$.

Practical stabilizers (one or more):

• denominator normalization analogous to linear attention, e.g.

  \[ z^{(h)}_{ij}=q^{(h)\top}_{ij}\!\!\sum_{k,\ell}\gamma^{(h)}_{(i,j),(k,\ell)}k^{(h)}_{k\ell}, \qquad \tilde{o}^{(h)}_{ij}=o^{(h)}_{ij}/(z^{(h)}_{ij}+\varepsilon), \]

• post-head or post-block normalization (e.g. LayerNorm/RMSNorm),
• constrain $\rho_x^{(h)},\rho_y^{(h)}$ away from $1$ (or regularize effective window size).

At $50k$ resolution this is a core training-stability issue, not only a numerical detail.

Final output (per head)

After computing $S^{(h)}_{ij}$, output

and combine heads in the usual way (concatenate or sum, then a linear projection).