让已稳定词汇跳过计算，显著降低扩散语言模型的推理开销。
Stopping Computation for Converged Tokens in Masked Diffusion-LM Decoding
- 检测未掩码词位置的后验分布是否稳定，稳定则锁定该位置。
- 在LLaDA-8B上减少30%~50%算法浮点运算量，生成质量不变。
- 适合追求高效推理的扩散语言模型应用开发者。
掩码扩散语言模型通过迭代采样逐步解掩码生成序列，但每一步仍对所有位置重新计算注意力与前馈层，即使多数已解掩码的词基本固定，造成大量计算浪费。本文提出SureLock：当某未掩码位置的后验分布跨步骤趋于稳定（满足锁定条件）时，将其锁定——此后跳过其查询投影与前馈子层计算，仅缓存其注意力键值，供其他位置继续访问。该方法将每步主要计算成本从$O(N^2d)$降至$O(MNd)$，其中$N$为序列长度，$M$为未锁定位置数，$d$为模型维度。实际中$M$随迭代递减，带来显著节省。在LLaDA-8B上，相比无锁定的同采样器，SureLock降低30%~50%算法浮点运算量，且生成质量保持相当。我们还提供理论分析，证明仅监控锁定步的局部KL即可控制最终词概率的偏差。项目页面见https://daioba.github.io/surelock。
原文摘要 · Abstract (English)
Masked Diffusion Language Models generate sequences via iterative sampling that progressively unmasks tokens. However, they still recompute the attention and feed-forward blocks for every token position at every step -- even when many unmasked tokens are essentially fixed, resulting in substantial waste in compute. We propose SureLock: when the posterior at an unmasked position has stabilized across steps (our sure condition), we lock that position -- thereafter skipping its query projection and feed-forward sublayers -- while caching its attention keys and values so other positions can continue to attend to it. This reduces the dominant per-iteration computational cost from $O(N^2d)$ to $O(MNd)$ where $N$ is the sequence length, $M$ is the number of unlocked token positions, and $d$ is the model dimension. In practice, $M$ decreases as the iteration progresses, yielding substantial savings. On LLaDA-8B, SureLock reduces algorithmic FLOPs by 30--50% relative to the same sampler without locking, while maintaining comparable generation quality. We also provide a theoretical analysis to justify the design rationale of SureLock: monitoring only the local KL at the lock step suffices to bound the deviation in final token probabilities. Our project page is available at https://daioba.github.io/surelock .
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。