Most post-training research for language models assumes an autoregressive architecture. The growing interest in diffusion-based language models, which generate text by iteratively denoising over the full sequence rather than left-to-right token prediction, has naturally prompted researchers to ask whether standard reward-maximising objectives transfer cleanly to this setting. A new paper from Ahmadi, Parthasarathi, and Cui argues they do not, and proposes a principled alternative. The paper is available at arXiv:2605.13935.
The Problem: Trajectory Locking
Reward-maximising post-training methods, whether PPO-style policy gradient or simpler REINFORCE variants, share a common objective: push probability mass toward outputs that receive high reward. For autoregressive models this is already known to cause mode collapse under certain conditions, but the problem takes a specific structural form in diffusion models that the authors call trajectory locking.
A diffusion language model generates text via a stochastic denoising chain. Many distinct denoising trajectories can lead to the same final output, and many distinct outputs can be correct for a given prompt. Reward-driven updates, when applied naively, concentrate probability mass onto a narrow subset of denoising paths associated with high-reward outputs seen during training. The result is a model that, under repeated sampling, produces far less diverse solutions than the base model, even when diversity would improve aggregate performance. Pass@k metrics, which measure whether at least one correct answer appears in k samples, are particularly sensitive to this failure.
This is not simply a restatement of the general mode-collapse problem. The denoising trajectory is itself a structured object, and over-concentrating on specific trajectories reduces the effective coverage of the output space in ways that are harder to detect and correct than in the autoregressive case.
Key Contributions
The authors propose TraFL (Trajectory Flow baLancing), a post-training objective grounded in the trajectory balance framework originally developed for GFlowNets. Rather than maximising expected reward, TraFL trains the policy to sample from a reward-tilted target distribution anchored to a frozen reference model. This is a distributional objective: the goal is proportional sampling from high-reward regions, not pure maximisation.
Several practical challenges arise when applying trajectory balance to diffusion language models specifically:
- The standard trajectory balance loss requires computing flow quantities over the full denoising chain, which is expensive and not straightforwardly defined for discrete sequence models.
- The normalisation constant in the reward-tilted distribution is prompt-dependent and intractable to compute exactly.
- Diffusion models operate on continuous noise schedules but produce discrete outputs, requiring careful treatment of the sequence-level reward signal.
TraFL addresses these with a diffusion-compatible sequence-level surrogate loss and a learned prompt-conditioned normalisation network. The normalisation network is trained jointly and estimates the partition function for each prompt, allowing the trajectory balance condition to be enforced without exhaustive enumeration. This is a non-trivial engineering contribution; getting the normalisation right is precisely what separates GFlowNet-style training from standard policy gradient in practice.
Results
The empirical evaluation covers mathematical reasoning and code generation benchmarks, with results reported across multiple sampling budgets (varying k in pass@k). The headline finding is that TraFL is the only post-training method in the comparison that improves over the base diffusion model in every benchmark-length setting tested. Competing reward-maximising baselines show inconsistent behaviour: gains at low k that erode or reverse at higher k, consistent with the trajectory locking hypothesis.
Key results include:
- TraFL maintains improvements over the base model on Minerva Math in held-out evaluation, where other methods degrade.
- On LiveCodeBench, TraFL achieves the best performance across all difficulty splits, suggesting the distributional objective generalises rather than overfitting to training reward signals.
- The advantage over baselines grows with sampling budget, which is the expected signature of better coverage: more diverse trajectories mean more chances to find a correct solution.
The held-out evaluation results are particularly worth attention. It is easy to construct post-training methods that improve benchmark numbers while degrading generalisation; the fact that TraFL's gains transfer to Minerva Math and LiveCodeBench, both of which are distinct from the training distribution in meaningful ways, provides some evidence that the method is doing something structurally useful rather than reward hacking.
Limitations and Open Questions
The paper's scope is necessarily narrow in some respects. The evaluation is confined to diffusion language models, and the specific architecture choices matter. It is not immediately clear how sensitive TraFL is to the quality of the learned normalisation network, or what happens when the reference model is itself poorly calibrated. The authors do not report ablations on the normalisation component in the abstract, though these would be important for practitioners considering adoption.
There is also a broader question about the compute cost of TraFL relative to simpler baselines. Trajectory balance objectives typically require more forward passes than standard policy gradient, and the joint training of the normalisation network adds overhead. For researchers working with large diffusion language models, this cost profile matters.
The connection to GFlowNet literature is intellectually productive but also raises questions the paper may not fully resolve. GFlowNet training is known to be sensitive to the choice of reward function and to require careful tuning of the flow matching conditions. Whether the practical robustness of TraFL holds across reward functions of varying quality, including learned reward models rather than ground-truth verifiers, is an open question.
From a broader perspective, this work sits at an interesting intersection. Diffusion language models have been advancing rapidly, with models like MDLM and more recent masked diffusion architectures showing competitive performance with autoregressive models on certain tasks. Post-training alignment for these models is largely unsolved, and the autoregressive playbook does not transfer cleanly. TraFL represents one of the first principled attempts to address this gap, and the trajectory balance framing is a natural fit given the sequential structure of the denoising process.
Whether distributional objectives will become the standard for diffusion LM post-training remains to be seen, but the evidence here that reward maximisation causes structural problems specific to this architecture is a useful diagnostic contribution independent of the proposed solution. Researchers working on diffusion language models should read this paper carefully, even if they end up pursuing different remedies.
Full paper: arXiv:2605.13935