STILL: Content-Aware Intra-Layer Hybrid Attention for Efficient LLM Linearization

Introduction

Softmax attention's quadratic complexity remains the primary bottleneck for long-context LLMs. Linear attention offers linear scaling but training from scratch requires thousands of GPU hours. Linearizing existing pre-trained LLMs is the most cost-effective path, yet current methods suffer from two critical flaws: (1) they route tokens by fixed positional windows rather than content importance, and (2) their feature maps distort the pre-trained model's norm structure, causing catastrophic long-context degradation.

Problems with Existing Intra-Layer Hybrid Attention

Prior work (LoLCATs, Liger, KL-LS) uses sliding-window routing: tokens inside a window receive Softmax attention (SA), outside tokens are compressed by linear attention (LA) or discarded (CSA). This positional heuristic assumes key dependencies are local, but critical evidence often lies far beyond the window. Once an important token is pushed to the LA branch, its representation is irreversibly compressed. Additionally, mainstream linearization (e.g., Hedgehog) inserts learnable MLPs before the feature map, freely rescaling feature norms. Since pre-trained Softmax attention is norm-aware — vector magnitude modulates representation strength — this rescaling breaks alignment with the original representation geometry.

STILL Architecture

STILL (Selecting Tokens for Intra-Layer Hybrid Attention) addresses both issues with three core components:

Content-Aware Routing with Fixed Budget

Each attention head receives a fixed KV cache budget m independent of sequence length. W slots are reserved for the most recent sliding-window tokens (preserving fine-grained local interaction). The remaining m−W slots are allocated to historically salient tokens selected by a content-based score. Each completed chunk retains λ = ⌊(m−W)/T_max⌋ tokens permanently; unselected tokens are accumulated into the linear attention recurrent state. This replaces positional splicing with content-driven competition for the "far" slots.

Delayed Chunk-Level Selection for Unified Training-Inference Parallelism

Per-token routing breaks parallelism. STILL adopts delayed chunk-level selection: when a new chunk begins, the previous chunk's tokens are scored and the top- λ are chosen. Both training (with chunked parallelism) and inference (prefill and decoding) share identical routing rules and memory updates, ensuring a unified architecture.

Self-Saliency Score: Predicting Future Utility from Local Context

The routing signal answers: Will a token, once it becomes distant history, still be attended by future queries? Global attention answers this but is quadratic. STILL's insight: this can be predicted cheaply inside the local window.

For each token, STILL computes two sliding-window attention distributions: a^diag (including self-interaction) and a^nodiag (excluding self-interaction). The Self-Saliency Score is the KL-type sensitivity between them:

Intuitively, a token whose attention distribution is highly sensitive to its own presence (strong self-interaction) tends to maintain high global attention rank. High-score tokens go to SA; the rest to LA. Routing is driven by content, not position.

Theoretical Guarantee: Local-Global Consistency

Proposition 1 : As ε → 0⁺, the Self-Saliency Score satisfies (1) monotonicity: ∂Score_t/∂p_local > 0; (2) order-preservation: Score_i > Score_j implies E[p^(i)_global] > E[p^(j)_global]. Thus local scores preserve the expected global ranking order.

Empirical Validation

Visualization shows tokens highlighted by local scores remain high under full attention. Overlap between locally selected top-224 tokens and globally attended distant keys reaches 62.9%–84.8% across shallow, middle, and deep layers. Word-cloud analysis reveals SA most frequently selects high-scope function words: negations (not, no, never), modals (will, would, may), and quantifiers (all, any, nothing) — precisely the tokens requiring precise, context-dependent modeling.

NP-Map: Norm-Preserving Feature Map

To prevent norm distortion, STILL decouples direction and magnitude. An MLP f(⋅) refines only the direction; the original pre-trained norm ‖x‖ is re-injected:

This ensures the linear kernel strictly aligns with pre-trained representation strength while satisfying non-negativity. A gating mechanism further boosts the linear branch's rank.

Experiments

Setup

Two-stage pipeline (Attention Transfer + low-rank linearization) on Cleaned Alpaca, 20M tokens, sequence length 1024. Total cost: 0.04B tokens . Backbones: Llama 3 8B, Llama 3.1 8B, Llama 3.2 1B/3B, Mistral 7B.

Commonsense & General Reasoning

On PIQA, ARC-E/C, HellaSwag, WinoGrande, MMLU, STILL achieves the best average among linearization methods. On Llama 3.1 8B, MMLU rises from Liger's 46.9 to 61.3 ; non-MMLU average 74.7 surpasses the original full-attention model (73.2) by 1.5 points.

S-NIAH Needle-in-a-Haystack

Under equal or smaller cache budgets, STILL dominates 0.5K–4K lengths. At 4K: STILL 86.2 vs LoLCATs 8.8, Liger 0.0, SWA 5.8 — a +77.4</sup> gain over the strongest baseline.

RULER Long-Context Benchmark

Eight tasks requiring long-range memory and state tracking (4K context). STILL (128 window + 896 total cache) scores 47.9 vs LoLCATs 7.2, Liger 2.1, Mamba2 44.7.

BABILong Reasoning

Strict 256-token total cache. STILL sweeps all lengths: QA1 4K from 3 to 10 , QA5 from 6 to 24 .

Compatibility with Layer-Wise Hybrid (KL-LS)

Embedding STILL into KL-LS (Llama 3.2 1B, 25% full-attention layers) yields perfect 100 on S-NIAH-1 at 4K/8K/16K, S-NIAH-2 74.0 at 16K (LoLCATs 8.6), RULER total 27.20 (Zebra-Llama 13.37).

Efficiency

Measured on single H200, 1K–64K:

Prefill : Softmax OOM after 16K; STILL scales to 64K. At 16K, latency is 17% of Softmax.

Decoding : Softmax OOM after 32K; STILL memory near constant, at 16K only 20% of Softmax, stable throughput.

Ablation Studies

Module ablation (BABILong 1K) : removing gating −5.8, removing NP-Map −4.9, both −8.3; further removing self-saliency selection drops accuracy from 27.5 to 12.5 (−15.0).

Routing metric comparison (S-NIAH) : Self-saliency 42.6 vs diagonal attention 23.1, q·k dot product 22.8, key norm 12.9, uniform 12.3, random 12.8. Diagonal attention approximates short context but fails in long-range sparse retention.

Chunk size : C=64 optimal; C=16→128 reduces training time 35.4% and memory 21.7%; per-token routing (C=1) triples decoding selection latency and halves training speed.

Backbone generalization : Llama 3.2 1B/3B and Mistral 7B all show superior results (3B MMLU 52.4 vs Liger 32.1; Mistral 7B average 70.9 vs LoLCATs 69.1, Liger 65.1).

Conclusion

STILL redefines true intra-layer hybrid attention by (1) routing tokens via content importance rather than position, with fixed cache budget and unified chunk-parallel training/inference; (2) predicting future utility of historical tokens through locally computed Self-Saliency Scores, backed by local-global consistency theory and 62.9%–84.8% empirical overlap; (3) preserving pre-trained norm statistics via NP-Map's direction-magnitude decoupling. With only 0.04B training tokens, STILL not only surpasses the original full-attention model on commonsense reasoning but also resurrects long-context capabilities previously collapsed in linearized models (S-NIAH-1 perfect 100 at 16K), while delivering 80% memory reduction and 83% prefill speedup. Code: https://github.com/ZacharyMeng/STILL, Paper: https://arxiv.org/abs/2602.02180