提出三项设计原则,让模型训练后可处理超长文本。
Understanding and Improving Length Generalization in Hierarchical Sparse Attention Models
- 用非线性编码器+专用标记提取块信息,提升检索能力。
- 通过绕行残差路径稳定融入全局信息,避免被局部信息覆盖。
- 预训练时强制稀疏选择,解决训练与测试长度不一致问题。
有效处理长上下文是语言模型的关键挑战。标准Transformer因二次复杂度和较差的长度外推能力受限,而滑动窗口注意力和状态空间模型则因固定记忆大小牺牲了对完整上下文的利用。分块稀疏注意力成为实现极端长度泛化的有前景范式,但其成功的核心机制尚不明确。本文通过统一框架与全面消融实验,揭示三个关键设计原则:(1) 使用具有表达力的非线性分块编码器及专用CLS标记生成可检索表示;(2) 采用绕行残差路径,稳定融合检索到的全局信息,防止被局部残差流覆盖;(3) 在预训练中强制稀疏选择,弥合训练与测试阶段的分布差距。我们提供了块内信息处理与标志性生成的理论依据。结合这些原则,我们在无需微调的情况下达到新基准,在RULER和BABILong数据集上将4K训练长度的模型成功外推至3200万标记。研究结果为构建未来高性能长上下文语言模型提供了清晰且实证支持的设计准则。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。