Optimizing Remasking Schedules for Reasoning in Discrete Diffusion Models

Zeyneb N. Kaya, Radostin Cholakov, Nicole H. Ma

Keywords: Diffusion LLMs, Reinforcement Learning, Efficiency, Adaptive Compute

TL;DR: We introduce a learned remasking scheduler for discrete diffusion LLMs that improves performance and inference efficiency by enabling dynamic compute and reasoning allocation.

Abstract: Discrete diffusion language models (DLLMs) have emerged as a new paradigm of language modeling that offers improved inference efficiency and nonlinear generation and reasoning. While standard methods rely on fixed or heuristic schedules (e.g., random or confidence-based), we present LeADS, a framework that enables dynamic inference-time control for DLLMs with a learned remasking scheduler optimized for downstream performance. LeADS chooses what tokens are denoised at each diffusion step based on the internal representations of the model and dynamically allocates compute for token efficiency. On mathematical reasoning tasks, LeADS achieves 19.2% relative improvement (12 pp) over low-confidence based denoising schedules and reduces required diffusion steps by up to 15.3%.

Submission Number: 7