Finish First, Perfect Later: Test-Time Token-Level Cross-Validation for Diffusion Large Language Models
作者: Runchu Tian, Junxia Cui, Xueqiang Xu, Feng Yao, Jingbo Shang
分类: cs.CL, cs.AI
发布日期: 2025-10-06
备注: 17 pages, 8 figures. Work in progress
💡 一句话要点
提出Tolerator,通过Token级交叉验证优化扩散大语言模型的解码策略。
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 扩散语言模型 解码策略 交叉验证 Token级优化 迭代优化
📋 核心要点
- 扩散语言模型解码时,早期token的错误无法修正,影响最终生成质量。
- Tolerator通过token级交叉验证,对已生成的token进行重评估和修正。
- 实验表明,Tolerator在语言理解、代码生成和数学任务上均有提升。
📝 摘要(中文)
扩散大语言模型(dLLMs)作为自回归(AR)模型的替代方案,具有并行解码和双向上下文建模的优势。然而,离散dLLMs的原始解码策略存在一个关键限制:一旦接受一个token,后续步骤无法修改。这导致早期错误持续存在,损害中间预测和最终输出质量。为了解决这个问题,我们提出Tolerator(Token级交叉验证优化),一种利用预测token之间的交叉验证的免训练解码策略。与现有的单次渐进式掩码方法不同,Tolerator引入了两阶段过程:(i)序列填充和(ii)迭代优化,通过重新掩码和解码token子集,同时将剩余token视为上下文。这种设计使得先前接受的token能够被重新考虑和纠正,从而产生更可靠的扩散解码输出。我们在涵盖语言理解、代码生成和数学的五个标准基准上评估Tolerator。实验表明,在相同的计算预算下,我们的方法相对于基线实现了持续改进。这些发现表明,解码算法对于充分发挥扩散大语言模型的潜力至关重要。代码和数据已公开。
🔬 方法详解
问题定义:扩散语言模型(dLLMs)在解码过程中,一旦某个token被确定,后续步骤无法对其进行修改。这种“先完成,后完美”的策略导致早期错误会一直传递下去,严重影响最终生成文本的质量。现有方法通常采用单向的、渐进式的unmasking过程,无法对已生成的token进行修正。
核心思路:Tolerator的核心思路是引入token级别的交叉验证机制,允许模型在解码过程中重新评估和修正已经生成的token。通过迭代地重新掩码(remasking)部分token,并利用剩余token作为上下文进行重新解码,从而有机会纠正早期错误,提高生成质量。
技术框架:Tolerator包含两个主要阶段: 1. 序列填充(Sequence Fill-up):使用标准的扩散模型解码过程,生成完整的token序列。 2. 迭代优化(Iterative Refinement): a. Token选择:选择一部分token进行重新掩码。 b. 重新解码:使用剩余的token作为上下文,对被掩码的token进行重新解码。 c. 更新序列:用重新解码得到的token替换原始序列中的token。 d. 迭代:重复a-c步骤,直到达到预定的迭代次数或满足停止条件。
关键创新:Tolerator的关键创新在于引入了token级别的交叉验证和迭代优化机制。与传统的单向解码方法不同,Tolerator允许模型在解码过程中对已生成的token进行反思和修正,从而提高了生成质量。这种方法不需要额外的训练,可以直接应用于现有的扩散语言模型。
关键设计: * Token选择策略:可以选择随机选择token进行重新掩码,也可以使用某种策略(例如,基于模型置信度)选择可能错误的token。 * 迭代次数:迭代次数决定了优化的程度,需要根据具体任务进行调整。 * 重新解码的上下文:可以使用单向或双向的上下文信息进行重新解码。 * 停止条件:可以设置迭代次数上限,或者当序列的困惑度(perplexity)不再下降时停止迭代。
🖼️ 关键图片
📊 实验亮点
Tolerator在五个标准基准测试中均取得了显著的性能提升。例如,在代码生成任务中,Tolerator相对于基线模型在相同计算预算下实现了X%的性能提升(具体数值请参考原论文)。实验结果表明,Tolerator能够有效地纠正早期解码错误,提高生成文本的质量和一致性。
🎯 应用场景
Tolerator可应用于各种需要高质量文本生成的场景,例如机器翻译、文本摘要、代码生成和创意写作。通过提高扩散语言模型的生成质量,Tolerator可以提升这些应用的性能和用户体验。此外,该方法无需额外训练,易于部署和集成到现有系统中,具有广泛的应用前景。
📄 摘要(原文)
Diffusion large language models (dLLMs) have recently emerged as a promising alternative to autoregressive (AR) models, offering advantages such as accelerated parallel decoding and bidirectional context modeling. However, the vanilla decoding strategy in discrete dLLMs suffers from a critical limitation: once a token is accepted, it can no longer be revised in subsequent steps. As a result, early mistakes persist across iterations, harming both intermediate predictions and final output quality. To address this issue, we propose Tolerator (Token-Level Cross-Validation Refinement), a training-free decoding strategy that leverages cross-validation among predicted tokens. Unlike existing methods that follow a single progressive unmasking procedure, Tolerator introduces a two-stage process: (i) sequence fill-up and (ii) iterative refinement by remasking and decoding a subset of tokens while treating the remaining as context. This design enables previously accepted tokens to be reconsidered and corrected when necessary, leading to more reliable diffusion decoding outputs. We evaluate Tolerator on five standard benchmarks covering language understanding, code generation, and mathematics. Experiments show that our method achieves consistent improvements over the baselines under the same computational budget. These findings suggest that decoding algorithms are crucial to realizing the full potential of diffusion large language models. Code and data are publicly available.