What are we trying to model?

In structured linear attention, each head maintains a recurrent state that looks like an outer-product memory:

\[S_t \approx \sum_{\tau \le t} g_{\tau:t} v_\tau k_\tau^\intercal\]

where each token contributes a value-key outer product, and model later reads from state using a query.

If the head dimension is $d$, then the recurrent state is roughly $d \times d$. Increasing head size from $64 \to 256$ makes state matrix $16 \times$ larger.

We can compromise by keeping $64$-dimensional sub-blocks but arrange them into a richer tiled state.

Modeling Idea

These are ways of spending head dimension:

  • d_head = 256: full wide head. have many key/value subcomponents interacting with each other. This preserves cross-interactions $v_i/k_j$ for many combinations. Note n_head$=1$.
  • d_head=64: much cheaper but only keeps a diagonal-style interaction. The risk is that it throws away cross-subspace interactions. n_head$=4$.
  • d_head=64 x 4: we try to recover some of the lost expressivity of d_head$=256$ using smaller subheads. \[T = \sum_t k_t^{(1)} \otimes k_t^{(2)} \otimes v_t\]

    which makes the state more like a 3D tensor. We give each key/value channel extra second-key slots, making it behave like 256 structured features.

    • This is better than plain d_head=64 as it reduces collision / information mixing inside each small head. With normal d_head=64, the state is roughly \[S = \sum_t k_t \otimes v_t\]

      so two different token types that have similar $k_t$ will have values mixed into the same slot.

      Now the token is indexed by the pair $(k_i, a_j)$, so instead of $1$ slot $k_1 \rightarrow v$, we get multiple structured slots $(k_1, \{a,b,c,d\})$.

    • It gets some of the effective address space of a 256-wide head, but in a cheaper and more factorized form. This can reduce memory/compute cost and may act like useful regularization, because the model is forced to represent interactions through products of smaller key spaces rather than an unconstrained dense key space.

      State size of d_head=256$=65536$, d_head=64$=4 \cdot 64^2 = 16384$, and d_head 64 x 4$= 64 \cdot 4 \cdot 64 = 16384$.

    Note for a query, we have $q^{(1)} \in \mathbb{R}^{64}$ and $q^{(2)} \in \mathbb{R}^4$. The output, using $T \in \mathbb{R}^{64 \times 4 \times 64}$ where $T_{i,j,r} = \sum_{t} k_{t,i}^{(1)} k_{t,j}^{(2)} v_{t,r}$, is

    \[y_r = \sum_{i=1}^{64} \sum_{j=1}^4 q_i^{(1)} q_j^{(2)} T_{i,j,r}\]

    Substituting in the definition of $T$,

    \[ y_r = \sum_{i=1}^{64} \sum_{j=1}^{4} q_i^{(1)} q_j^{(2)} \sum_t k^{(1)}_{t,i} k^{(2)}_{t,j} v_{t,r}\]

    Rearranging the sums gives

    \[ y_r = \sum_t \left( \sum_{i=1}^{64} q_i^{(1)} k^{(1)}_{t,i} \right) \left( \sum_{j=1}^{4} q_j^{(2)} k^{(2)}_{t,j} \right) v_{t,r}\]

    So in vector form,

    \[ y = \sum_t \Big(q^{(1)} \cdot k_t^{(1)}\Big) \Big(q^{(2)} \cdot k_t^{(2)}\Big) v_t. \]

Core Hypothesis

IMPORTANT: Constant state size is sufficient for long context, but constant needs to be big enough

In transformers, we can pretrain on short context ($\sim4k$), then continued pretraining (CPT) on long context ($\sim 1M$ or $256\times$ extension)

  • Essentially pretrain with state size $K$, then long context extension with state size $256*K$.

Can we do something similar for Linear Attention?

Pretrain with state size $M$ $\rightarrow$ Long context extension with state size e.g. $16 \cdot M$.

  • Can use 2nd key/query in 3D Linear Attention for good init (e.g. one-hot 2nd keys is an initializer from 2D).

6/10

Cross Pairing vs Diagonal Pairing

\[ \text{\textbf{Cross-pairing} means the state keeps all combinations: }v_i k_j \quad \text{\textbf{Diagonal pairing} keeps only matched pairs: }v_i k_i \]
NOTE Motivation is if cross-pairing helps, model needs interactions between different subspaces, else cheaper structured version is probably sufficient

Head-Size vs Recurrent-State-Size Tradeoff

Diagonal pairing (3D state) seems to be better than cross pairing (bigger headsize).

Setup: GDN with various head sizes,  200M params, 10B tokens

Increasing head_dim improves validation loss, but the gains saturate quickly. For GDN, h64 and h128 are already within about \(+0.10\%\) of h256, despite using much smaller recurrent states.

6/17

KV-Pairing, Gating Granularity, Second Key/Query

