Alleviating Forgetfulness of Linear Attention by Hybrid Sparse Attention and Contextualized Learnable Token Eviction

TL;DR

Linear attention offers $O(1)$ per-step computation but suffers from forgetfulness when past information is compressed into a fixed-size state. The paper proposes contextualized Learnable Token Eviction (LTE) and two hybrids, laLTE and laNSA, to restore selective access to past tokens while preserving constant-time, constant-space operation. LTE uses a lightweight CNN to predict per-head retention scores and drives end-to-end eviction under a budget, while laNSA leverages Native Sparse Attention for input-adaptive selection; both are implemented with efficient Triton/Flash Attention kernels. Across short-context and retrieval-intensive benchmarks, the hybrids substantially close the gap to full transformers under memory and compute constraints, demonstrating the practical value of learned past-token access for retrieval tasks.

Abstract

Linear-attention models that compress the entire input sequence into a fixed-size recurrent state offer an efficient alternative to Transformers, but their finite memory induces forgetfulness that harms retrieval-intensive tasks. To mitigate the issue, we explore a series of hybrid models that restore direct access to past tokens. We interleave token mixers with intermediate time and space complexity between linear and full attention, including sparse attention with token eviction, and the query-aware native sparse attention. Particularly, we propose a novel learnable token eviction approach. Combined with sliding-window attention, an end-to-end trainable lightweight CNN aggregates information from both past and future adjacent tokens to adaptively retain a limited set of critical KV-pairs per head, maintaining linear attention's constant time and space complexity. Efficient Triton kernels for the sparse attention mechanisms are provided. Empirical evaluations on retrieval-intensive benchmarks support the effectiveness of our approaches.

Figures (5)

• Figure 1: Hierarchy of token mixers across different time and space complexity per step. More complete and direct access to past tokens entails higher time and space costs.
• Figure 2: Left: Illustration of our KV eviction scheme, combining LTE, an attention sink (size $s=2$), and SWA ($w=12$), with 4 KV heads. A CNN stack is trained end-to-end to predict a per-token, per-head retention score $r_{j,h}$ to decide whether a KV-pair will be evicted when moved out of SWA. With recent KVs cached by SWA, the CNN can read short-range past and future context. $r_{j}$ can be computed once the CNN receptive field is fully covered by SWA. Right: Resulting per-head A-plus-column sparsity pattern; colored tiles indicate selected tokens at each step.
• Figure 3: Illustration of the KV cache layout in LTE inference, coloring identical to Figure \ref{['fig:LTE_token']}: a circular SWA segment for in-window KVs, and an out-of-window segment with capped capacity for LTE-retained KVs and the attention sink. LTE scoring is batched and deferred as late as possible.
• Figure 4: Comparing time costs between LTE and SWA or full attention with Flash Attention-2, on (a) kernels with an artificial KV cache during prefilling; and hybrid layers with real inputs during (b) prefilling and (c) decoding with a 4K-token prompt, using our 1.4B models with batch size=32.
• Figure 5: Average retention rate per head per layer produced by a laLTE model.