On Trajectory-Aware Training for Masked Diffusion Language Models

📄 arXiv: 2609.37974v1 📥 PDF

作者: Manuel Madeira, Amitis Shidani, Alice Bizeul, Victor Turrisi, Louis Béthune, Bhavika Devnani, Dan Busbridge, Pierre Ablin, João Monteiro

分类: cs.LG, cs.AI, cs.CL

发布日期: 2026-09-29

💡 一句话要点

提出PUMBA框架以解决掩蔽扩散语言模型的训练与推理不一致问题

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)

关键词: 掩蔽扩散模型 轨迹感知训练 自然语言处理 文本生成 深度学习

📋 核心要点

  1. 现有掩蔽扩散模型在训练和推理阶段条件不一致，导致模型性能受限。
  2. 本文提出PUMBA框架，通过轨迹感知训练，允许步骤间信息传递，优化训练过程。
  3. 实验结果显示，PUMBA在性能和功能评估次数上均优于标准微调方法，提升显著。

📝 摘要（中文）

掩蔽扩散模型（MDMs）通过逐步解掩多个标记生成文本，但训练和推理条件不同。模型在随机掩蔽序列上训练，而推理则遵循由模型自身预测形成的轨迹。本文提出PUMBA，一个统一的轨迹感知训练框架，通过在策略诱导轨迹的连续步骤上训练去噪器，允许步骤间信息传递，并通过时间反向传播共同优化。研究表明，较松的训练-推理对齐能更好地接近推理时的掩蔽模式，连续信息传递优于离散梯度估计，且随着时间反向传播跨越更多步骤，性能得到提升。PUMBA在LLaDA-8B的监督微调中表现出色，显著减少了功能评估次数。

🔬 方法详解

问题定义：本文旨在解决掩蔽扩散模型在训练与推理阶段条件不一致的问题。现有方法在随机掩蔽序列上训练，而推理则依赖于模型的预测轨迹，导致模型在推理时的表现不佳。

核心思路：PUMBA框架的核心思想是通过轨迹感知训练，使得模型在训练过程中能够利用连续步骤的信息，从而改善训练与推理之间的对齐。通过这种设计，模型能够更好地适应推理时的条件。

技术框架：PUMBA的整体架构包括多个模块：首先是策略诱导轨迹生成模块，然后是去噪器训练模块，最后是通过时间反向传播的优化模块。这些模块协同工作，实现信息的有效传递和优化。

关键创新：PUMBA的主要创新在于其轨迹感知训练机制，允许在训练过程中跨步骤传递信息，这与传统方法的独立训练方式形成鲜明对比。

关键设计：在技术细节上，PUMBA采用了连续信息传递的策略，优化了损失函数设计，并在网络结构上进行了调整，以支持跨步骤的信息流动。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

实验结果表明，PUMBA在全画布生成中，性能相当于同规模自回归模型，但功能评估次数减少了最多22%；在块扩散生成中，减少了最多26%。这些结果显示了PUMBA在效率和性能之间的良好平衡。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等。通过优化掩蔽扩散模型的训练过程，PUMBA能够提升模型在实际应用中的表现，具有重要的实际价值和未来影响。

📄 摘要（原文）

Masked diffusion models (MDMs) generate text by unmasking several tokens per step, but they are trained and sampled under different conditions. The model is trained on randomly masked sequences, whereas inference follows a trajectory shaped by the model's own predictions. Additionally, each step has no access to what the previous one computed. Recent methods narrow these limitations from separate angles, leaving open how these choices interact. We introduce PUMBA, a unified framework for trajectory-aware training that trains the denoiser on consecutive steps of policy-induced trajectories, passes information between steps, and optimizes them jointly by backpropagation through time. A controlled study of this design space shows that i) exact train--inference alignment fails due to local overfitting, whereas a looser alignment still brings training masks closer to those seen at inference; ii) passing continuous information outperforms discrete gradient estimators through the commitment at each step; and iii) performance improves as backpropagation through time spans more steps, which we support theoretically. Combined, these components match the best checkpoint of a same-size autoregressive model. Building on these findings, we scale PUMBA to supervised fine-tuning of LLaDA-8B, where it improves the trade-off between performance and number of function evaluations (NFEs) in both full-canvas and block diffusion generation. At matched performance, it needs up to 22% fewer NFEs than standard fine-tuning with twice the budget in full-canvas generation, and up to 26% fewer than standard fine-tuning for the same number of steps in block diffusion.