PaperlayerView exact revision on arXiv

Source-linked AI summary

Who Should Lead Decoding Now? Tracking Reliable Trajectories for Ensembling Masked Diffusion Language Models

Heecheol Yun, Joonhyung Park, Joowon Kim, Eunho Yang

arXiv:2606.16281v1cs.CLcs.AI

TL;DR

Ensembling heterogeneous MDLMs is underexplored because their non-left-to-right decoding produces incompatible intermediate states. TIE relays confidence-guided trajectories between models and consistently improves performance across diverse reasoning tasks without additional training.

  • Problem

    Ensembling heterogeneous MDLMs remains underexplored, and their flexible non-left-to-right decoding makes shared next-token coordination difficult.

  • Method

    TIE iteratively scores models’ confidence-guided decoding trajectories and relays the highest-scoring partially denoised state across models.

  • Results

    TIE consistently improves performance over individual MDLMs across general reasoning, mathematics, coding, and planning tasks.

  • Takeaways & Limitations

    Without additional training, confidence-guided intermediate exchange enables MDLMs to recover from suboptimal trajectories and contribute complementary strengths during denoising.

  • Takeaways & Limitations

    TIE is less effective when constituent-model performance gaps exceed 15%, and the experiments cover only two-model ensembles.

Abstract

from arXiv · show

Masked Diffusion Language Models (MDLMs) have emerged as a distinct paradigm for sequence generation. As MDLMs become diverse in capabilities and knowledge coverage, an important question is how to combine their knowledge. Toward this, we first investigate the unique decoding dynamics of MDLMs. We find that successful generations exhibit stable confidence dynamics over answer-relevant positions, while unreliable trajectories can often be corrected by injecting promising intermediate states from other models. Guided by this observation, we propose $\textbf{TIE}$ ($\textbf{T}$rajectory-based $\textbf{I}$terative $\textbf{E}$nsembling), a knowledge fusion framework in which MDLMs iteratively identify reliable decoding trajectories and relay them across models. TIE tracks confidence dynamics over answer-relevant positions to determine which model currently follows a more reliable trajectory and selectively transfers partially denoised sequences across models. As the model on the more promising trajectory often changes across denoising steps, TIE allows different models to contribute complementary strengths at different stages of generation. Strong performance across diverse reasoning tasks, along with our analyses, suggests that TIE offers a practical approach to the underexplored problem of MDLM ensembling.

1 Introduction

MDLM ensembling is underexplored because their flexible, non-left-to-right decoding prevents conventional next-token ensemble strategies from transferring cleanly across models. TIE addresses this by tracking confidence dynamics over answer-related tokens and iteratively relaying promising intermediate trajectories, enabling models to contribute complementary strengths across denoising steps.

  • MDLMs generate sequences through parallel denoising and iterative remasking, making them compelling alternatives to autoregressive language models.
  • Ensembling heterogeneous MDLMs remains underexplored because models can occupy different partially denoised sequences at each step, preventing a shared next-token definition.Conventional approaches based on averaging next-token distributions or routing to the more confident model are therefore not directly applicable.
  • Answer-related tokens tend to exhibit more stable denoising trajectories, allowing confidence dynamics to identify which model is more likely to produce a correct answer before generation finishes.
  • TIE iteratively relays trajectories from the currently more confident model to other MDLMs, transferring reliable intermediate responses to guide deviating models toward promising generation paths.The process monitors confidence dynamics over answer-related tokens and repeats knowledge transfer periodically throughout generation.
  • TIE improves over individual MDLMs across general reasoning, mathematics, coding, and planning, with effectiveness strongest when constituent models have comparable and strong individual capabilities.

2 Preliminaries

