Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models
Authors: Saba Ahmadi, Prasanna Parthasarathi, Yufei Cui
Summary
arXiv:2605. 13935v1 Announce Type: new Abstract: Diffusion language models are a promising alternative to autoregressive models, yet post-training methods for them largely adapt reward-maximizing objectives.
Relevance
Read next because Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models overlaps with clean result "LoRA persona trained on alone emits at 23.5% when a co-trained partner learns ..., vs 0% control on Qwen2.5-7B-Instruct (MODERATE confidence)", clean result "Leakage rate is a usable signal for recovering trigger-shaped phrases on Gaperon-1125-1B without knowing the hidden trigger itself (MODERATE confidence)", clean result "Language-mismatch LoRA SFT on Qwen2.5-7B leaks the trained completion language into bystander directives the model was never trained on, absent under same-language SFT (LOW confidence)". Matching terms: code, strong, rect, under, correct, eval, rate, length. Source: arxiv cs.LG (Machine Learning).
Threat model
Potential threat/caveat for clean result "LoRA persona trained on alone emits at 23.5% when a co-trained partner learns ..., vs 0% control on Qwen2.5-7B-Instruct (MODERATE confidence)": this item discusses failure, evaluation, benchmark.
Abstract
arXiv:2605.13935v1 Announce Type: new Abstract: 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.