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

TL;DR

STILL introduces a novel intra-layer hybrid attention method for linearizing large language models, improving token selection accuracy and efficiency while maintaining performance on reasoning tasks and long-context benchmarks.

Contribution

The paper proposes STILL, a new intra-layer hybrid linearization framework with a self-saliency score and norm-preserved feature map, enhancing token selection and preserving pretrained representations.

Findings

Matches or surpasses original model on reasoning tasks

Achieves up to 86.2% improvement on long-context benchmarks

Efficient hardware implementation with chunk-wise parallelization

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.