发现扩散语言模型的注意力焦点不稳,提出动态修剪策略提升推理效率。
Sink-Aware Pruning for Diffusion Language Models
- 识别扩散模型中不稳定的注意力焦点并自动修剪
- 无需重训练,在相同计算量下优于现有剪枝方法
- 适合追求高效推理的扩散模型应用开发者
扩散语言模型(DLMs)因迭代去噪带来高推理开销,亟需高效剪枝。现有剪枝启发式方法多沿用自回归语言模型(AR LLMs),通常保留注意力聚焦词(sink tokens),因其在AR模型中为稳定全局锚点。我们发现该假设在DLMs中不成立:注意力焦点位置在整个生成轨迹中显著波动(通过不同时间步主导焦点位置的转移程度衡量),表明焦点常为瞬时存在,结构重要性低于AR模型。基于此,我们提出【Sink-Aware Pruning】,可自动识别并修剪DLM中不稳定的焦点。无需重训练,该方法在相同计算开销下实现更优的质量-效率平衡,超越强基线剪枝方法。代码已开源:https://github.com/VILA-Lab/Sink-Aware-Pruning。
原文摘要 · Abstract (English)
Diffusion Language Models (DLMs) incur high inference cost due to iterative denoising, motivating efficient pruning. Existing pruning heuristics largely inherited from autoregressive (AR) LLMs, typically preserve attention sink tokens because AR sinks serve as stable global anchors. We show that this assumption does not hold for DLMs: the attention-sink position exhibits substantially higher variance over the full generation trajectory (measured by how the dominant sink locations shift across timesteps), indicating that sinks are often transient and less structurally essential than in AR models. Based on this observation, we propose ${\bf \texttt{Sink-Aware Pruning}}$, which automatically identifies and prunes unstable sinks in DLMs (prior studies usually keep sinks for AR LLMs). Without retraining, our method achieves a better quality-efficiency trade-off and outperforms strong prior pruning baselines under matched compute. Our code is available at https://github.com/VILA-Lab/Sink-Aware-Pruning.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。