Self-Distilled Trajectory-Aware Boltzmann Modeling: Bridging the Training-Inference Discrepancy in Diffusion Language Models
Preprint 2026

Diffusion Language Models (DLMs) have recently emerged as a promising alternative to autoregressive language models, offering stronger global awareness and highly parallel generation. However, post-training DLMs with standard Negative Evidence Lower Bound (NELBO)-based supervised fine-tuning remains inefficient: training reconstructs randomly masked tokens in a single step, whereas inference follows a confidence-guided, multi-step easy-to-hard denoising trajectory. Recent trajectory-based self-distillation methods exploit such inference trajectories mainly for sampling-step compression and acceleration, often improving decoding efficiency without substantially enhancing the model's underlying capability, and may even degrade performance under full diffusion decoding.

We propose Trajectory-Aligned Optimization via Boltzmann Modeling (TABOM), a self-distilled trajectory-based post-training framework that aligns training with the easy-to-hard structure of inference. TABOM models the inference unmasking preference as a Boltzmann distribution over predictive entropies and derives a tractable pairwise ranking objective to align the model's certainty ordering with the observed decoding trajectory. Empirically, TABOM achieves substantial gains in new domains, expands the effective knowledge boundary of DLMs, and significantly mitigates catastrophic forgetting compared with standard SFT.

🔍 The key difficulty is a hidden mismatch between training and inference. During NELBO-style SFT, a DLM receives a randomly masked sequence and is asked to reconstruct all masked tokens in one shot. This objective is simple and scalable, but it implicitly treats all masked positions uniformly. In other words, a very easy token and a very hard token contribute to the same reconstruction objective without considering their order in the actual decoding process.

⚡ Inference behaves very differently. A DLM does not recover all tokens at once. It repeatedly estimates token confidence or entropy, unmasks a subset of high-confidence tokens, and then uses these newly recovered tokens as context for later decisions. This creates an easy-to-hard trajectory: easy tokens should become certain first, while difficult tokens are resolved later with richer context.

• Uniform training bias: all masked tokens are reconstructed under the same objective.
• Easy-to-hard inference bias: high-confidence tokens are decoded earlier and shape later context.
• Core question: can self-distilled trajectories be used for genuine knowledge acquisition, not merely faster decoding?

💡 Self-distilled trajectories are attractive because they lie on the pretrained model's own distributional manifold. But if we still train them with a uniform reconstruction loss, the model sees the trajectory data without learning the ordering preference behind the trajectory. This is the central issue TABOM targets.

📉 Our preliminary diagnostics compare offline ground-truth SFT and self-distilled SFT across code generation and mathematical reasoning. The observation is consistent: self-distilled trajectories reduce the optimization barrier and help avoid catastrophic forgetting, but they still provide only marginal gains when trained with the same NELBO objective. The bottleneck is not only the data source, but also the training objective.

🧭 TABOM keeps the useful part of self-distillation while changing what the model is asked to learn. Instead of only reconstructing tokens from trajectory states, it also learns the relative certainty ordering encoded by the trajectory. Intuitively, if token A is decoded before token B in an entropy-guided trajectory, the model should assign token A lower entropy than token B under the corresponding context.

As a foundation, TABOM uses self-distilled trajectory states as reconstruction contexts instead of relying only on uniformly random masks. The key difference, however, is not just the data: TABOM explicitly models and optimizes the inference-time ordering preference.

Boltzmann Modeling of Inference Preference

🔥 We view the easy-to-hard unmasking behavior as a probability distribution over possible unmasked states. States containing easier, lower-entropy tokens should receive higher probability. TABOM formalizes this preference as a Boltzmann distribution over predictive entropies, turning training-inference alignment into a principled distribution matching problem.

Energy-Based Pairwise Ranking

🎯 Ideally, training should make the model-induced unmasking distribution match the Boltzmann target above. This gives a KL alignment objective:

Directly optimizing this KL objective is intractable because it requires sampling from the global target distribution and computing the partition function. TABOM therefore derives a local pairwise ranking loss that shapes the entropy landscape so that earlier decoded tokens become more certain than later decoded tokens within a local timestep window.

✨ This local ranking view is important in practice. Comparing tokens decoded at vastly different stages can be noisy because their difficulties are inherently different. TABOM therefore ranks tokens inside a local decoding window, which better matches the step-wise nature of diffusion inference.

We evaluate TABOM on Dream-7B-Instruct and LLaDA-8B-Instruct across two post-training domains: code generation and mathematical reasoning. Each fine-tuned model is evaluated not only on its in-domain tasks, but also on out-of-distribution tasks. This matters because standard SFT can improve a target domain while damaging the pretrained model's broader capabilities.

🏆 TABOM consistently achieves the strongest in-domain average across both base models and training domains. More importantly, it preserves or improves OOD performance, avoiding the catastrophic forgetting often observed when standard SFT is applied to new-domain data.

📌 The table highlights the main story of the paper. SFT-SD is safer than ground-truth SFT because it stays closer to the pretrained model's manifold, but its gains are often small. TABOM keeps this stability while turning the self-distilled trajectories into larger performance improvements.

🔬 To check whether a model really learns the easy-to-hard structure, we introduce Trajectory Discrimination Score (TDS). At each decoding step, TDS measures the variance of predictive entropy across currently masked tokens. A high TDS means the model assigns clearly different uncertainty levels to different tokens; a low TDS means the model treats masked tokens almost uniformly.

Higher TDS indicates stronger discrimination between easy and hard tokens at the same decoding step. TABOM produces a more discriminative entropy landscape, providing mechanism evidence that its gains come from better trajectory alignment rather than simply reusing self-generated samples.

✅ In other words, TABOM does not merely expose the model to more intermediate states. It changes the shape of the model's uncertainty landscape so that the learned model behaves more like the inference process it will actually use.