arXiv:2605.13935cs.LGcs.CL2026-05

解决扩散语言模型后训练中路径锁定问题,提升生成多样性与稳定性。

Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models

  • 引入轨迹平衡目标,防止奖励驱动更新过度集中于少数生成路径。
  • 在数学推理和代码生成任务中,所有长度设置下均优于基础模型,且随采样预算增加性能持续提升。
  • 适用于需要高可靠性和多样性的生成场景,尤其适合对生成质量要求高的研究者。

扩散语言模型是自回归模型的有前景替代方案,但其后训练方法大多沿用奖励最大化目标。我们识别出一种核心缺陷——轨迹锁定:基于奖励的采样更新过度集中在少数去噪路径上,导致重复采样时对其他正确解的覆盖不足。为此,我们提出TraFL(轨迹流平衡),一种以冻结参考模型为锚点的奖励倾斜目标分布的轨迹平衡目标。通过设计适配扩散模型的序列级代理损失和学习的提示依赖归一化,使其在扩散语言模型中可实践。在数学推理与代码生成基准测试中,TraFL是唯一在所有评估长度设置下均优于基线模型的方法,且随着采样预算增加性能持续提升。改进效果可迁移至未见评估:在Minerva Math上始终优于基线,在LiveCodeBench每个难度分组上均为最强方法。

原文摘要 · Abstract (English)

Diffusion language models are a promising alternative to autoregressive models, yet post-training methods for them largely adapt reward-maximizing objectives. We identify a central failure mode in this setting we call trajectory locking: sampled reward-driven updates over-concentrate probability mass onto a narrow set of denoising paths, reducing coverage of alternative correct solutions under repeated sampling. To address this, we propose TraFL (Trajectory Flow baLancing), a trajectory-balance objective that trains the policy toward a reward-tilted target distribution anchored to a frozen reference model. We make this practical for diffusion language models with a diffusion-compatible sequence-level surrogate and a learned prompt-dependent normalization. Across mathematical reasoning and code generation benchmarks, TraFL is the only evaluated post-training method that improves over the base model in every benchmark-length setting, with gains that persist as the sampling budget increases. The improvements transfer to held-out evaluations: TraFL stays above the base model on Minerva Math and is the strongest method on every LiveCodeBench difficulty split.

扩散模型语言生成强化学习后训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。