dSpark: Teaching Diffusion Language Models Who to Unlock First via Schedule Distillation

One of the biggest selling points of diffusion language models (dLLMs) is their ability to generate multiple tokens in parallel, but in practice there is often an awkward catch: even when each token is individually easy to predict, it can still take many decoding rounds to "fill in" the whole sentence. The problem lies in the order of "commitment"—which token is determined first directly affects the predictability of the remaining positions.

Research from VILA-Lab proposes dSpark, a framework based on schedule distillation. The core idea is: don't just focus on "which token is easy to predict," but ask "which token should be revealed first so that other positions become predictable faster."

Why "easy to predict" does not equal "should come out first"

The decoding process of diffusion language models can be understood as gradual unmasking. Which positions are chosen to "place a piece" at each step determines the difficulty of every subsequent step.

Intuitively, determining high-confidence tokens first seems safest. But the paper points out a counterintuitive phenomenon: a highly predictable token may provide little help to the context, while revealing an uncertain token early can instead make several other positions suddenly predictable. In other words, the value of a token is not only its own confidence, but also its "information spillover" to the remaining positions.

This is like a crossword puzzle: filling in an obvious short word first may not help you solve the long words, while cracking a key crossing letter first can open up an entire region. Most existing dLLM decoding strategies do not explicitly model this "chain unlocking" effect.

The relationship between diffusion language models (dLLMs) and masked language models

The backbone of a dLLM is highly similar to that of BERT-style masked language models: during training, a portion of tokens is randomly masked, and the model restores them based on the remaining context. The difference lies in the inference stage—BERT usually makes only one prediction, whereas a dLLM starts from full masking or a high masking rate and iterates repeatedly, gradually reducing the mask until the whole sentence is "filled in." In each iteration, the model produces distributions for all masked positions in parallel, then selects a subset of positions to "place" (unmask) according to some strategy; the selected tokens are fixed from then on and are not revised. Therefore, the computation of a single forward pass is approximately that of one parallel prediction, and the total time depends on the number of iterations. The fewer the rounds, the greater the speedup, and the number of rounds is directly determined by "which positions are selected in each round." This is also why scheduling strategy becomes the core lever for dLLM inference efficiency.

dSpark's two-step approach: construct a schedule, then distill it into the model

dSpark's approach is divided into two stages.

First is constructing a "good" unmasking schedule. The scheduling strategy designed by the researchers prioritizes revealing tokens that make other positions easier to predict, thereby reducing overall uncertainty early in decoding. This is equivalent to computing a smarter path offline.

Then this path is distilled into the model itself. An offline schedule alone is not enough—it needs to align with the model's actual parallel sampling behavior to truly accelerate. dSpark uses distillation to teach the model faster parallel sampling and introduces a KL divergence regularization term to constrain generation quality, avoiding sacrificing output for speed.

The ingenuity of this design lies in turning "decoding order" from a heuristic rule at inference time into a capability learned inside the model.

The meaning of schedule distillation

Distillation usually refers to using the output distribution of a large model (teacher) to guide the training of a small model (student). Schedule distillation replaces the "teacher" with "a precomputed decoding trajectory": first construct offline a high-quality unmasking order (i.e., which positions to reveal at each step), then have the model imitate the parallel sampling behavior corresponding to this trajectory during training. The motivation is that although the offline schedule is theoretically better, it relies on the confidence of a fixed model at a specific moment, whereas the model's distribution during actual parallel sampling changes with the tokens already revealed, creating a mismatch. Distillation essentially translates the "intent" of the offline schedule into a strategy executable inside the model, so that at inference time there is no need to run a complex schedule search—the model's own predictions can closely approximate that good path.

The role of the KL divergence regularization term in generation quality control

KL divergence measures the difference between two probability distributions. In dSpark's distillation objective, in addition to making the model match the behavior of the scheduling trajectory, a KL regularization term is added to constrain the model's prediction distribution at the current step from deviating too far from a reference distribution (usually the output of the original model or teacher distribution at the corresponding position). The purpose of this term is to prevent the model from overconfidently guessing in pursuit of "fewer steps": once the distribution at some positions is sharpened too much, incorrect tokens will be fixed too early and cannot be rolled back. KL regularization acts as a soft constraint of "don't stray too far" on the acceleration process, balancing speed and generation quality. This is also an important reason why the paper can maintain performance on structure-sensitive tasks such as GSM8K.

Results: speedup on mathematical reasoning and code generation

The paper evaluates on two types of tasks: mathematical reasoning and code generation, both of which have high requirements for generation quality and structure.

On GSM8K, compared with various baseline methods, dSpark achieves a 10.12× speedup on LLaDA-8B-Instruct and a 5.12× speedup on Dream-7B-Instruct, while maintaining or even improving task performance.

It should be noted that these two speedup numbers correspond to different base models and baseline combinations, and should not be simply understood as "a tenfold increase on the same model." But even conservatively, speedups of this magnitude are highly meaningful for the practical adoption of diffusion language models.

The code is open-sourced at https://github.com/VILA-Lab/dSpark.

A direction worth watching

For diffusion language models to truly compete with autoregressive models, inference efficiency is an unavoidable hurdle. dSpark's contribution is not inventing a new sampler, but changing the question: from "what to predict" to "when to commit."

This scheduling perspective may also inspire other parallel generation scenarios, such as draft scheduling in speculative decoding and ranking strategies for multi-token prediction. Of course, schedule distillation requires additional offline construction and distillation costs, and its generalization to larger models and more tasks remains to be further verified by the community.

One-line Summary

dSpark uses schedule distillation to optimize the unmasking order of diffusion language models, achieving several-fold inference speedups.

Paragraph Summary

When diffusion language models generate tokens in parallel, the number of decoding rounds is often too high, and the bottleneck lies in the order in which tokens are revealed first. VILA-Lab proposes dSpark, pointing out that the value of a token lies not only in whether it is easy to predict itself, but also in whether revealing it can make other positions predictably chain-react. The method has two steps: first, offline construct an unmasking schedule that prioritizes revealing tokens with high "information spillover," then distill this schedule to align with the model's parallel sampling behavior, with KL divergence regularization to constrain quality. On GSM8K, LLaDA-8B-Instruct achieves a 10.12× speedup and Dream-7B-Instruct achieves a 5.12× speedup, while task performance is maintained or even improved. This work transforms decoding order from a heuristic rule at inference time into a learned capability of the model.