Technical Summary: Self-Distilled Trajectory-Aware Boltzmann Modeling (TABOM)

1. Problem Statement

Diffusion Language Models (DLMs) offer advantages over autoregressive models, such as superior global awareness and highly parallel token generation. However, post-training DLMs via standard Supervised Fine-Tuning (SFT) faces a critical training-inference discrepancy:

• Training: Standard SFT uses a Negative Evidence Lower Bound (NELBO) objective with uniform random masking, requiring the model to reconstruct all masked tokens simultaneously in a single step. This instills a "uniform inductive bias," treating all tokens as equally difficult regardless of their inherent predictability.
• Inference: DLM inference follows a confidence-guided, multi-step "easy-to-hard" trajectory. Tokens are unmasked iteratively, starting with those having the lowest predictive entropy (highest certainty) and progressing to more uncertain tokens.

The Dilemma:
Recent attempts to bridge this gap using self-distilled trajectories (where the model learns from its own decoding traces) have yielded limited results.

1. Standard SFT on Ground Truth (SFT-GT): Improves in-domain performance but causes severe catastrophic forgetting of out-of-distribution (OOD) capabilities.
2. Standard SFT on Self-Distilled Data (SFT-SD): Preserves OOD capabilities (avoiding forgetting) but yields only marginal in-domain gains.
   The paper posits that while self-distilled trajectories offer a lower optimization barrier, naively applying standard NELBO objectives to them fails to exploit the structured "easy-to-hard" information embedded within the trajectories, leading to sub-optimal performance.

2. Methodology: TABOM

The authors propose Trajectory-Aligned Optimization via Boltzmann Modeling (TABOM), a post-training framework that aligns the training objective with the inference-time decoding dynamics.

2.1 Theoretical Foundation: Boltzmann Unmasked Distribution

The authors theoretically demonstrate that the inference unmasked distribution under an easy-to-hard decoding schedule can be modeled as a Boltzmann distribution over the ideal predictive entropy of tokens.

• Proposition: The probability of an unmasked set $U$ is proportional to $\exp(-\beta \sum_{r \in U} H_{ideal}(x_r))$, where $H_{ideal}$ is the predictive entropy.
• Implication: This formulation explicitly captures the inductive bias that the model prefers unmasking "easy" (low entropy) tokens first. The goal is to minimize the Kullback-Leibler (KL) divergence between the model's predicted distribution and this target Boltzmann distribution.

2.2 Tractable Surrogate: Pairwise Ranking

Directly optimizing the KL divergence is intractable due to the need for global sampling and computing a partition function over all possible unmasked states. TABOM introduces a tractable surrogate objective based on Pairwise Ranking:

• Mechanism: Instead of maximizing likelihood globally, the method enforces relative rankings between pairs of unmasked states. If token $r$ is decoded before token $s$ in a trajectory (implying $r$ is "easier"), the model's predicted entropy for $r$ must be lower than for $s$.
• Loss Function: A pairwise hinge loss is applied over local decoding windows. For a set of newly unmasked tokens $\Delta$, the loss penalizes cases where the entropy of an "easier" token is not sufficiently lower than that of a "harder" token:
  $$L_{rank} = \max(0, h_\theta(r) - h_\theta(s) + \gamma)$$
• Local Window: To avoid noise from comparing tokens with vastly different difficulties (e.g., first vs. last step), the ranking is applied strictly within a fixed local window size $W$ (e.g., $W=32$).

2.3 Final Objective

The total TABOM loss combines:

1. Trajectory-Aware Reconstruction: A standard cross-entropy loss where the model predicts newly unmasked tokens $\Delta$ given the context of previously unmasked tokens (replacing uniform masking).
2. Pairwise Ranking Regularization: The hinge loss described above to align the entropy landscape with the easy-to-hard trajectory.

$$L_{TABOM} = L_{reconstruction} + \lambda L_{rank}$$

3. Key Contributions

1. Theoretical Insight: The paper formally defines the inference unmasked distribution as a Boltzmann distribution over predictive entropies, identifying the "uniform inductive bias" of standard NELBO as the root cause of the training-inference mismatch.
2. Methodological Innovation: TABOM introduces a pairwise ranking objective that aligns the model's certainty ordering with observed decoding trajectories without requiring intractable partition function computations.
3. Empirical Resolution: The framework resolves the "forgetting vs. marginal gains" dilemma. It achieves substantial in-domain performance improvements while preserving out-of-distribution capabilities, effectively mitigating catastrophic forgetting.
4. Diagnostic Metric: The authors introduce the Trajectory Discrimination Score (TDS) to quantify whether a model preserves token-level uncertainty differences along the decoding trajectory.

4. Experimental Results

Experiments were conducted on two DLMs (Dream-7B and LLaDA-8B) across Code Generation (HumanEval, MBPP) and Mathematical Reasoning (GSM8K, MATH500) domains.

• Performance Gains: TABOM consistently outperformed baselines (SFT-GT, SFT-SD, dInfer, T3D).
  • On Dream-7B (Code), TABOM achieved a +7.70 gain on HumanEval and +2.60 on MBPP, compared to marginal gains for SFT-SD (+1.00, +1.20).
  • Crucially, TABOM maintained OOD performance (e.g., GSM8K remained high at 81.73 vs. the 52.33 drop seen in SFT-GT).
• Robustness: Under 2-token parallel decoding, TABOM maintained high performance, whereas SFT-GT suffered severe degradation, highlighting TABOM's alignment with the model's natural inference distribution.
• Ablation Studies:
  • Removing the ranking loss resulted in limited gains, confirming its necessity for aligning the entropy landscape.
  • Applying ranking globally (without local windows) introduced noise and degraded performance, validating the need for local window constraints.
• Trajectory Discrimination: TDS analysis showed that TABOM significantly increased the variance of predictive entropy across decoding steps (e.g., HumanEval TDS rose from 0.035 in SFT-SD to 0.711 in TABOM), proving that the model successfully learned to distinguish between easy and hard tokens, unlike standard SFT which treats them uniformly.

5. Significance and Claims

The paper claims that TABOM represents a fundamental shift in how self-distilled trajectories are utilized for DLMs:

• Beyond Efficiency: While previous trajectory-based methods focused on sampling-step compression (acceleration), TABOM uses these trajectories for genuine knowledge acquisition and performance enhancement.
• Alignment over Shortcut: By modeling the inference behavior as a Boltzmann distribution and enforcing entropy-based ranking, TABOM aligns the model's optimization landscape with its actual generation dynamics.
• Mitigation of Catastrophic Forgetting: The method successfully expands the effective knowledge boundary of DLMs into new domains without sacrificing pre-trained capabilities, a common failure mode of standard SFT.

The authors conclude that TABOM provides a robust solution to the training-inference discrepancy in Diffusion Language Models, offering a path toward more effective and efficient post-training strategies that respect the inherent "easy-to-hard" structure of generative processes.