Blog Tags GitHub
Back to blog
sparse-attention efficient-attention pruning vlm

Long-Context Modeling with Dynamic Hierarchical Sparse Attention for On-Device LLMs

面向设备端LLM的动态分层稀疏注意力长上下文建模

Long-Context Modeling with Dynamic Hierarchical Sparse Attention for On-Device LLMs

一、论文概述

项目内容
标题Long-Context Modeling with Dynamic Hierarchical Sparse Attention for On-Device LLMs
作者(未在摘要中列出)
机构-
论文https://arxiv.org/abs/2510.24606
发布2025-10

二、核心思想

问题定义

注意力机制的二次计算成本阻碍了长上下文LLM的可扩展性，特别是在资源受限的环境中。现有静态稀疏方法（如滑动窗口或全局token）虽然利用了注意力的稀疏性来降低成本，但由于其静态性，难以适应注意力的内容依赖变化。

现有方法的局限：

  • 静态稀疏方法：滑动窗口、全局token等方法无法适应内容变化
  • 动态方法：依赖预定义模板或启发式机制，通用性受限
  • 精度损失：可能剪枝掉上下文重要的token

解决方案概述

本文提出动态分层稀疏注意力 (DHSA)，一个数据驱动的框架，在不重新训练的情况下动态预测注意力稀疏性：

  1. 分层稀疏预测：将序列分割为变长chunk，计算chunk表示，上采样到token级别
  2. 动态边界检测：基于输入序列自适应确定chunk边界
  3. 鲁棒chunk表示：长度归一化聚合，避免chunk长度变化引入的偏差

核心优势：

  • 在Gemma2上匹配密集注意力的精度
  • 减少20-60%的prefill延迟
  • 减少35%的峰值内存使用
  • 相比块稀疏注意力实现6-18%的相对精度提升

DHSA框架概览

三、技术架构

核心公式

1. 分层稀疏预测

Step 1: Chunk级预测

将token序列 T=[t0,t1,...,tL−1]\mathbf{T}=[\mathbf{t}_{0},\mathbf{t}_{1},...,\mathbf{t}_{L-1}] 分割为 NcN_c 个非重叠chunk {C0,C1,...,CNc−1}\{\mathbf{C}_{0},\mathbf{C}_{1},...,\mathbf{C}_{N_{c}-1}\}，由边界索引 B={b0,b1,...,bNc}\mathcal{B}=\{b_{0},b_{1},...,b_{N_{c}}\} 定义，其中 0=b0<b1<...<bNc=L0=b_{0}<b_{1}<...<b_{N_{c}}=L。

构建chunk级相似度矩阵 Sc∈RNc×Nc\mathbf{S}_{c}\in\mathbb{R}^{N_{c}\times N_{c}}，其中 (Sc)l,k(\mathbf{S}_{c})_{l,k} 表示chunk Cl\mathbf{C}_{l} 和 Ck\mathbf{C}_{k} 之间交互的预测重要性。

Step 2: Token级选择

从 Sc\mathbf{S}_{c} 上采样得到token级相似度矩阵 St∈RL×L\mathbf{S}_{t}\in\mathbb{R}^{L\times L}：

fupsample(Sc,B)≔Stf_{\text{upsample}}(\mathbf{S}_{c},\mathcal{B})\coloneqq\mathbf{S}_{t}

对于每个chunk对 (Cl,Ck)(\mathbf{C}_{l},\mathbf{C}_{k})，子矩阵 (St)[bl:bl+1],[bk:bk+1](\mathbf{S}_{t})_{[b_{l}:b_{l+1}],[b_{k}:b_{k+1}]} 被赋值为 (Sc)l,k(\mathbf{S}_{c})_{l,k}。

然后应用TopK选择生成token级稀疏掩码 M\mathbf{M}，每个query token的预算为 NbN_b。

2. 动态边界检测

将chunking公式化为边界检测问题，对于每个位置 i∈[0,L−1]i\in[0,L-1]，定义边界指示函数 δ(i)\delta(i)。

编码器：对于每个候选位置 ii，提取两个局部窗口：

kleft=fMHA([ki−w+1,⋯ ,ki]),kright=fMHA([ki+1,⋯ ,ki+w])\mathbf{k}_{\text{left}}=f_{\text{MHA}}([\mathbf{k}_{i-w+1},\cdots,\mathbf{k}_{i}]), \quad \mathbf{k}_{\text{right}}=f_{\text{MHA}}([\mathbf{k}_{i+1},\cdots,\mathbf{k}_{i+w}])

其中 ww 是窗口大小，kj\mathbf{k}_{j} 是token jj 的key向量，fMHAf_{\text{MHA}} 是带池化的多头注意力模块。

边界预测器架构：

  • 共享编码器：处理左右窗口
  • 特征融合模块：融合左右特征
  • MLP分类器：预测边界概率

3. 鲁棒Chunk表示

