Test-Time Scaling in Diffusion LLMs via Hidden Semi-Autoregressive Experts

2025-10-06

AI Summary

Diffusion language models (dLLMs) implicitly learn a semi-autoregressive mixture-of-experts structure during training, yet existing inference methods employ a single fixed scheduling strategy, failing to exploit this inherent diversity and thus limiting performance. Method: We propose HEX (Hidden Ensemble of semi-autoregressive eXperts), a training-free test-time scaling method that performs parallel generation across multiple block sizes and aggregates predictions via majority voting—enabling multi-path ensemble inference. Contribution/Results: HEX is the first approach to explicitly uncover and leverage the intrinsic semi-autoregressive expert structure of dLLMs, mitigating failure modes associated with rigid scheduling. Evaluated on GSM8K, MATH, ARC-C, and TruthfulQA, HEX achieves new state-of-the-art accuracies of 88.10%, 40.00%, 87.80%, and 57.46%, respectively—significantly outperforming existing inference paradigms.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Planning, Routing, and Scheduling: Planning with Language Models
• Search and Optimization: Learning to Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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.

Problem

• Optimizing inference performance in diffusion-based large language models
• Leveraging latent semi-autoregressive experts during test-time generation
• Overcoming limitations of fixed inference schedules through ensemble methods

Innovation

• Ensembles diverse block-sized generation paths
• Uses majority vote over semi-autoregressive experts
• Performs training-free inference via schedule ensembling