Linear Transformers with Learnable Kernel Functions

Attention without the quadratic bill

Full attention computes a similarity score between every query and every earlier key, $\operatorname{sim}(q_i,k_j)=\exp(q_i^{\top}k_j/\sqrt{d})$, so the cost of a forward pass grows as $\mathcal{O}(N^2)$ with sequence length $N$. A linear transformer swaps the exponential for a factorized kernel, $\operatorname{sim}(q,k)=\phi(q)^{\top}\phi(k)$. Once the score factorizes, the sums re-associate: everything the model has read collapses into a running state $S_i=\sum_{j\le i}\phi(k_j)v_j^{\top}$, and the output at position $i$ is a read of that state, $y_i\propto\phi(q_i)^{\top}S_i$.

The price of that speed is paid entirely by the feature map. Whatever $\phi$ keeps, the state keeps; whatever $\phi$ fails to suppress lingers as noise for every later query. Choosing the kernel is not an implementation detail. It is the mechanism.

We build on Based, a linear-attention hybrid introduced by Arora and colleagues that mixes short convolutions with a second-order polynomial kernel. Before our work it was the strongest subquadratic performer on associative recall, and it still lost ground to full attention as contexts grew and hidden states shrank. That remaining gap is what we set out to understand.

The probe that separates recall from fluency

Perplexity is a blunt instrument for this question, because a model can lower it with local pattern-matching while never learning precise long-range routing. Multi-Query Associative Recall, or MQAR, is the sharper probe: a sequence writes down key-value pairs, distractors accumulate, and much later the keys reappear as questions.

The probe is unforgiving at scale. At sequence length 2,048 with a hidden size of 64, every subquadratic architecture we tested collapses on MQAR while full attention still solves it; in our runs Mamba and RWKV fail the task across all tested sizes. Whatever separates these architectures, this task measures it directly.

A kernel that cannot say zero

Based scores a pair with a truncated Taylor expansion of the exponential, $\operatorname{sim}(q,k)=1+q^{\top}k+\tfrac{1}{2}(q^{\top}k)^2$. Written in the dot product $t=q^{\top}k$, this is a parabola whose minimum sits at $t=-1$, and the minimum value is $0.5$. No pair of tokens, however unrelated, can score below it.

Across a few dozen positions that floor is a rounding error. Across two thousand it is a chorus: every irrelevant token contributes at least half a unit of unnormalized weight, and the one token that should be retrieved must outshout the crowd. Subtracting the constant would drop the floor to zero, but the root would stay pinned at $t=-1$, an arbitrary point that queries and keys would then have to organize their whole geometry around.

Let training place the parabola

In ReBased we stop hand-picking the curve. We reduce the kernel to an element-wise square and put a learnable affine map in front of it, $\phi(x)=(\gamma\,\operatorname{norm}(x)+\beta)^2$, with separate $\gamma$ and $\beta$ for queries and keys. The result is still non-negative, so the linear-attention rearrangement survives untouched, but training now decides where each parabola touches zero.

Normalization is the quiet half of the trick: it pins the feature scale so a single learned $\gamma$ and $\beta$ fit the whole distribution. Our ablation pulls the pieces apart. A raw square is actually worse than Based at length 2,048, with mean accuracy 0.44 against 0.51. Normalization alone lifts it to 0.58, the affine map alone reaches 0.60, and the full combination gets 0.61 there and 0.63 at length 256.

Recall improves where the floor used to bite

The same ordering holds across our full sweep, sequence lengths 128 to 2,048 and hidden sizes 64 to 512: ReBased matches or beats every subquadratic baseline we tested, with the margin widest exactly where the kernel floor bites, long sequences and small states. Full attention remains ahead, and at length 2,048 with hidden size 64 it is the only architecture that still solves the task.

Attention maps say the same thing quantitatively. Over 10,000 MQAR examples we measured intersection-over-union between each model's attention pattern and the ground-truth retrieval positions: full attention scores 0.999, an almost exact lookup. Based reaches 0.157 and ReBased 0.173, sharper than its predecessor but still a diffuse approximation of the routing that attention performs exactly.

Real text moves the same direction

On the Pile we split perplexity by whether a token requires recalling something seen earlier in the context. ReBased improves both slices: associative perplexity moves from 3.27 to 3.25 and the non-associative rest from 37.02 to 36.73, for 12.90 overall against 12.99 for Based. Full attention sits at 11.98, so the ranking that MQAR predicted transfers to real text at the 160M-parameter scale we trained.

One-shot evaluation moves the same direction. The LM Evaluation Harness macro average rises from 41.3 to 42.1, and every SuperGLUE task we report improves, with the largest jumps on RTE and BoolQ. The margins are small but uniformly signed: the kernel change does not trade general capability for recall.

What we did not claim

We tested academic-scale models, and we make no production claims: no inference-latency numbers, no evidence about behavior after much longer pretraining, and copying-heavy workloads remain untested. The recurrent state also still scales with hidden size, so linear in sequence length does not mean free.

What we can defend is narrow and useful. A fixed kernel imposes an avoidable floor on how thoroughly a linear transformer can ignore a token; letting training place the zero removes that floor and reliably helps recall. The distance that remains to full attention is not fine print to hide. It is the measurement that tells us where the next kernel should go.