MDLMs define sequence generation through a masked forward diffusion process and a learned reverse process that updates only masked positions. During inference, ancestral sampling iteratively denoises a fully masked sequence, with token-unmasking policies determining which positions are revealed.

  • Masked diffusion language models: MDLMs represent intermediate sequence states with masked tokens and apply a predefined noise schedule in the forward diffusion process.The formulation distinguishes a clean token sequence, a one-hot mask representation, and an intermediate noisy state.
  • Masked diffusion language models: For earlier diffusion levels, the posterior q(z_s|z_t) is analytically expressible, and it becomes deterministic when z_t is unmasked.This posterior structure supports the reverse-process parameterization used by MDLMs.
  • Masked diffusion language models: The learned reverse process approximates the posterior only on masked positions, while preserving tokens that have already been unmasked.A neural network f_θ estimates the clean-token distribution from the noisy state and diffusion time.
  • Ancestral sampling and unmasking: Inference discretizes diffusion into T denoising steps, begins from a fully masked sequence, and samples successive reverse states while applying token-unmasking policies.Prior policies include confidence-based, thresholding, and KL-divergence criteria.

3 Towards Effective MDLM Ensembling

MDLM ensembling is motivated by two decoding observations: correct trajectories have more stable, confident answer-token dynamics, and sharing partial reliable trajectories can steer weaker models toward correct answers. TIE operationalizes these insights by evaluating and relaying promising intermediate trajectories across models during decoding.

  • Decoding dynamics: Correct decoding trajectories show more stable and confident answer-token dynamics than incorrect trajectories.Incorrect samples have roughly twice the token change count of correct samples, while correct samples maintain higher top-1 probability and probability margin and lower entropy.
  • Cross-model collaboration: Partial reliable trajectories from one model can guide another model’s subsequent decoding toward a correct answer.This tests cross-model collaboration by giving an initially incorrect model an early trajectory portion generated by another model.
  • Analysis setup: The analysis uses MMLU and GSM8K with LLaDA-1.5 and Dream-7B-Instruct under semi-autoregressive decoding with block size 16 and low-confidence remasking.Generation lengths are 128 tokens for MMLU and 256 tokens for GSM8K.
  • Decoding dynamics: Token Change Count C(n) measures cumulative top-1 answer-token changes across decoding steps, with lower values indicating greater stability.The count is accumulated over masked answer-token positions at each step and consecutive decoding steps.

4 TIE: Trajectory-based Iterative Ensembling

TIE is an MDLM-specific ensemble method that uses confidence dynamics at answer-relevant positions to iteratively generate, assess, and relay partial decoding trajectories. It selects the final response by answer-token stability, with top-1 probability resolving ties.

  • Core procedure: TIE cycles through independent trajectory generation, confidence-based assessment, and relaying the highest-scoring trajectory across all constituent models.Each model decodes for n steps, trajectories are scored, and the selected trajectory replaces every model’s current partial trajectory.
  • Trajectory Generation: During each round, every model independently decodes for n steps, while answer-token positions are unmasked only after reasoning-token positions are fully unmasked.Here, n is the ensemble interval.
  • Trajectory Assessment: TIE scores trajectories over masked answer-token positions using negative token change count, top-1 probability, negative entropy, or probability margin.Higher scores indicate greater confidence, and the metrics target stable and confident decoding behavior over answer tokens.
  • Trajectory Assessment: For logit-based metrics, TIE averages each trajectory’s score across all constituent models to reduce dependence on differing confidence calibration.This cross-model scoring favors trajectories that are consistently confident across models rather than only favored by their source model.
  • Final Selection: After decoding, TIE selects the candidate with the lowest answer-token change count, breaking ties with the highest top-1 probability.The generation-assessment-relay cycle repeats until any constituent model completes decoding.

5 Experiments

