Source-linked AI summary
Test-Time Scaling in Diffusion LLMs via Hidden Semi-Autoregressive Experts
Jihoon Lee, Hoyeon Moon, Kevin Zhai, Arun Kumar Chithanar, Anit Kumar Sahu, Soummya Kar, Chul Lee, Souradip Chakraborty, Amrit Singh Bedi
TL;DR
dLLMs can generate tokens in flexible orders, but how to use that flexibility at inference remains an open question. HEX is a training-free method that ensembles different block schedules with majority voting. It reports GRPO-level results across four benchmarks, with reasoning accuracy gains of up to 3.56×.
Problem
The best way to use dLLMs’ flexible masking strategies at inference remains open, and confidence-based approaches can overlook sequential structure learned from training data.
Source: §1 · pp. 1–2Method
HEX is a training-free inference method that marginalizes across block schedules and aggregates predictions through majority voting.
Source: §1 · p. 2Results
HEX achieves GRPO-level results on GSM8K, MATH, ARC-C, and TruthfulQA, with reasoning accuracy improving by up to 3.56× (from 24.72% to 88.10%).
Source: §1 · p. 3Takeaways & Limitations
The sequence of masking plays a critical role in dLLM inference, and different schedules can activate distinct semi-autoregressive experts learned during training.
Source: §1 · p. 2, §7 · p. 10HEX requires more computation at test time, has mainly been evaluated on reasoning tasks, and lacks an established theoretical understanding.
Source: §7 · p. 11
Abstract
from arXiv · showhide
Diffusion-based large language models (dLLMs) are trained flexibly to model extreme dependence in the data distribution; however, how to best utilize this information at inference time remains an open problem. In this work, we uncover an interesting property of these models: dLLMs trained on textual data implicitly learn a mixture of semi-autoregressive experts, where different generation orders reveal different specialized behaviors. We show that committing to any single, fixed inference time schedule, a common practice, collapses performance by failing to leverage this latent ensemble. To address this, we introduce HEX (Hidden semiautoregressive EXperts for test-time scaling), a training-free inference method that ensembles across heterogeneous block schedules. By doing a majority vote over diverse block-sized generation paths, HEX robustly avoids failure modes associated with any single fixed schedule. On reasoning benchmarks such as GSM8K, it boosts accuracy by up to 3.56X (from 24.72% to 88.10%), outperforming top-K margin inference and specialized fine-tuned methods like GRPO, without additional training. HEX even yields significant gains on MATH benchmark from 16.40% to 40.00%, scientific reasoning on ARC-C from 54.18% to 87.80%, and TruthfulQA from 28.36% to 57.46%. Our results establish a new paradigm for test-time scaling in diffusion-based LLMs (dLLMs), revealing that the sequence in which masking is performed plays a critical role in determining performance during inference.
1 Introduction
The introduction asks how to choose dLLM inference schedules when confidence-based masking can overlook structure learned from sequential language data. It identifies hidden semi-autoregressive experts and proposes HEX, which ensembles schedules without training and reports strong reasoning results.
1 Introduction
- Confidence-based methods can overlook sequential structure in training data and may perform worse than random unmasking.The introduction identifies why this failure occurs and how to avoid it as an open question.Source: §1 · pp. 1–2
- dLLMs implicitly learn semi-autoregressive experts during training, and varying block size can activate different experts with distinct masking preferences.These preferences often reflect left-to-right or autoregressive-like generation orders.Source: §1 · p. 2
- HEX marginalizes across block schedules and aggregates predictions by majority vote, avoiding reliance on a single decoding path.The method is training-free and treats block size and order as latent variables for test-time scaling.Source: §1 · p. 2
- The paper frames the inference problem as a limitation of current strategies that overlook implicit training-time structural biases, risking suboptimal performance and unstable generation.The figure compares HEX with Top-K, Top-K margin [1], random expert selection [2], and training-based GRPO [3].Source: §1 · p. 2, Figure 1 · p. 2
- HEX achieves GRPO-level results on GSM8K, MATH, ARC-C, and TruthfulQA; reasoning accuracy rises by up to 3.56× (from 24.72% to 88.10%).It also reports MATH improving from 16.40% to 40.00%, ARC-C from 54.18% to 87.80%, and TruthfulQA from 28.36% to 57.46%.Source: §1 · p. 3
2 Related Work
Related work traces dLLMs from discrete masked diffusion approaches to large-scale models that can rival similarly sized autoregressive systems. It contrasts this progress with sparse dLLM inference-time research, where many gains have come from training-time methods.
2 Related Work
- Large-scale dLLMs, including DiffuLLaMA, Dream, and LLaDA, can rival similarly sized autoregressive models, including on complex reasoning.The field moved from continuous diffusion over latent text, which faced scalability and discretization challenges, toward masked diffusion for discrete language.Source: §2 · p. 3
- Earlier discrete masked-diffusion approaches [12] [13] offered a more tractable alternative to continuous diffusion over latent text representations [11].The passage also notes that diffusion models have achieved state-of-the-art performance in image generation [8].Source: §2 · p. 3
- dLLM inference-time methods remain sparse, while many performance gains have come from training-time approaches such as GRPO and temporal consistency reinforcement.Autoregressive inference-time scaling has been studied through approaches including chain-of-thought prompting and self-consistency.Source: §2 · p. 3
3 Problem Formulation
The formulation treats a masked diffusion language model as a conditional denoiser trained to recover masked tokens from visible context. At inference, the model must follow a trajectory of masking decisions, and the schedule matters because training leaves some conditionals poorly learned.
3 Problem Formulation
- A masked diffusion language model specifies a conditional denoiser that predicts masked tokens x[M] given the unmasked context x[M c].The training objective maximizes the likelihood of the ground-truth tokens at masked positions.Source: §3 · p. 3
- Training samples a masked-token count uniformly and selects a subset of positions of that size, averaging prediction losses over masked tokens.This objective can be viewed as averaging next-token losses over token permutations, an any-order objective [2].Source: §3 · p. 3
- Inference generates text step by step by choosing which tokens to reveal along a trajectory of masks that partitions the output positions.The ideal goal is to select a trajectory that maximizes the natural log-likelihood functional, but this is empirically intractable.Source: §3 · pp. 3–4
- Because training includes all mask patterns, some poorly contextualized conditionals are ill-posed or rarely observed, leaving the model's performance sensitive to the inference schedule.Strong performance depends on schedules aligned with sequential biases learned during training.Source: §3 · p. 4
- The open question is how to choose an inference strategy that reflects what the model learned when the optimal masking trajectory is under-specified.Source: §3 · p. 4
4 Limitations of SoTA and Our Key Insight
On GSM8K, random unmasking outperforms reverse Top-K margin decoding, which can collapse to [AfterEoT]. The paper attributes such schedule failures to differing masked-context experts and proposes aggregating their predictions rather than committing to one path.
From failure to mechanism
- HEX treats semi-autoregressive schedules as queries to different conditional experts and aggregates their predictions to avoid premature commitment to a single context.The paper frames this aggregation as marginalizing over contexts.Source: §4 · pp. 4–5
- Ill-posed masked contexts can be poorly learned and disproportionately favor special tokens such as [AfterEoT], making local confidence unreliable.The training objective averages across varied maskings, including contexts with little useful information.Source: §4 · pp. 4–5
- Consecutive greedy choices do not guarantee maximizing J, while globally searching all possible trajectories is practically infeasible.Source: §4 · p. 4
Failure of existing inference methods for reasoning tasks
- 50.87% vs 24.72% accuracy: random decoding outperforms Top-K margin on GSM8K, where reverse unmasking often collapses to [AfterEoT].Over 55.5% of Top-K margin runs suffered this collapse.Source: §4 · p. 4, Figure 2 · p. 5
Our Key Insight: dLLM is an implicit mixture of experts
- Most contexts concentrate probability on the correct token “Bell,” but some are flat or incorrect, so not every conditional acts as an expert.The figure shows that the distribution for “Bell” changes across masking conditions.Source: §4 · p. 5, Figure 3 · p. 6
5 Hidden Semi-autoregressive Experts for Test-time Scaling
HEX approximates aggregation over latent semi-autoregressive experts by sampling block schedules and voting over their predictions. Semi-AR schedules preserve left-to-right prefix structure while allowing within-block denoising, avoiding the collapse seen with fully parallel decoding.
5 Hidden Semi-autoregressive Experts for Test-time Scaling
- Because likelihood marginalization is infeasible, HEX samples semi-AR schedules as proxies for experts and uses majority vote to aggregate their predictions.Each schedule queries one expert; the mode of the sampled predictions approximates the final decision.Source: §5 · p. 6
- Semi-AR decoding reveals consecutive blocks left-to-right while denoising within each block, preserving prefix structure while allowing parallelism inside blocks.Source: §5 · p. 7
- Table 1 reports 0% collapse and much higher accuracy for semi-AR decoding than non-AR decoding.Moderate left-to-right block sizes avoid the [AfterEoT] collapse observed with a single large block.Source: Table 1 · p. 7, §5 · p. 7
- Across block sizes, answers can diverge; selecting the most frequent answer favors Andrew over alternatives such as Michael or David.Source: Figure 4 · p. 7
6 Experiments
Experiments evaluate HEX on four reasoning benchmarks, its accuracy-compute scaling, and ablations of schedule diversity and candidate selection. HEX leads the reported benchmark comparisons, while more voting samples improve accuracy and reduce ties.
6.1 Setup
- Experiments use LLaDA-8B-Instruct, four reasoning benchmarks, and task accuracy; HEX draws 25 samples across block sizes [8, 16, 32, 64, 128].Five samples are drawn per block size at temperature 0.9.Source: §6.1 · p. 8
6.2 Main Results: HEX Establishes a New State-of-the-Art
- HEX leads all four benchmarks and beats d1 (GRPO): GSM8K 88.10% vs. 79.80%, MATH 40.00% vs. 37.20%, and ARC-C 87.80% vs. 82.68%.The reported comparison includes training-free and fine-tuned baselines.Source: §6.2 · pp. 8–9
6.3 Analysis of Scaling and Compute Trade-off
- Across all four benchmarks, more voting samples monotonically improve accuracy and reduce tie rates, while sampling more trajectories linearly raises compute cost.The paper describes this as a tunable accuracy-compute trade-off without retraining.Source: §6.3 · p. 9, Figure 6 · p. 9
6.4 Ablation Studies
- On GSM8K, increasing dynamic trajectories from 5 to 30 raises accuracy from 81.96% to 84.15% and reduces the tie rate by more than half.Structured diversity—a fixed block set with multiple seeds—performs better than dynamic schedules.Source: §6.4 · p. 9
- On ARC-C, performance is worse than Random decoding (70.05% vs. 60.84%); HEX’s frequency-based majority vote achieves 74.57%, and ensemble agreement is identified as key to HEX’s success.Source: §6.4 · p. 10
- Ties default to the output from the smallest block size; combining frequency and log-likelihood offers no clear advantage.Source: §6.4 · p. 10
7 Conclusion and Limitation
The conclusion presents HEX as a training-free way to combine schedules and use dLLMs’ hidden semi-autoregressive experts. The authors note higher test-time computation, a focus on reasoning evaluations, and no theoretical understanding of HEX.
7 Conclusion and Limitation
- HEX combines multiple generation schedules through expert voting, requiring no extra training and outperforming standard methods and GRPO on challenging reasoning benchmarks.The conclusion describes this as using the model’s hidden team rather than relying on a single schedule.Source: §7 · p. 10
- The evaluation has focused mainly on reasoning tasks, while creative generation and a theoretical understanding of HEX remain future-work areas.The authors also identify increased test-time computation as a limitation.Source: §7 · p. 11
A.1 Qualitative analysis of Baselines vs. HEX
- Figure 7 presents generated responses from different decoding strategies for a prompt about the cost of lollipops and candies.The supplied example asks for the cost of 10 lollipops and 10 candies; the figure caption does not describe the responses’ contents.Source: Figure 7 · p. 14
A.2 Qualitative analysis of semi-AR vs. non-semi-AR
The qualitative analysis links confidence-based decoding problems to generation order: non-semi-AR decoding can collapse into repeated [AfterEoT] tokens, and answer-first generation can prevent later reasoning from influencing the answer. The appendices describe tie handling, dynamic block settings, and related sampling approaches.
A.2 Qualitative analysis of semi-AR vs. non-semi-AR
- With block size 256, removing positional constraints can make generation collapse into [AfterEoT] tokens; block size 32 instead enforces left-to-right progress across blocks.The figure describes two tokens being unmasked per step, with output length 256 and 128 steps.Source: Figure 8 · p. 15
- Confidence-based decoding can produce an answer before its justification, leaving subsequent reasoning unable to influence the answer unless those answer tokens are revised.The paper also notes that answer-token probabilities can change substantially during reasoning, so confidence in an early answer is unreliable [26].Source: §A.2 · p. 16
B.1 Additional context of related work
- The related-work discussion contrasts reward-signal-dependent Best-of-N selection with majority voting and reward-guided decoding; Table 4 compares BoN and semi-AR experts on GSM8K.The caption identifies output length and block sizes in the table, but provides no cell values.Source: §B.1 · p. 16, Table 4 · p. 16
B.2 Details of dynamic HEX block settings
- Dynamic block sizes are randomly chosen and adjusted to match total output length; with 128 unmasking steps over 256 tokens, block sizes are multiples of two.Source: Figure 10 · p. 17
B.3 Experimental results of HEX’s tie-breaking methods
The section describes the default rule for resolving tied majority-vote outputs and reports evaluations of alternative tie-breaking methods. It also notes a decline in performance as block size approaches output length.
B.3 Experimental results of HEX’s tie-breaking methods
- The default tie-break selects the output from the smallest block size, because performance tends to degrade as block size approaches output length.Likelihood-based selection did not provide significant benefit.Source: §B.3 · p. 17, Figure 11 · p. 18
- Table 5 compares negative-log-likelihood selection, smallest-block-size selection, and treating a tie as correct when any tied option is correct.The passage says the TIED: any results highlight that HEX majority voting works well across datasets, without giving values here.Source: Table 5 · p. 17
- Figure 11 reports accuracy across tasks as block size increases in increments of 16 and notes a sharp drop in the 128–256 range.The block count is set to 256 // B when B divides 256, and otherwise to (256 // B) + 1.Source: Figure 11 · p. 18
B.4 Analysis of decoding latency across inference methods
The latency analysis reports that using more HEX candidates improves performance and lowers the tie rate, alongside an increase in inference cost. The further analysis finds that no single semi-autoregressive block size consistently dominates across prompts.
B.4 Analysis of decoding latency across inference methods
- More candidate samples in HEX majority voting improve performance and lower the tie rate, while increasing inference cost.Source: §B.4 · p. 18
- Table 6 reports GSM8K generation time in seconds and the corresponding scaling factor for different inference methods.Source: Table 6 · p. 18
B.5 Further semi-AR analysis
- Across GSM8K prompts, no semi-AR block size consistently dominates, and it is rare for outputs from all tested block sizes to fail together.The figure caption notes that block size 256 is the exception to the absence of a consistent pattern; varying block sizes change accumulated context and outputs.Source: Figure 12 · p. 19, Figure 13 · p. 19
B.6 Ablation study of HEX across different output lengths
HEX maintains robust performance across output lengths of 128, 256, and 512, surpassing other inference methods and GRPO-trained models. Its majority-vote results also exceed the average performance of the individual semi-autoregressive experts.
B.6 Ablation study of HEX across different output lengths
- Across output lengths of 128, 256, and 512, HEX surpasses the other inference methods and GRPO-trained models.The ablation reports robust performance irrespective of output length.Source: §B.6 · p. 20, Table 7 · p. 20
- Across output lengths of 128, 256, and 512, HEX exceeds the mean accuracy of the individual samples used in its majority vote.The authors interpret this gap as evidence of diverse inference behaviors among hidden semi-autoregressive experts.Source: §B.6 · p. 21, Table 8 · p. 21
- Table 8 labels output length before brackets and lists the block sizes used inside brackets.Its bottom three rows report HEX; its top three rows report the mean accuracy of majority-vote samples.Source: Table 8 · p. 21