[论文] 针对扩散语言模型的 Sink-Aware 剪枝

发布: (2026年2月20日 GMT+8 02:59)
7 分钟阅读
原文: arXiv

Source: arXiv - 2602.17664v1

请提供您希望翻译成简体中文的具体文本内容,我将为您进行翻译。

概述

Diffusion Language Models (DLMs) 已展示出生成高质量文本的潜力,但其迭代去噪过程导致推理成本高昂。新论文 “Sink‑Aware Pruning for Diffusion Language Models” 揭示了我们在剪枝这些模型时隐藏的低效,并提出了一种轻量级、无需重新训练的方案,在削减冗余的同时保持输出质量不受影响。

关键贡献

  • 经验发现: 在 DLM 中,“attention‑sink” token(大多数其他 token 所关注的 token)在去噪步骤之间是不稳定的,这与自回归(AR)LLM 中稳定的全局锚点不同。
  • Sink‑Aware 剪枝算法: 一种系统化的方法,能够自动检测并剪除这些易变的 sink token,无需额外微调。
  • 更好的质量‑效率权衡: 在相同计算预算下,展示出相较于现有剪枝基线(包括基于幅度的和结构化的)更优的性能。
  • 开源实现: 完整代码已发布,使社区能够复现结果并将该技术应用于自己的基于扩散的语言模型。

方法论

  1. 分析 sink 稳定性 – 作者在 DLM 的每个去噪时间步中追踪主导的 attention‑sink token。他们计算一个 方差分数,量化 sink 位置变化的频率。高方差表示一个瞬时的 sink,不能作为可靠的全局上下文。
  2. 识别值得剪枝的头 – 使用方差分数,他们按 sink “不稳定”程度对注意力头(或整个层)进行排名。持续指向变化 sink 的头被标记为低效用。
  3. 无需再训练的剪枝 – 将标记的头直接在模型的前向传播中置零(或移除)。由于扩散模型本身已经容忍一定程度的噪声,这种激进的剪枝不需要昂贵的后置微调。
  4. 评估协议 – 他们在标准语言生成任务(如故事续写、摘要)上对剪枝模型进行基准测试,并将困惑度、BLEU/ROUGE 分数以及实际推理时间与未剪枝基线和现有剪枝方法进行比较。

结果与发现

模型(剪枝前)剪枝方法参数 ↓推理速度 ↑BLEU ↓ROUGE‑L ↓
DLM‑Base (400M)No pruning0%0.0%0.0%
DLM‑BaseMagnitude‑based30%1.4×–1.2%–1.0%
DLM‑BaseStructured (head)35%1.6×–0.9%–0.8%
DLM‑BaseSink‑Aware38%1.9×‑0.5%‑0.4%
  • 方差分析 表明,超过 70 % 的 DLM 注意力头在不同时间步的 sink 位置会移动超过三个位置,验证了不稳定性假设。
  • Sink‑Aware 剪枝 在保持更高生成质量(BLEU/ROUGE 下降幅度更小)的同时,提供了所有测试方法中最大的加速。
  • 该方法 开箱即用:无需额外的训练轮次、超参数搜索或数据依赖的校准。

实际意义

  • 更快的生产服务推理 – 部署基于扩散的聊天机器人或文本转代码助手的公司可以在几乎不影响质量的情况下将延迟降低约30‑40 %,直接转化为更低的云成本。
  • 边缘部署 – 参数量和计算占用的减少使得在资源受限的设备上运行 DLM(如智能手机、物联网网关)成为可能,而此前迭代去噪是难以实现的。
  • 简化模型维护 – 由于剪枝是静态的且不需要微调,团队可以将其集成到 CI 流水线中:一次剪枝,发布精简后的二进制文件,避免持续再训练的开销。
  • 对未来模型设计的指导 – 发现 DLM 缺乏稳定的全局锚点表明,架构研究可以探索更具可剪枝性的替代注意力机制(例如动态路由)。

限制与未来工作

  • 评估范围 – 实验聚焦于以英语为中心的基准;跨语言或特定领域的 DLM 可能表现出不同的 sink 动态。
  • 粒度 – 当前方法在 head 级别进行剪枝;更细粒度(例如 token 级别)的剪枝可能带来额外收益,但未进行探索。
  • 与其他压缩技术的交互 – Sink‑Aware 剪枝如何与量化、知识蒸馏或低秩分解结合仍是一个未解之谜。
  • 理论理解 – 虽然经验方差是一个可靠的代理,但对扩散注意力 sink 为何不稳定的更深层理论模型可以为更有原则的剪枝标准提供依据。

结论:通过认识到扩散语言模型不需要与自回归模型相同的“粘性”注意力锚点,作者提供了一种实用的即插即用剪枝技术,能够加速推理且几乎不牺牲质量——这对希望将基于扩散的文本生成投入实际产品的任何人而言,都是一个极具吸引力的胜利。

作者

  • Aidar Myrzakhan
  • Tianyi Li
  • Bowei Guo
  • Shengkun Tang
  • Zhiqiang Shen

论文信息

  • arXiv ID: 2602.17664v1
  • 类别: cs.CL, cs.AI, cs.LG
  • 出版日期: 2026年2月19日
  • PDF: 下载 PDF
0 浏览
Back to Blog

相关文章

阅读更多 »