Experiments across eight benchmarks and four MDLMs show that TIE selects stronger decoding trajectories and benefits from intermediate, confidence-based knowledge fusion. Analyses find that TIE works best with comparable strong models, while ablations show sensitivity to ensemble interval and compatibility with decoding acceleration.

  • Trajectory selection: TIE improves over individual models across diverse domains, with token change count achieving the best results on four of eight benchmarks.Token change count tracks the full decoding history and remains robust when applied to code-specialized MDLMs.
  • Model compatibility: TIE is most effective when constituent models have comparable, strong capabilities; large performance gaps can cause weaker models to introduce noisy trajectory signals.The paper attributes this to stronger models having more reliable confidence dynamics for trajectory assessment.
  • Dynamic contribution: The highest-scoring model changes across decoding stages, allowing intermediate-generation ensembling to combine complementary strengths and outperform post-generation ensembling.Repeated trajectory aggregation progressively moves generation toward better final answers.
  • Ablations: Across ensemble intervals, TIE generally outperforms individual models, but intervals that are too small or too large weaken trajectory assessment or aggregation frequency.Short intervals may lack sufficient information for reliable assessment, whereas long intervals limit knowledge fusion.
  • Ablations: TIE is also evaluated for compatibility with MDLM decoding acceleration methods, including thresholding and top-k unmasking.Thresholding unmasks tokens above a fixed top-1 probability threshold, while top-k unmasking selects the k most confident tokens per step.

6 Related Work

Prior ensembling work primarily targets autoregressive language models, aggregating outputs at output, span, or token granularity. These methods assume participating models generate the next span or token at the same position, limiting direct applicability when generation order differs or is non-autoregressive.

  • Autoregressive Language Models Ensemble: Autoregressive language model ensembling methods aggregate model outputs at different granularities, including complete responses, spans, and individual tokens.Output-level methods first generate complete responses independently before combining them; finer-grained methods aggregate during generation.
  • Autoregressive Language Models Ensemble: Span-level methods iteratively select promising spans from candidates proposed by multiple models, often using perplexity from other models.Span candidates may consist of sequences of words or tokens.
  • Autoregressive Language Models Ensemble: Span-level and token-level ensembles require participating models to aggregate the next span or token at the same position.This requirement makes them unsuitable for non-autoregressive generation or settings where models use different token-generation orders.

7 Conclusion

TIE is a knowledge fusion framework for heterogeneous Masked Diffusion Language Models that collaborates through confidence-guided intermediate exchange. Without additional training, it improves performance across diverse reasoning tasks by enabling trajectory recovery and complementary contributions during denoising.

  • TIE rethinks collaboration among heterogeneous Masked Diffusion Language Models through continual intermediate exchange guided by confidence dynamics over answer-related tokens.
  • TIE helps models recover from suboptimal trajectories and contribute complementary strengths throughout denoising.
  • Without additional training, TIE consistently improves performance across diverse reasoning tasks.

Limitations · A Experimental Details

TIE is effective across varied domains and generation settings, but its effectiveness declines when constituent models differ too substantially in performance. The paper identifies model routing as a potential way to strengthen TIE in such cases.

  • Limitations: TIE has been effective across various domains and generation settings.
  • Limitations: Several aspects of TIE remain open for improvement.
  • Limitations: TIE becomes less effective when the performance gap between constituent models is excessively large.
  • Limitations: Performance gaps greater than 15% are given as an example of this limitation.
  • Limitations: Large performance disparities are described as a common challenge in LLM ensembling.
  • Limitations: Model routing is suggested as a mechanism that could further strengthen TIE.

A.1 Dataset Details … A.4 Answer-token positions