为了避免chunk长度变化引入的偏差，应用长度归一化聚合：

  • 将chunk内的token嵌入聚合为chunk表示
  • 按chunk大小的平方根缩放平均嵌入
  • 确保不同长度chunk的表示具有可比性

实现细节

算法流程 (Algorithm 1: DHSA Prefill Stage)：

  1. 输入token序列和边界预测器
  2. 动态检测chunk边界
  3. 计算chunk级相似度
  4. 上采样到token级
  5. TopK选择生成稀疏掩码
  6. 应用稀疏注意力计算

四、核心创新

创新点说明理论/实验依据
分层稀疏预测chunk级预测→token级选择减少计算复杂度
动态边界检测基于内容自适应确定chunk边界适应内容变化
长度归一化聚合按chunk大小平方根缩放避免长度偏差
数据驱动框架无需重新训练即可预测稀疏性即插即用
端到端优化边界预测器可端到端训练自动学习最优边界

五、实验结果

实验设置

模型：Gemma2-2b-it 基准：

  • Needle-in-a-Haystack Test
  • LongBench

基线方法：

  • Dense Attention (密集注意力)
  • Block Sparse Attention (块稀疏注意力)
  • Sliding Window Attention (滑动窗口注意力)

Needle-in-a-Haystack结果

Needle-in-a-Haystack对比 Needle-in-a-Haystack DHSA

Figure 2: Gemma2-2b-it上的Needle-in-a-haystack结果（最大上下文长度=8k）。

关键发现：

  • DHSA在budget=1k时匹配密集注意力的精度
  • 动态chunking优于静态chunking
  • 在不同上下文位置都能有效检索信息

LongBench性能

方法NrtvQAQasperMf-enHotpotQA2WikiMQAMusiqueGovReportQMSumMultiNewsTriviaQASAMSum
Dense22.3735.3237.3241.6332.0519.0527.0821.0825.4887.0041.26
Block sparse16.7426.1532.8335.7431.9314.4426.2019.5425.3086.1240.38
DHSA20.6930.2034.9838.7831.9615.9026.7520.7425.3887.0341.46
DHSA (+bs)19.6428.6731.2039.5032.9714.3926.7220.1225.5787.7440.68

关键结果：

  • DHSA在大多数任务上匹配或接近密集注意力的精度
  • 相比块稀疏注意力，DHSA实现6-18%的相对精度提升
  • DHSA (+bs) 表示跨层共享边界的变体

效率分析

延迟对比 内存对比

关键发现：

  • DHSA减少20-60%的prefill延迟
  • 减少35%的峰值内存使用
  • 在单个24GB GPU上支持更长上下文

边界检测可视化

静态vs动态chunking 动态chunking

Figure 5: 静态chunking与动态chunking的对比。

关键发现：

  • 动态chunking能够识别语义边界
  • 在不同层表现出不同的chunking模式
  • 局部注意力层和全局注意力层有不同的chunking特征

六、相关工作

稀疏注意力方法

方法类别代表方法特点局限
静态稀疏Sliding Window, Global Tokens固定模式无法适应内容变化
块稀疏Block Sparse Attention规则分块忽略语义边界
动态稀疏DHSA内容自适应需要训练边界预测器

长上下文优化

方法关键特性本文对比
FlashAttention计算优化正交方法，可组合
KV Cache压缩减少内存不同优化维度
线性注意力降低复杂度近似方法

七、总结

核心贡献

  1. DHSA框架：提出数据驱动的动态分层稀疏注意力框架
  2. 动态边界检测：基于内容自适应确定chunk边界
  3. 长度归一化聚合：避免chunk长度变化引入的偏差
  4. 显著效率提升：减少20-60%延迟和35%内存使用
  5. 精度保持：匹配密集注意力精度，优于块稀疏注意力

技术影响

  • 设备端部署：使长上下文LLM在资源受限设备上可行
  • 即插即用：无需重新训练即可集成到现有模型
  • 通用性：适用于不同模型架构和任务
  • 效率-精度权衡：提供灵活的预算控制

局限性

  1. 需要训练边界预测器
  2. 边界检测增加了一定开销
  3. 在极短序列上可能无法充分体现优势
  4. 依赖于chunk级相似度的质量

八、参考资源

九、图片索引

图片说明文件名
Figure 1DHSA框架概览framework-overview.jpg
Figure 2Needle-in-a-haystack结果对比needle-haystack-dense.jpg, needle-haystack-dhsa.jpg
Figure 3边界预测器架构boundary-predictor.jpg
Figure 4自动标注策略auto-labeling.jpg
Figure 5静态vs动态chunking对比static-chunking-local.jpg, dynamic-chunking-local.jpg
Figure 6不同方法的haystack对比haystack-comparison-dense.jpg, haystack-comparison-dhsa-1k.jpg
Figure 7延迟对比latency-comparison.jpg
Figure 8内存对比memory-comparison.jpg