Understanding and Improving Length Generalization in Hierarchical Sparse Attention Models

TL;DR

The paper tackles the core challenge of long-context processing by dissecting chunk-based sparse attention. It argues that extreme length generalization hinges on three design principles: a non-linear Chunk Encoder with a dedicated CLS token for disentangled retrieval, a Bypassing Residual Path for stable integration of retrieved information, and enforced selection sparsity during pre-training to bridge train-test distribution gaps. Through a unified framework and thorough ablations, the authors demonstrate state-of-the-art training-free extrapolation from 4K to 32M tokens on RULER and BabiLong, supported by theoretical motivation and diagnostic analyses of retrieval and integration. The findings yield concrete, empirically-grounded guidelines for designing future long-context language models and reveal the critical balance between retrieval prominence and information integration. The work has practical impact on building scalable, long-context LMs with predictable zero-shot length generalization across diverse tasks.

Abstract

Effectively processing long contexts is a critical challenge for language models. While standard Transformers are limited by quadratic complexity and poor length extrapolation, alternative architectures like sliding window attention and state space models sacrifice the ability to effectively utilize the full context due to their fixed-size memory. Chunk-based sparse attention has emerged as a promising paradigm for extreme length generalization, yet the key architectural principles underpinning its success are not yet fully understood. In this work, we present a systematic dissection of these models to identify the core components driving their performance. Through a unified framework and comprehensive ablation studies, we demonstrate that a combination of three design principles is critical: (1) an expressive, non-linear Chunk Encoder with a dedicated CLS token to produce representations for retrieval; (2) a Bypassing Residual Path to stably integrate retrieved global information without it being overridden by the local residual stream; and (3) enforced selection sparsity during pre-training to bridge the train-test distribution gap. We provide a theoretical motivation for intra-chunk information processing and landmark generation. By combining these principles, we establish a new state-of-the-art for training-free length extrapolation, successfully generalizing models trained on a 4K context to 32 million tokens on RULER and BABILong. Our findings provide a clear and empirically-grounded set of design principles for developing future, highly-capable long-context language models.

Figures (5)

• Figure 1: DRT Architecture. The DRT hu2025efficientlengthgeneralizableattentioncausal model is primarily composed of Transformers with Sliding Window Attention (SWA), with layers divided into upper layers and lower layers. Lower decoders are composed of SWA followed by an FFN. At the architectural midpoint, a chunking layer splits representations from lower decoders into chunks, and encode them by chunk to build a global memory, each chunk of which contains a landmark vectors ($\mathbf{lmk}$) and encoded chunks ($\mathbf{E}$). Upper decoders each contains a local self-attention layer for local context, followed by an HSA to attend selected global memory and then an FFN.
• Figure 2: Design of Encoder: (a): Encoder w/o CLS (b): Encoder with a learnable CLS token.
• Figure 3: BabiLong Evaluation Results: Our best DRT model demonstrates robust long-context extrapolation, maintaining high accuracy up to 8M tokens. In stark contrast, baseline models exhibit clear limitations. Full-attention methods collapse to near-zero accuracy shortly after their training length. Prior sparse methods like Landmark Attention show better generalization but also fail decisively at longer contexts ($\geq$64k). Sliding Window Attention plateaus at a low accuracy consistent with random guessing, while Mamba2's performance degrades substantially.
• Figure 4: Analysis of Retrieval and Utilization Mechanisms. Subplots (a) and (b) show retrieval accuracy (Recall@Top-K) on the S-N and MQ-N tasks, respectively. The results highlight that the Bypass mechanism is the most critical component for successful retrieval, consistently outperforming variants without it. The learnable encoder and CLS token also provide noticeable improvements. Subplot (c) plots the conditional answer accuracy on the S-N task, given that the correct chunk was retrieved at a specific rank (by weight). Across all models, integration precision drops sharply as the rank of the correct chunk decreases. This demonstrates a stricter requirement: a chunk must not only be successfully retrieved (recall) but also be assigned a high-ranking weight to be effectively utilized by the model.
• Figure 5: RAMba Architectures. An MLP is inserted into upper decoder layer in RAMB-MLP as Mamba does not have explicit MLP layer.