The appendix specifies dataset selection and evaluation conventions, hardware and prompt implementations, and the answer-token positions used to measure confidence dynamics. It also defines filtering rules for answer positions, including EOS exclusion and an eight-position limit in longer code-generation answers.

  • A.1 Dataset Details: MMLU* contains 35 subjects where both LLaDA and Dream achieve over 60% accuracy.The subset is used to examine whether ensembling stronger models yields larger gains.
  • A.1 Dataset Details: Evaluation uses the test split when ground-truth answers are available and otherwise uses the validation split; MBPP uses its sanitized version.The sanitized MBPP version filters out low-quality samples.
  • A.2 Hardware: When using TIE, each model is loaded onto a separate RTX 3090 GPU with bfloat16 precision.
  • A.3 Prompts: Multiple-choice and math prompts follow simple-evals, Countdown follows Wang et al. (2026), and HumanEval and MBPP prompts are provided.
  • A.3 Prompts: For coding tasks, the prompt replaces the Answer: suffix with ```python and requests only the implemented Python function without explanation.
  • A.4 Answer-token positions: Answer-token positions follow an appended Answer: suffix: preceding tokens are reasoning positions, following tokens are answer positions, and answer positions unmask only after reasoning positions are fully unmasked.
  • A.4 Answer-token positions: Confidence dynamics exclude answer positions whose top-1 token is EOS because those positions do not represent the model’s actual answer.
  • A.4 Answer-token positions: In domains with substantially longer answer sequences, confidence dynamics are evaluated only on the first eight masked answer positions to reduce noise from later tokens.

B Effect of Token Change Count Normalization · C Effect of Cross-Model Scoring

The paper normalizes token change counts by the number of masked answer-token positions because models can otherwise appear less stable simply by having more answer positions. It also uses cross-model scoring to reduce calibration bias and favor trajectories consistently supported across constituent models.

  • B Effect of Token Change Count Normalization: TIE uses a normalized token change count in its history-based trajectory assessment.The normalization is applied instead of directly using the unnormalized token change count.
  • B Effect of Token Change Count Normalization: Normalizing by masked answer-token positions accounts for differences in answer length across models.Without this adjustment, models with more answer-token positions may have higher token change counts even when their decoding dynamics are not genuinely less stable.
  • B Effect of Token Change Count Normalization: Compensating for differences in masked answer-token positions improves performance.Table 8 reports the effect of normalizing the token change count.
  • C Effect of Cross-Model Scoring: Cross-model scoring is important for logit-based trajectory scoring because different models can have different calibration.The relevant scoring functions are top-1 probability, entropy, and probability margin.
  • C Effect of Cross-Model Scoring: Source-model-only scoring can bias trajectory comparisons toward overly confident models regardless of actual reliability.This issue arises when logit-based scoring functions disproportionately favor a model’s trajectories.
  • C Effect of Cross-Model Scoring: Cross-model scoring evaluates each trajectory under all constituent models and selects the one with the highest average confidence.This favors trajectories consistently supported across the models rather than only by their source model.
  • C Effect of Cross-Model Scoring: Cross-model scoring outperforms source-model-only scoring in Table 9.The comparison is between TIE with cross-model scoring and TIE without cross-model scoring.

D Ablation on Final Response Selection Strategies · E When Does the Ensemble-Leading Model Change? · F Qualitative examples

The final response must be selected after trajectory generation, with the best strategy depending on constituent-model performance gaps. TIE’s leading model changes mainly early in decoding, enabling early knowledge exchange, while qualitative examples compare individual and ensembled responses.

  • D Ablation on Final Response Selection Strategies: After the final trajectory-generation step, TIE requires a separate strategy to select one response from the M candidate responses.No trajectory relay occurs after all answer-token positions have been unmasked.
  • D Ablation on Final Response Selection Strategies: TIE evaluates three final-response strategies: lowest TCC, best model, and most selected.These select the response with the lowest ˜C(T) m, from the best-performing individual model, or from the most frequently selected model, respectively.
  • D Ablation on Final Response Selection Strategies: All three final-response strategies achieve strong performance, but their relative effectiveness depends on the performance gap between constituent models.The optimal strategy is not universal across relative model-strength settings.
  • D Ablation on Final Response Selection Strategies: When constituent models differ substantially, selecting the best-performing model’s response is favored, whereas comparable models favor the response with the lowest ˜C(T) m.The comparison uses token change count for trajectory assessment.
  • E When Does the Ensemble-Leading Model Change?: The model producing the highest-scoring trajectory changes frequently throughout generation, with changes occurring more often during early decoding stages.Figure 3 measures the proportion of decoding steps in which the highest-scoring model changes.
  • E When Does the Ensemble-Leading Model Change?: Early changes allow knowledge exchange before a clear answer trajectory emerges, enabling TIE to progressively converge toward higher-quality responses.The ensemble-leading model can therefore vary across decoding stages rather than remaining fixed.
  • F Qualitative examples: Qualitative examples compare responses from individual models with responses produced by ensembling those same models through TIE.These comparisons are presented in Tables 11 and 12.
Loading 2606.16281v1…