LLaDA2.1：通过Token编辑加速文本扩散模型推理 LLaDA2.1: Speeding Up Text Diffusion via Token Editing

LLaDA2.1的方法包含训练和推理两个完整流程。训练阶段分三步进行：（1）继续预训练（CPT）：使用混合M2T和T2T的双重流目标，在文档级别的数据上训练。M2T流让模型学习在[MASK]位置预测正确token（草稿能力），T2T流让模型从随机噪声扰动中恢复原始token（编辑能力）。同时采用多轮前向（MTF）数据增强技术，通过在不同时间步暴露模型于更多编辑场景来增强编辑能力。（2）监督微调（SFT）：同样使用混合M2T和T2T目标，但在QA任务的变体上进行，进一步强化模型的指令遵循能力。（3）强化学习（RL）对齐：采用EBPO框架，利用ELBO作为精确序列对数似然的代理，通过裁剪的代理目标 $J_{EBPO}(\theta)$ 优化策略。具体地，对离散化时间步 $\{t_n\}_{n=1}^N$ 和权重 $\{w_n\}$，构造复合输入 $z_n = y_{t_n} \oplus y_0$，在单次前向传播中并行计算所有块条件概率。推理阶段采用阈值解码与显式编辑机制结合的算法：在每个时间步 $t$，根据公式 $\Gamma_t = \{i | x_i^t = [MASK] \text{ and } p_\theta(v_i^t|x^t) > \tau_{mask}\}$ 和 $\Delta_t = \{i | x_i^t \neq v_i^t \text{ and } p_\theta(v_i^t|x^t) > \tau_{edit}\}$ 确定更新索引，然后在 $\Gamma_t \cup \Delta_t$ 的并集上应用更新。可选地启用多块编辑（MBE）机制，允许模型在新块完成后回溯修订之前已生成的块。

LLaDA2.1的技术新颖性体现在多个层面。首先，在解码算法层面，本文首次将token编辑（T2T）作为一等公民引入dLLM的解码过程，打破了传统吸收态框架的不可逆约束。这种'Draft-and-Edit'范式与先前的置信度重掩码（remasking）方法有本质区别：重掩码只是将低置信度token重新掩码等待后续步骤重新生成，而T2T编辑允许模型直接将一个token替换为另一个更合适的token，实现了真正的'修正'而非简单的'重新生成'。其次，在训练策略层面，混合M2T和T2T的双重流目标是全新的设计，使模型在单一参数空间内同时具备草稿和编辑两种能力，这与传统的纯M2T训练有本质区别。MTF数据增强技术进一步丰富了编辑场景的多样性。第三，在RL训练层面，EBPO框架通过ELBO代理和向量化似然估计（Vectorized Likelihood Estimation）解决了dLLM策略梯度估计的不可计算性问题，首次实现了dLLM上的大规模RL训练，支持前所未有的上下文长度和训练规模。第四，在系统层面，可配置的双阈值机制将速度-质量权衡转化为用户可控的连续体，这是一个优雅的工程创新。

尽管LLaDA2.1取得了显著进展，仍存在几个值得关注的弱点。第一，S Mode在某些基准上的质量下降较为明显，例如在BIG-Bench Extra Hard上Flash从27.86下降到33.51（反向提升），但在Mini上从16.47下降到15.30；在BIRD-SQL上Flash从45.76下降到42.18（-3.58分），Mini从39.34下降到37.32。这表明激进的解码策略对不同任务类型的影响不一致，当前的阈值配置可能需要更细粒度的领域自适应机制。改进方向：可以探索基于任务类型的自适应阈值调节，或引入轻量级的任务分类器来自动选择最优模式。第二，论文虽然提出了MBE机制，但对其计算开销的分析不够详细——Table 4显示MBE使平均TPF从5.82下降到5.14（Flash），但没有提供端到端的延迟数据。改进方向：需要更全面的效率分析，包括实际延迟、内存占用和批处理能力。第三，RL训练的具体细节（如奖励函数设计、训练数据规模、超参数选择）描述相对简略，可能影响复现性。改进方向：应提供更详细的RL训练配方和消融研究。第四，论文没有与最新的AR模型（如DeepSeek-R1、Claude等）进行直接的速度-质量对比，使得速度优势的实际意义不够明确。

论文提出了几个重要的未来研究方向。首先，作者明确提出将编辑能力整合到强化学习训练中是关键的下一步——当前的RL训练主要优化M2T能力，如果能在RL阶段也优化T2T编辑能力，可能进一步提升模型的自纠错性能。其次，探索编辑能力在更多场景中的应用，例如多轮对话中的上下文修正、长文档生成中的一致性维护等。第三，研究更智能的编辑策略——当前的T2T编辑基于简单的置信度阈值，未来可以探索基于语义理解的智能编辑，例如识别事实性错误并进行针对性修正。第四，从本文的成果可以延伸出几个有前景的方向：（1）将可编辑解码范式扩展到其他类型的扩散模型（如连续扩散）；（2）探索编辑能力与思维链（Chain-of-Thought）推理的结合，使模型能够在推理过程中自我质疑和修正；（3）研究编辑机制在多模态扩散模型中的应用；（4）开发更高效的MBE实现，使其能够在保持质量提升的同时最小化速度损失。此外，论文提到的速度-质量权衡与数据分布的关系（结构化数据vs通用对话）也值得深入研究，这可能揭示扩散模型的内在归纳偏置。

从复现性角度来看，本文的情况较为复杂。有利因素包括：论文提供了详细的训练三阶段流程（CPT、SFT、RL），使用了已知的开源框架（dFactory用于CPT/SFT，AReaL用于RL，SGLang用于推理），并且有明确的模型发布计划（LLaDA2.1-Mini和LLaDA2.1-Flash）。不利因素包括：（1）论文未明确说明模型权重和代码是否开源，这对于社区复现至关重要；（2）100B规模的模型训练需要巨大的算力投入，即使使用AReaL框架的分布式优化，普通研究机构也难以复现；（3）RL训练涉及的EBPO框架、向量化似然估计等技术细节在论文中描述相对简略，完整的实现可能需要参考AReaL和ASystem的内部文档；（4）训练数据的具体组成和规模未详细说明，而这对结果的影响可能很大；（5）双阈值 $\tau_{mask}$ 和 $\tau_{edit}$ 的具体取值在论文中未明确给出，需要从实验中搜索。总体而言，完整复现本文结果的难度较高，但核心算法思想（T2T编辑、双阈值解码）可以在较小规模上验证。