💻 computer science

Long-Context Modeling with Dynamic Hierarchical Sparse Attention for Memory-Constrained LLM Inference

本文提出了动态分层稀疏注意力（DHSA），这是一种数据驱动框架，通过分层路由在线预测注意力稀疏性，从而在有限硬件上实现内存高效的长上下文大语言模型推理，同时保持接近稠密方法的精度，并相较于现有稀疏方法获得显著加速。

原作者： Siheng Xiong, Joe Zou, Faramarz Fekri, Yae Jee Cho

发布于 2026-05-29

📖 1 分钟阅读☕ 轻松阅读

原作者： Siheng Xiong, Joe Zou, Faramarz Fekri, Yae Jee Cho

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在试图阅读一本庞大的百科全书，以查找一个特定的事实，比如“秘鲁的首都是什么？”在标准的语言模型（LLM）中，计算机就像一位非常详尽但缓慢的图书管理员。为了回答你的问题，这位图书管理员会阅读百科全书的每一页，将其与你的问题进行比对，然后决定如何作答。

如果这本百科全书有 10 万页，那么图书管理员每回答一个问题都需要完成海量的工作。这不仅昂贵、缓慢，而且经常会导致计算机内存崩溃（就像试图一次性抱住 10 万本书）。

本文介绍了一种名为DHSA（动态分层稀疏注意力）的新方法。这相当于将那位图书管理员升级为一位聪明且自适应的侦探，他确切地知道应该跳过哪些页面。

以下是其工作原理，分解为简单的概念：

1. 问题：“二次方”瓶颈

本文指出，当前的人工智能模型遭受“二次方成本”的困扰。这意味着，如果你将文本长度加倍，计算机需要完成的工作量并不会仅仅加倍，而是会翻四倍。

类比：想象一下试图在人群中寻找一位朋友。如果有 10 个人，你只需查看 10 张脸。如果有 100 个人，你不仅仅要看 100 张脸；你必须查看每一个人，并将他们与彼此进行比对，以看清谁在跟谁说话。这种情况很快就会变得混乱且缓慢。

2. 旧方案：“僵硬的网格”

过去试图解决这一问题的方法是使用静态稀疏注意力。

类比：想象图书管理员决定只阅读每第 10 页，或者无论故事内容如何，只阅读每一章的第一页和最后一页。
缺陷：这就像使用饼干模具。有时重要的信息恰恰就在你切掉的地方！如果“针”（即答案）位于你决定跳过的书页部分，你就会失败。本文表明，当文本变得非常长时，这些僵化的方法往往会遗漏重要细节。

3. 新方案：DHSA（聪明的侦探）

DHSA 之所以不同，是因为它是动态且分层的。它不使用固定的规则；而是先“阅读”文本，以决定什么是重要的。

步骤 A：“分块”侦探（动态边界）

DHSA 不是将书籍切成等大小的切片（例如每片 10 页），而是查看内容。

类比：想象文本是一部电影。僵化的方法会将电影切成 10 分钟的片段，即使场景转换发生在第 9 分钟。DHSA 足够聪明，能够看到场景转换，并在故事转折的确切位置切断电影。它将属于一起的句子（如一个段落或一段代码块）分组为“块”。
工作原理：它使用一个小型、轻量级的辅助工具来扫描文本，并指出：“好的，这句话结束了一个想法，而新的一句开始了不同的话题。”它在那里画了一条线。

步骤 B：“摘要”策略（分层路由）

一旦文本被分组为这些智能块，模型还不会查看块内的每一个单词。

类比：想象你有 50 章。与其阅读每一章的每一个字，侦探首先会阅读章节摘要。它会问：“哪 5 章最有可能包含答案？”
过程：
1. 它为每个块创建一个“摘要”。
2. 它将你的问题与这些摘要进行比对。
3. 它挑选出看似最相关的几个“摘要”块。
4. 只有在此之后，它才会回过头来，阅读那些选定块中的具体单词。

4. 为什么这很重要

本文声称，这种方法解决了三个主要问题：

节省内存：因为模型只关注文本的一小部分（约 6% 到 12% 的单词），所以它可以将巨大的书籍（多达 10 万个单词）放入单个标准计算机显卡（如游戏显卡）中。如果没有这种方法，计算机就会因内存不足而崩溃。
速度快：通过跳过不相关的部分，模型回答问题要快得多。本文显示，在处理非常长的文本时，其速度比旧方法快10 倍。
准确：与那些如果答案位置不对就会遗漏答案的“僵硬网格”方法不同，这位聪明的侦探几乎能像通读全书一样找到“大海捞针”。在测试中，它的准确率显著高于其他“跳过”方法。

总结

本文提出了一种方法，使人工智能模型无需超级计算机即可处理海量文本。DHSA不像旧方法那样盲目地阅读所有内容，或使用僵化的、一刀切的跳过规则，而是像一位聪明的编辑。它首先识别文本中自然的“段落”，然后快速扫描“目录”以找到最相关的部分，最后仅深入挖掘那些特定部分。

这使得标准计算机能够阅读和理解像小说或法律合同一样长的文档，且速度快，不会耗尽内存。

技术摘要：面向内存受限大语言模型推理的动态分层稀疏注意力

问题陈述

自注意力机制的二次计算复杂度（ $O(L^2)$ ）为大语言模型（LLM）扩展至长上下文带来了巨大的瓶颈，尤其是在严格的硬件内存预算下。尽管 prior 研究已证实 LLM 中的注意力矩阵本质上具有稀疏性，但现有解决方案面临显著局限：

