STILL: Selecting Tokens for Intra-Layer Hybrid Attention to Linearize LLMs

TL;DR

STILL tackles the quadratic bottleneck of pretrained LLM attention by introducing a content-aware, in-layer hybrid approach that selects salient tokens for Softmax Attention (SA) and summarizes the rest with Linear Attention (LA). Key ideas are a Self-Saliency Score that maintains local-global routing consistency and a Norm-Preserved Feature Map (NP-Map) that reinjects pretrained norms to preserve representation geometry, all within a chunk-wise, delayed selection framework to boost hardware efficiency. Empirically, STILL matches or surpasses full-attention baselines on commonsense and general reasoning tasks and achieves large gains on long-context benchmarks like RULER and BABILong, while reducing memory and improving decoding throughput. The combination of self-saliency-based routing, norm-preserving linearization, and chunk-wise computation yields a scalable, long-context-capable LLM linearizer suitable for deployment with LoRA fine-tuning and beyond.

Abstract

Linearizing pretrained large language models (LLMs) primarily relies on intra-layer hybrid attention mechanisms to alleviate the quadratic complexity of standard softmax attention. Existing methods perform token routing based on sliding-window partitions, resulting in position-based selection and fails to capture token-specific global importance. Meanwhile, linear attention further suffers from distribution shift caused by learnable feature maps that distort pretrained feature magnitudes. Motivated by these limitations, we propose STILL, an intra-layer hybrid linearization framework for efficiently linearizing LLMs. STILL introduces a Self-Saliency Score with strong local-global consistency, enabling accurate token selection using sliding-window computation, and retains salient tokens for sparse softmax attention while summarizing the remaining context via linear attention. To preserve pretrained representations, we design a Norm-Preserved Feature Map (NP-Map) that decouples feature direction from magnitude and reinjects pretrained norms. We further adopt a unified training-inference architecture with chunk-wise parallelization and delayed selection to improve hardware efficiency. Experiments show that STILL matches or surpasses the original pretrained model on commonsense and general reasoning tasks, and achieves up to a 86.2% relative improvement over prior linearized attention methods on long-context benchmarks.

Figures (6)

• Figure 1: Illustration of Different Hybrid Linear Attention Designs. LoLCATs assigns tokens within a local window to softmax attention (SA), while the remaining tokens are routed to linear attention (LA) in a position-biased manner. Liger also routes tokens within the local window to SA, but LA handles all global context. In contrast, our approach not only applies SA to tokens within the window, but also selectively promotes certain tokens outside the window to SA, while the remaining tokens are processed using LA.
• Figure 2: Architecture of STILL for Intra-Layer Hybrid Linear Attention and Train Strategy. a) The diagram shows attention computation, chunk-level routing, and hybrid aggregation to outputs. STILL first computes the Self-Saliency Score within sliding-window attention for each token. Top-scoring tokens are routed to the softmax attention, while remaining tokens are processed via linear attention. Chunk-wise selecting replaces per-token selecting, enabling parallel training and inference. b) STILL follows the two-stage procedure.
• Figure 3: Distribution of Self-Saliency Scores across Input Tokens. Red points in the left panel mark tokens with the highest scores under sliding-window attention (routed to SA), while red points in the right panel highlight the same token indices as in left, illustrating strong local-global consistency.
• Figure 4: Word Cloud of Tokens Most Frequently Selected for the SA Branch by Self-Saliency Score. The selected set is dominated by functional operators (e.g., negation, modality, and quantification), suggesting that SA is preferentially allocated to tokens whose contribution depends on precise, context-specific interactions.
• Figure 5: Prefilling Efficiency. STILL demonstrates favorable prefill latency and memory consumption from 1K to 64K.