We want to see how to get a $4 \times$ larger state from hs = 64.

Tested on three axis:

  • KV-pairing: all $k_i$ with $v_j$ combinaions (corss pairing) or $k_i$ with $v_i$ 4 times (diagonalKV).
  • Gating Granularity: every $64 \times 64$ tile separate forget gate or single forget gate for entire $256 \times 256$.
  • Second Key/Query choice: OnesQK: $q^{(2)} = k^{(2)} = 1$, sigmoidQK: $q^{(2)}, k^{(2)} \in (0,1)^E$, softmaxQK: learned with $\sum_{j} q_j^{(2)} = \sum_{j}k_j^{(2)} = 1$.
ON GLA

Takeaways:

  • expanded $196k$-state variants beat small $49k$ hs=64 baseline
  • tile-wise forget gate is better: red beats blue.
  • CrossKV is not obviously better despite being denser. Comes from more interference, making DiagonalKV/3D-state competitive.
GDN

Pretrain GDN on 4k seq, extend to 64k

  • Pretrain GDN on short context: $4k$.
  • Extend to long context: $64k$, or $16\times$ longer sequences.
  • For the GDN long-context extension, try expanding recurrent state size by $2\times$, $4\times$, and $8\times$.
  • Compare against a transformer baseline extended with RoPE scaling.
RoPE scaling
RoPE scaling: rescale token positions before applying RoPE so a model trained on short context can handle longer context.
Example. Model trained on \(4\text{k}\) tokens, want \(1\text{M}\) tokens. Use scale factor \[ s = \frac{1\text{M}}{4\text{k}} = 256. \]

Original:

\[ \theta_i(p)=p\omega_i. \]

Position interpolation:

\[ p'=\frac{p}{s} \quad\Rightarrow\quad \theta_i'(p)=\frac{p}{s}\omega_i=\frac{\theta_i(p)}{s}. \]

Makes long context possible, but local position resolution gets compressed.

head_dim=128

6/24

Long Context Extension with State Size Extension

Pretrain GDN on 100B tokens with 4k seq. length.

Then do long context extension to 64k seq. length on 5B tokens with state size extension.

On 20B tokens:

It gets worse over time... (bug?)

Quadratic Linear Attention

16x smaller in state! Try a context increase < 2x to see if it matches performance...

With transformers under GQA architecture, this could have a even bigger gap.

7/1

Zero-Shot

Evaluated 1B GDN/Transformer models, which are pretrained on 100B tokens (4k seq length), then long context extended for 5B (64k seq length).

With proper long context mix (books, code, papers, pretraining), 0-shot degradation is minimal:

why expect degradation
long-context extension can cause forgetting / distribution drift. Continuing training on a different context-length setup, the model may become better at long-context behavior but slightly worse on standard short-context zero-shot tasks.
  • Pretrained GDN: $58.4$ -> Long Context GDN (with 1x,2x,..., 32x state extension): $57.8-58.0$
  • Pretrained Transformer: $57.0$ -> Long Context Transformer: $55.1$ (most decrease comes from a single task)

With books-only for long-context extension, degradation is much worse:

  • Long context extended GDN 16x: $51.1$
  • Long context extended Transformer: $50.9$

RULER

Books-only long-context finetuning teaches some long-context retrieval behavior, but not robustly. It improves easy/simple NIAH tasks, but many harder/multi-key/multi-value tasks remain weak.

NIAH scores from RULER.

Top 2 rows are long context extended GDNs, with different state size expansion. Bottom row left is long context extended transformer. bottom mid and right are GDN/transformer before long context extension.

ideas

Dynamic State Growth for Two-Query Linear Attention

The fixed two-query state is

\[ T_t = \sum_{\tau \le t} k_\tau^{(1)} \otimes k_\tau^{(2)} \otimes v_\tau, \]

with readout

\[ y_t = \sum_{\tau \le t} (q_t^{(1)} \cdot k_\tau^{(1)}) (q_t^{(2)} \cdot k_\tau^{(2)}) v_\tau . \]

To make the state grow dynamically, replace the fixed second-key axis with a Fenwick/log-linear memory hierarchy:

\[ \{S_t^{(0)}, S_t^{(1)}, \dots, S_t^{(L(t)-1)}\}, \qquad L(t)=O(\log t). \]

Each \(S_t^{(\ell)}\) is a normal fixed-size GDN/linear-attention state summarizing one temporal bucket:

\[ S_t^{(\ell)} \approx \sum_{\tau \in B_t^{(\ell)}} g_{\tau:t}^{(\ell)} v_\tau k_\tau^\top . \]

The second query becomes a level router:

\[ y_t = \sum_{\ell=0}^{L(t)-1} q_t^{(2)}(\ell) S_t^{(\ell)} q_t^{(1)}. \]

Thus the model keeps the original \(64\)-dimensional content-addressing computation, but the number of memory states grows as the context grows.