静态稀疏方法：如 Longformer 和 BigBird 等方法依赖于固定模式（例如滑动窗口、膨胀注意力），无法适应注意力分布中随输入变化的差异，往往导致在多样化任务上的性能下降。
动态启发式方法：最近的动态方法（例如 MInference、H2O）利用预定义模板或启发式规则（例如“垂直 - 斜线”模式）来加速推理。然而，这些方法缺乏捕捉高度依赖输入的注意力稀疏性的灵活性，限制了其在不同任务和输入上的通用性。

核心挑战在于开发一种输入自适应、计算高效且能够在无需重新训练 LLM 主干的情况下保留因果依赖的稀疏注意力机制。

方法论：动态分层稀疏注意力（DHSA）

作者提出了动态分层稀疏注意力（DHSA），这是一个专为标准仅解码器 Transformer 设计的即插即用模块。DHSA 通过在保持 LLM 主干冻结的同时在线预测注意力稀疏性来运行。该框架采用两步分层路由过程，以降低识别重要 token 交互的复杂度。

1. 分层路由公式

DHSA 将问题从对所有 $L \times L$ 个 token 对进行评分，简化为两步过程：

步骤 1：块级预测：将输入 token 序列划分为可变长度的块。一个轻量级的边界预测器基于局部键向量窗口动态确定块边界。随后构建块级相似度矩阵，以估计块之间交互的重要性。
步骤 2：Token 级选择：基于块级相似度分数，系统将每个查询块路由到一组紧凑的相关键 token 索引。该选择将高分块扩展为其组成 token，同时强制执行因果约束，从而生成一个保留最具影响力的查询 - 键交互的稀疏注意力掩码。

2. 关键技术组件

动态边界检测：一个基于神经网络的预测器（包含编码器、特征融合和 MLP）分析局部键窗口以检测语义边界。它使用源自基础模型稠密注意力模式的自动软标签进行训练，避免了人工标注的需求。该预测器使用非极大值抑制（NMS）来最终确定边界，确保边界分离良好且对噪声具有鲁棒性。
鲁棒的块表示：为了有效处理可变长度的块，DHSA 采用长度归一化的池化方案。它计算嵌入的前缀和，除以实际块长度，随后进行长度归一化。这防止了零填充稀释平均表示，并确保在不同大小的块之间具有稳定的相似度估计。
硬件无关的后端：DHSA 实现了两个互补的后端以支持多样化的硬件：
- 一个PyTorch SDPA 后端，用于在模型系列和平台（包括 CPU）之间实现广泛的兼容性。
- 一个分块在线 softmax 后端，用于在 GPU 上实现高效率，利用流式 softmax 避免实例化完整的注意力矩阵。

主要贡献

分层路由公式：本文将稀疏化重构为块级路由问题。通过预测块 - 块相似度矩阵并将其用于将注意力路由到一组紧凑的 token，DHSA 在显著降低计算开销的同时保留了因果语义。
内容感知分割：作者引入了一种动态边界预测器，根据内容变化而非固定长度对序列进行分割。这结合了一种长度鲁棒的表示方案，确保无论块大小如何，相似度估计都能保持稳定。
实用的即插即用模块：DHSA 被设计为一个插件模块，插入在查询/键投影与注意力后端之间。它无需对 LLM 主干进行微调，支持多种开源权重模型系列，并可在 GPU 和 CPU 环境中运行。

实验结果

使用 LLaMA-3.1-8B（4-bit）、Qwen2.5-3B 和 Gemma-2-2B 等模型，在**“大海捞针”（Needle-in-a-Haystack）、LongBench和RULER**基准上进行了广泛评估。

准确性：DHSA 在高度稀疏的设定下保持了接近稠密的准确性。在 LongBench 上，在可比的预填充成本下，其相对于块稀疏注意力实现了12–20% 的相对准确性提升。在“大海捞针”测试中，DHSA 在固定密度为 6.25% 的情况下，在100K 上下文长度下仍保持高检索准确性，而静态方法在关键信息落在其受限范围之外时会出现性能急剧下降。
效率：
- 加速比：借助内存高效的分块后端，与稠密注意力相比，DHSA 在 128K 上下文长度下实现了高达10 倍的预填充加速。
- 内存约束：在单张 24GB GPU 上，DHSA 使 LLaMA-3.1-8B（4-bit）能够扩展至100K 上下文长度，而在该设定下稠密注意力会因内存不足（OOM）错误而失败。
- 批处理：在内存受限的设置中，DHSA 通过轻量级 for 循环对批次进行顺序处理，避免了批处理 FlashAttention-2 中常见的 OOM 故障，并提供更快的推理速度。

意义与主张

本文将 DHSA 定位为面向内存受限的长上下文 LLM 推理的高效且适应性强的解决方案。其主要意义在于弥合了注意力理论稀疏性与实际、输入自适应实现之间的差距。

适应性：与基于模板的方法不同，DHSA 从 token 嵌入中在线学习内容自适应的稀疏性，无需手动调整即可在多样化任务上提供加速。
可扩展性：通过将每层注意力成本从 $O(L^2)$ 降低至约 $O(L \cdot N_b)$ （其中 $N_b$ 为 token 预算），DHSA 使得在消费级硬件上处理极长上下文成为可能。
通用性：该框架与模型无关且兼容硬件，支持 GPU 和 CPU 后端，使其成为在多样化开源权重模型系列中部署长上下文能力的实用工具。

作者得出结论，DHSA 有效地解决了计算效率与检索准确性之间的权衡，为长上下文建模提供了一种比静态和启发式稀疏注意力方法更稳健的替代方案。