扩散语言模型的感知下沉剪枝法

摘要

扩散语言模型因迭代去噪过程导致推理成本高昂，亟需高效剪枝方法。现有剪枝启发式方法主要沿袭自自回归大语言模型，通常保留注意力汇聚令牌，因为自回归模型中的汇聚令牌可作为稳定的全局锚点。本文发现该假设不适用于扩散语言模型：在完整生成轨迹中，注意力汇聚位置的方差显著更高（通过主导汇聚位置在时间步间的偏移程度衡量），表明扩散模型的汇聚点常具有瞬时性，其结构重要性低于自回归模型。基于此发现，我们提出**汇聚感知剪枝法**，可自动识别并剪枝扩散模型中不稳定的汇聚点（先前研究通常为自回归大语言模型保留汇聚点）。无需重新训练，本方法在匹配计算量下实现了更优的质量-效率平衡，超越了现有强基准剪枝方法。代码已开源：https://github.com/VILA-Lab/Sink-Aware-Pruning。

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 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.

扩散语言模型的感知下沉剪枝法

Sink-Aware Pruning for Diffusion Language Models

摘要

Support