扩散语言模型的注意力汇聚感知剪枝 Sink-Aware Pruning for Diffusion Language Models
针对扩散语言模型中注意力汇聚不稳定的特点，提出感知汇聚的剪枝方法
前置知识
扩散语言模型 (Diffusion Language Models, DLMs)
扩散语言模型是一类非自回归的文本生成模型，借鉴了连续域扩散模型的思想，将文本生成建模为迭代去噪过程。与自回归模型逐个生成 token 不同，DLM 在每个去噪时间步同时更新整个 token 序列或其潜在表示，直到序列收敛。代表模型包括 LLaDA、Dream、MMaDA 等。这种双向迭代去噪机制使得模型能够并行生成文本，但也因此需要多次前向传播，推理计算成本显著高于自回归模型。
本文的核心研究对象就是 DLM，理解其迭代去噪生成机制是理解全文的前提，特别是去噪时间步 (denoising timestep) 的概念贯穿方法和实验。
注意力汇聚 (Attention Sinks)
注意力汇聚是指一小部分 token 位置（通常是序列开头的 BOS 或系统提示词等早期 token）在多个注意力层和头中持续吸引不成比例的大注意力权重，尽管这些 token 本身的语义价值有限。Xiao et al. (2023) 首次在大语言模型中系统描述了这一现象，将其归因于 Softmax 归一化机制：当某个 query 没有强匹配时，注意力质量必须被分配到某处，全局可见的早期 token 自然成为冗余注意力的倾倒场。在自回归模型中，这些汇聚位置表现为稳定的全局锚点。
本文的核心创新点就是发现 DLM 中的注意力汇聚行为与自回归模型截然不同——汇聚位置是不稳定的、短暂的，这一发现直接支撑了论文的剪枝策略。
后训练剪枝 (Post-Training Pruning)
后训练剪枝是一种模型压缩技术，直接从预训练模型中识别并移除不重要参数，无需完整重训练。典型方法包括 Wanda（基于权重幅值与输入激活范数的乘积评估重要性）和 SparseGPT（将剪枝建模为逐层重建问题，使用二阶信息迭代优化）。相比结构化剪枝移除整个神经元或注意力头，非结构化剪枝可精确到单个权重，灵活性更高但需要硬件支持。
本文提出的 Sink-Aware Pruning 是在 Wanda 和 SparseGPT 基础上叠加的改进，理解这两个基线方法的原理才能看懂本文的技术贡献。
汇聚方差 (Sink Variance)
本文引入的度量指标，包含两个维度：空间方差衡量注意力在整个生成轨迹上对各位置分配的不均匀程度；时间方差衡量汇聚位置在不同时间步间的漂移程度。低时间方差表示汇聚位置锁定在同一区域，高时间方差表示汇聚位置随时间步大幅迁移。在论文的数学表达中，空间方差定义为各位置平均注意力质量的方差，时间方差定义为注意力加权质心随时间步的变化方差。
这是本文方法论的核心——通过量化汇聚位置的稳定性来区分稳定锚点和瞬态汇聚，直接决定哪些位置应被剪枝。
研究动机
现有扩散语言模型的剪枝方法直接继承了自回归大语言模型的启发式规则，其中最典型的就是始终保留注意力汇聚 token。在自回归模型中，汇聚位置通常固定在序列开头（如 BOS token 或系统提示词），这些位置在所有生成步骤中持续充当稳定的全局锚点，帮助传播条件信息并稳定残差流动力学。因此 AR 剪枝方法（如 Wanda、SparseGPT）通常在重要性评估中保护这些位置，避免灾难性质量下降。然而，这种做法隐含地假设了自回归模型中的注意力行为及其重要 token 能够原封不动地迁移到扩散式生成中，但这个假设从未被验证过。由于扩散语言模型通过迭代去噪生成文本，每个时间步都对全序列进行注意力计算，其注意力组织随去噪过程不断演化——早期步骤在高噪声下需要解析全局结构，后期步骤在低噪声下需要精炼局部句法和语义——这与自回归模型的单向前向传播有本质区别。
本文的目标是本文的具体目标有两个层面：在分析层面，系统地量化和比较扩散语言模型与自回归大语言模型中注意力汇聚行为的差异，特别是汇聚位置随生成步骤变化的稳定性；在应用层面，基于分析发现提出一种针对扩散语言模型特性的剪枝策略，在不重训练的前提下改善模型压缩的质量-效率权衡，使剪枝后的 DLM 在匹配计算预算下优于现有基线方法。
与已有工作不同的是，本文的独特切入角度在于挑战了模型压缩领域的一个隐含假设——注意力汇聚是跨范式不变的重要特征。通过引入汇聚位置方差这一量化指标，论文揭示了自回归模型和扩散模型之间一个根本性差异：AR 模型的汇聚是空间集中且时间稳定的（高空间方差、近零时间方差），而 DLM 的汇聚是空间分散且时间短暂的（低空间方差、时间方差高出数个数量级）。这种差异意味着始终保留汇聚这一 AR 领域的最优策略，在扩散模型中反而会保留大量不稳定的、功能上非关键的参数，浪费有限的剪枝预算。因此，本文提出的是范式感知的剪枝——对 AR 模型保留汇聚，对 DLM 允许并鼓励剪除不稳定汇聚。
核心方法
Sink-Aware Pruning 的整体思路可以用一个直觉来理解：在自回归模型中，注意力汇聚就像固定的灯塔，始终引导着信息流动，因此不能拆除；但在扩散语言模型中，汇聚更像是闪烁的信号灯，位置不断变化，某些时间步出现的汇聚到下一时间步可能就消失了。既然这些闪烁的信号灯不提供稳定的结构支撑，那在压缩模型时就可以安全地剪除它们，把宝贵的参数预算留给真正承载语义信息的权重。技术路线分为三步：首先，在校准数据上计算每个 token 位置在所有去噪步骤中接收的注意力质量，量化汇聚位置的方差；然后，通过阈值检测和 sigmoid 松弛为每个位置生成软汇聚分数，分数高的位置被判定为不稳定汇聚；最后，用汇聚分数对原始激活进行加权抑制，将抑制后的激活代入现有的剪枝重要性指标（Wanda 或 SparseGPT），从而在剪枝决策中自动降低不稳定汇聚位置的权重。
本文的核心创新点在于发现了自回归模型中始终保留注意力汇聚这一启发式规则并不适用于扩散语言模型，并基于此提出了一种生成范式感知的剪枝策略。与已有方法的本质区别在于：Wanda 和 SparseGPT 等基线方法完全不考虑汇聚状态，仅基于权重幅值或二阶重建误差评估重要性；而此前专门为 LLM 设计的汇聚感知方法（如 StreamingLLM）则选择始终保留汇聚位置。本文首次提出，汇聚的可剪枝性取决于其在生成过程中的稳定性——通过量化汇聚位置的方差来自动判断，对 AR 模型低方差汇聚保留，对 DLM 高方差汇聚剪除。这种自适应机制使得同一个框架可以同时适用于两种生成范式。
方法步骤详情
Sink-Aware Pruning 包含四个连续步骤。第一步是注意力质量聚合：对每个去噪时间步 t，计算每个 token 位置 j 在所有层和头上接收到的累积注意力分数 m_t(j)，即对所有注意力头和层的注意力权重沿查询维度求和。第二步是汇聚位置检测与软分数计算：通过阈值准则判断每个位置在当前时间步是否为汇聚位置，并用 sigmoid 松弛得到连续的汇聚分数，再在校准时间步集合上平均得到步不变的汇聚估计。第三步是激活抑制：根据汇聚分数定义每个位置的权重因子 omega_j = 1 - phi_bar(j)，构建抑制激活，使不稳定汇聚位置的激活接近零。第四步是剪枝决策：将抑制激活代入 Wanda 重要性公式（权重幅值乘以输入激活范数）或 SparseGPT 的 Hessian 计算（使用抑制激活构建经验 Hessian），然后按原方法执行剪枝和权重重建。通过在 Hessian 中抑制汇聚 token，重建过程将容量分配给忠实复现语义有意义位置的输出，而不是保留那些过量激活会主导二阶统计量的 token 的行为。
技术新颖性
Sink-Aware Pruning 的技术新颖性体现在三个层面。首先，在分析方法上，论文首次系统引入了汇聚位置方差的量化框架，包含空间方差和时间方差两个互补维度，为理解不同生成范式下注意力汇聚的本质差异提供了严格的数学工具。其次，在方法设计上，通过软 sigmoid 松弛将离散的汇聚检测转化为可微的连续权重因子，使得汇聚感知信号可以无缝嵌入现有的剪枝重要性指标（Wanda 和 SparseGPT），而不需要重新设计整个剪枝流程——这是一种即插即用的改进思路。最后，在范式感知性上，该方法对 AR 模型（低汇聚方差）和 DLM（高汇聚方差）自动产生不同行为：对前者几乎不改变基线剪枝决策，对后者则主动降低不稳定汇聚位置的重要性，实现了跨生成范式的自适应。
实验结果
论文在多个扩散语言模型家族（LLaDA-8B、Dream-7B、LLaDA-1.5-8B、MMaDA-8B）和 8 个标准基准上进行了全面评估。核心发现是：Sink-Aware Pruning 在中等至高稀疏度（50%-75%）下一致地匹配或超越对应的基线方法，改善最为显著。具体而言，在 LLaDA-8B 上 50% 稀疏度下，Sink-Aware Wanda 的平均准确率为 53.18%（vs. 基线 Wanda 52.70%），Sink-Aware SparseGPT 为 52.36%（vs. 基线 52.34%）；在 75% 稀疏度下差距更大，Sink-Aware Wanda 达到 30.94%（vs. 基线 29.99%），提升了近 1 个百分点。在 Dream-7B 上，75% 稀疏度下 Sink-Aware Wanda 达到 30.52%（vs. 30.38%），Sink-Aware SparseGPT 达到 31.83%（vs. 31.72%）。在 LLaDA-1.5 上，75% 稀疏度下 Sink-Aware Wanda 的提升最为显著，达到 32.89%（vs. 31.03%），提升了 1.86 个百分点。结构化剪枝实验同样支持这一趋势：在 LLaDA 上 30% 结构化剪枝下，Sink-Aware 方法在 PIQA 上达到 0.6955（vs. 基线 0.6834），WinoGrande 上达到 0.6740（vs. 0.6630），ARC-E 上达到 0.7175（vs. 0.6907）；50% 结构化剪枝下差距更大，ARC-E 上达到 0.5279（vs. 0.4853）。这些结果支持了论文的核心论点：注意力汇聚并非普遍必须保留的 token，其效用取决于生成动态。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| MMLU (5-shot) | 准确率 (%) | LLaDA: 62.16 (Wanda), 60.79 (SparseGPT) at 50% sparsity | Wanda 61.43, SparseGPT 60.97 at 50% sparsity | Sink-Aware Wanda +0.73pp, Sink-Aware SparseGPT -0.18pp |
| PIQA (0-shot) | 准确率 (%) | LLaDA: 73.18 (Wanda), 72.95 (SparseGPT) at 50% sparsity | Wanda 72.63, SparseGPT 72.20 at 50% sparsity | Sink-Aware Wanda +0.55pp, Sink-Aware SparseGPT +0.75pp |
| GSM8K (5-shot) | 准确率 (%) | LLaDA: 55.88 (Wanda), 52.11 (SparseGPT) at 50% sparsity | Wanda 57.01, SparseGPT 53.53 at 50% sparsity | Sink-Aware Wanda -1.13pp, Sink-Aware SparseGPT -1.42pp |
| WinoGrande (0-shot) | 准确率 (%) | LLaDA: 65.27 (Wanda), 65.82 (SparseGPT) at 50% sparsity | Wanda 64.56, SparseGPT 64.64 at 50% sparsity | Sink-Aware Wanda +0.71pp, Sink-Aware SparseGPT +1.18pp |
| Average across 8 benchmarks (LLaDA) | 平均准确率 (%) | 75% sparsity: Wanda 30.94, SparseGPT 32.57 | 75% sparsity: Wanda 29.99, SparseGPT 32.57 | Sink-Aware Wanda +0.95pp at 75% sparsity |
局限与改进
论文坦诚地承认了若干局限性。首先，汇聚统计量是从固定的校准分布（WikiText-2 的 128 条序列，长度 2048 tokens）上估计的，分布偏移可能会降低汇聚检测的可靠性——如果实际推理数据的分布与校准集差异很大，汇聚位置的判定可能不准确。其次，论文主要评估的是后训练剪枝而没有进行恢复微调 (recovery finetuning)，结合轻量级后剪枝自适应可能进一步提升鲁棒性。第三，虽然包含了一个多模态 DLM（MMaDA），但在更大规模的多模态和长上下文设置上的验证仍然不足。此外，从独立观察来看，Sink-Aware Pruning 的改进在低稀疏度（25%）下有时是边际性的甚至是负向的（如 GSM8K 上 Sink-Aware Wanda 在 LLaDA 50% 稀疏度下反而下降 1.13pp），说明汇聚不稳定性并非在所有层、任务和模型家族中同等显著，这限制了该方法的普适增益。
独立分析的弱点
尽管 Sink-Aware Pruning 展示了一致的改进趋势，但从独立分析来看存在几个值得改进的弱点。第一，校准数据的固定性：当前方法使用 WikiText-2 的固定 128 条样本作为校准集，但在实际部署中，下游任务可能涉及代码生成、数学推理或对话等不同分布，汇聚模式可能因任务而异。改进方向是设计任务自适应的校准策略，或在运行时动态更新汇聚统计量。第二，时间方差阈值的敏感性：汇聚检测依赖于阈值和温度这两个超参数，但论文未提供其敏感性分析。不同模型和稀疏度可能需要不同的阈值设置，增加了调参负担。改进方向是引入自适应阈值机制。第三，在部分任务上（如 GSM8K 数学推理），Sink-Aware 变体偶尔会出现性能下降，说明在需要精确推理的任务中，某些瞬态汇聚可能仍承载关键信息。改进方向是分层设计汇聚策略——对不同层使用不同的汇聚保留率。第四，当前方法仅在非结构化和粗粒度结构化剪枝中评估，尚未探索与量化等其他压缩技术的联合优化。
未来方向
作者在论文中提出了几个未来研究方向。首先是逐层时间步自适应汇聚策略：当前方法对所有层使用相同的汇聚分数，但不同层的汇聚稳定性可能存在显著差异——早期层可能更需要保留汇聚以维持全局信息流，而后期层的汇聚可能更加瞬态。其次是与量化的联合优化：将 Sink-Aware 剪枝与低比特量化结合，进一步拓展质量-效率前沿。第三是更广泛的验证：在更大规模的多模态 DLM 和长上下文设置中验证方法的有效性。基于已有成果还可以延伸到以下方向：将汇聚感知思想应用于 DLM 的训练阶段（如在训练中鼓励或抑制特定的汇聚模式），探索在混合自回归-扩散模型（如 SDAR、Block Diffusion）中的汇聚行为，以及设计硬件友好的汇聚感知结构化剪枝算法以避免非结构化剪枝的硬件开销。
复现评估
论文的复现条件相对友好。代码已开源于 GitHub（https://github.com/VILA-Lab/Sink-Aware-Pruning），并在论文中明确说明了基线方法的实现基于 Wanda 和 SparseGPT 的开源代码库。校准数据使用公开的 WikiText-2 数据集，128 条随机采样序列截断至 2048 tokens。评估使用 lm-evaluation-harness 框架，保证了评估协议的一致性。模型方面使用了公开可用的预训练权重：LLaDA-8B、Dream-7B、LLaDA-1.5-8B 和 MMaDA-8B。不过，实验涉及多个 7-8B 规模模型在 8 个基准上的评估，需要相当的 GPU 算力；汇聚方差的统计分析还需要在多个去噪时间步上计算注意力矩阵，内存开销较大。总体而言，对于拥有充足算力的研究团队，复现难度中等。
论文图表
该图通过左右两组示意对比了标准注意力汇聚和本文 Sink-Aware Pruning 的行为。左侧展示了 DLM 中注意力汇聚位置随去噪步骤不断变化（用不同颜色的箭头和圆圈表示），右侧展示了 Sink-Aware Pruning 如何通过降低不稳定汇聚的权重来改善注意力分配。
这是论文的核心概念图，直观地传达了全文的关键发现——DLM 中的汇聚是不稳定的——以及方法的基本思想。
该图展示了 LLaMA-3-8B（AR 模型）和 LLaDA（DLM）在不同生成阶段（25%、50%、75%）的注意力热力图。每张子图中 y 轴为 token 位置，x 轴为注意力头/层的聚合注意力，颜色深浅表示接收的注意力质量。在 LLaMA 中，汇聚位置（深色纵向条带）在各步骤间保持稳定；而在 LLaDA 中，汇聚位置随扩散步骤显著偏移，表明更高的汇聚方差。
这是支撑论文核心论点的最关键实验证据，通过可视化直接对比了 AR 和 DLM 中汇聚行为的根本差异。
该图包含两个子图。(a) 展示了跨生成/去噪步骤的汇聚方差，LLaDA 和 Dream 的方差约为 0.0173 和 0.0149（乘以10的负2次方），而 Llama 3.1 和 Qwen 2.5 的方差仅为 3.4e-06 和 1.2e-05，DLM 比 AR 模型高出数个数量级。(b) 展示了每个 token 在所有步骤中接收总注意力质量的方差，AR 模型反而更高（112.50 vs 53.14），表明 AR 模型的汇聚空间上更集中但时间上更稳定。
这是论文定量分析的核心图表，用具体数字证实了 DLM 和 AR 模型汇聚方差的巨大差异，为方法提供了统计学支撑。