Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling

背景

现代大型语言模型（LLM）在处理长序列时面临两个核心瓶颈：一是稠密注意力（dense attention）的计算复杂度随序列长度呈平方增长（O(n²)）；二是模型在超出训练上下文长度时难以有效泛化（即长度外推能力差）。为突破这些限制，分块稀疏注意力（chunk-wise sparse attention）被提出作为一种替代方案——它将输入序列划分为若干块，仅让查询（query）与少量选中的块进行注意力计算，从而降低复杂度。然而，现有分块稀疏方法在块选择（chunk selection）上不够准确，导致性能无法达到全注意力的水平。针对这一不足，本文提出了一种新的层次化分块稀疏注意力机制——Hierarchical Landmark Sparse (HiLS) Attention，旨在实现高效且精准的长上下文建模。

核心内容

HiLS Attention 是一种端到端学习块选择的分块稀疏注意力机制。其核心设计思想是：将注意力层次化地分解。具体而言，每个查询独立地与每个被检索到的块进行注意力计算，以提取该块特有的信息；然后，根据每个块的检索得分（retrieval score），将这些信息融合成最终的输出。通过将检索得分直接纳入前向注意力计算，HiLS 能够利用语言建模损失（LM loss）对检索得分进行端到端优化，从而实现原生的稀疏训练（native sparse training）与检索学习。

实验结果表明：

• 在域内上下文长度（即训练时所见长度）上，HiLS Attention 的性能与全注意力相当，甚至在某些情况下更优。
• 在外推能力方面，HiLS Attention 能够外推至训练上下文长度的 64 倍以上，且块检索准确率高达 90%，远超全注意力模型。
• 现有的全注意力模型可以通过轻量级的持续预训练（lightweight continued pretraining）转换为 HiLS Attention 模型，在保持域内性能的同时获得超长上下文的泛化能力。

由于 HiLS 在注意力计算中仅需稀疏的键值（KV）访问和稀疏计算，它打破了以往效率与性能之间的权衡，使得长上下文 LLM 既能比全注意力模型更高效，又能在通用长上下文任务上表现更优。

关键要点

• 核心挑战：稠密注意力的二次复杂度与弱外推能力限制了 LLM 的长上下文处理。
• 现有问题：分块稀疏注意力方法虽降低计算，但块选择不精确导致性能下降。
• HiLS 设计：层次化注意力——每个查询独立与每个检索块计算注意力，再按检索得分融合；检索得分通过 LM 损失端到端学习。
• 优势一（域内性能）：HiLS 在训练长度上达到甚至超越全注意力的效果。
• 优势二（外推能力）：可外推至训练长度的 64 倍以上，检索准确率 90%。
• 优势三（模型转换）：现有全注意力模型仅需轻量持续预训练即可转换为 HiLS，兼顾域内性能与超长外推。
• 效率-性能平衡：稀疏 KV 访问与计算使 HiLS 在效率与效果上均优于全注意力模型。

意义与影响

HiLS Attention 的提出为构建能够处理无限上下文的 LLM 提供了切实可行的新范式。它不仅解决了现有稀疏注意力方法中块选择不准确的根本问题，还通过端到端训练统一了检索与注意力，使得稀疏注意力在性能上首次不输稠密注意力。更关键的是，HiLS 展现了超乎寻常的长度外推能力（64 倍以上），这意味着模型可以在极短训练上下文下泛化到极长序列，大幅降低训练长上下文模型所需的计算资源。同时，对现有全注意力模型的轻量转换路径，使得已有预训练模型可以快速获得超长上下文能力，而无需从头训练。这一工作有望推动长上下文 LLM 在实际应用（如文档理解、代码分析、对话记忆等）中的部署，并可能引发对注意力机制设计范式的重新思考——即如何将稀疏性与端到端可学习性有机结合，以兼顾效率与效果。