Technical Summary: Trajectory-Balance Post-Training for Diffusion Language Models

Problem: Trajectory Locking in Diffusion RL

Diffusion Language Models (dLLMs) offer a promising alternative to autoregressive models for reasoning and code generation. However, existing post-training methods for dLLMs largely adapt reward-maximizing Reinforcement Learning (RL) objectives (e.g., PPO, GRPO) designed for autoregressive models. The authors identify a central failure mode in this adaptation called trajectory locking.

In diffusion models, the reward is assigned only to the final denoised completion ($y$), while the model distribution is induced over the entire denoising trajectory ($\tau$). Because terminal rewards are "path-indifferent"—meaning any trajectory reaching the same correct $y$ receives the same reward—sampled policy-gradient updates tend to reinforce the specific trajectories that happen to be sampled early in training. This creates a self-reinforcing feedback loop where probability mass collapses onto a narrow set of denoising paths. Consequently, even if a model learns to produce a correct answer, it loses the ability to generate a diverse set of valid solutions (modes) under repeated sampling. The authors formalize this via Theorem 1, proving that terminal mode coverage is upper-bounded by trajectory diversity; thus, trajectory locking directly limits the model's solution coverage.

Methodology: TraFL (Trajectory Flow Balancing)

To address trajectory locking, the authors propose TraFL, a post-training objective grounded in the trajectory-balance principle from Generative Flow Networks (GFlowNets). Instead of maximizing reward directly, TraFL trains the policy to match a reward-tilted target distribution anchored to a frozen reference model.

The Objective

The target distribution $p^*(\tau | x)$ is defined as:
$$p^*(\tau | x) = \frac{1}{Z(x)} p_{\text{ref}}(\tau | x) \exp(\beta r(x, y))$$
where:

• $p_{\text{ref}}$ is a frozen reference model.
• $r(x, y)$ is the terminal reward.
• $\beta$ controls reward sharpness.
• $Z(x)$ is a prompt-dependent normalization term (partition function).

The training objective minimizes the squared residual between the current model's log-probability and the target:
$$\delta(\tau, x) = \log p_\theta(\tau | x) - \log p_{\text{ref}}(\tau | x) - \beta r(x, y) + \log Z(x)$$
$$L_{\text{TraFL}} = \mathbb{E}_{\tau \sim p_\theta} [\delta(\tau, x)^2]$$

Practical Implementation for Diffusion

Since dLLMs do not expose exact trajectory log-probabilities, the authors introduce two key components to make TraFL practical:

1. Diffusion-Compatible Sequence-Level Surrogate: They derive a tractable log-probability surrogate ($\log \hat{p}_\theta(y|x)$) based on a masked-reconstruction lower bound. This involves sampling a mask level $l$, corrupting $l$ tokens, and averaging the log-probabilities of the reconstructed tokens. A $1/l$ factor is used to normalize the score independent of the number of masked tokens.
2. Learned Prompt-Dependent Normalization: Unlike methods that estimate $Z(x)$ via softmax over a batch (which can induce concentration), TraFL learns a partition function $Z_\phi(x)$ via a small MLP. This term receives gradient signals at every step regardless of which completions are sampled, preventing the "batch-local" normalization failure seen in other distribution-matching approaches.

Key Contributions

1. Identification of Trajectory Locking: The paper formally identifies and proves that reward-maximizing objectives in dLLMs lead to a collapse of trajectory diversity, which in turn limits the coverage of distinct correct solution modes.
2. TraFL Objective: The proposal of a trajectory-balance objective that anchors the policy to a reference model, separating the role of reward (tilting terminal completions) from the reference model (allocating probability mass over plausible trajectories).
3. Practical Surrogate and Normalization: The development of a diffusion-compatible sequence-level surrogate and a learned normalization term that enables stable training without exact trajectory probabilities.
4. Empirical Superiority: TraFL is the only evaluated post-training method that consistently improves over the base model across all benchmark-length settings (GSM8K, MATH-500, HumanEval, MBPP) and sampling budgets.

Results

The authors evaluate TraFL on mathematical reasoning and code generation benchmarks against the base model (LLaDA-8B-Instruct) and strong baselines (ESPO, JUSTGRPO).

• In-Distribution Performance: TraFL improves Pass@k across all sampling budgets ($k \in \{1, 2, 4, 8, 16\}$) and temperatures. Notably, it improves Pass@1 (single-sample accuracy) while also scaling better with increased sampling budgets, unlike baselines that often sacrifice single-sample accuracy for multi-sample gains.
• Held-Out Generalization: Gains transfer to out-of-distribution benchmarks. On Minerva Math, TraFL stays above the base model. On LiveCodeBench, TraFL is the strongest method across all difficulty splits (Easy, Medium, Hard), significantly outperforming baselines on medium and hard problems.
• Diversity Analysis: Using an LLM-as-a-judge protocol, the authors show that TraFL's gains are associated with broader coverage of distinct solution strategies (e.g., different proof approaches or algorithmic decompositions) rather than just sharpening the score of a single mode. Baselines like ESPO, while improving accuracy, do not show significant diversity gains over the base model, consistent with the trajectory-locking hypothesis.

Significance and Claims

The paper claims that TraFL provides a principled solution to the "mode-seeking" failure of standard RL post-training for diffusion models. By shifting from reward maximization to trajectory-balance, the method preserves the diversity of denoising paths necessary to cover multiple valid solutions.

The authors state that TraFL is the only evaluated method that improves over the base model in every benchmark-length setting tested. They emphasize that the improvements are robust to sampling budget increases and generalize to held-out datasets, suggesting that the trajectory-balance objective effectively decouples the reward signal from the path-allocation mechanism. The work advocates for distribution-matching objectives over pure reward maximization for dLLM post-training, offering a practical framework that avoids the concentration failure modes inherent in standard policy-gradient approaches.

The paper remains modest regarding limitations, noting that the method relies on an approximate sequence-level surrogate rather than exact trajectory probabilities and uses coarse binary rewards (exact match or test pass), which may limit learning efficiency in scenarios with rare successes.