Guided Star-Shaped Mask Diffusion

Masked diffusion cannot take anything back

A masked diffusion language model generates by running corruption in reverse. The forward process replaces tokens with a special $m = \texttt{[MASK]}$ symbol according to a schedule $\alpha_t$, giving the marginal $q(x_t \mid x_0) = \mathrm{Cat}(x_t;\, \alpha_t x_0 + (1-\alpha_t)\, m)$. A denoiser $f_\theta$ is trained to look at a partially masked sequence and predict the clean text; sampling walks from fully masked to fully revealed.

The trouble hides in the reverse transition that standard MDLM sampling uses. Conditioned on the current state and the model's prediction, it is

The first case is the villain of this story: any token that is already revealed is copied forward with probability one. A wrong commitment made at step three is still there at step three hundred. That rigidity is tolerable when tokens are revealed one at a time, but parallel generation, the whole reason to use diffusion for text, reveals many tokens per step and inevitably plants inconsistencies that the sampler is then forbidden to repair.

Existing escape hatches each stumble. Random remasking, as in ReMDM by Wang and colleagues, revisits correct tokens as often as wrong ones and burns steps on chance. Confidence-based selection is built on sand, because the training loss $\mathbf{1}(x_t^i = m)\log p_\theta^i$ touches only masked positions, so the model's probabilities on already revealed tokens are never calibrated. And refiners trained to spot uniform random corruptions learn an easier task than catching the plausible, in-context mistakes a diffusion model actually makes.

Point every state at the data, not at its neighbor

Our fix changes the graph, not the model. In the star-shaped formulation, borrowed in spirit from Okhotin and colleagues, the latent states are conditionally independent given the data: $q(x_{1:T} \mid x_0) = \prod_{t=1}^{T} q(x_t \mid x_0)$. The reverse posterior then collapses to $q(x_{t-1} \mid x_t, x_0) = q(x_{t-1} \mid x_0)$, and a sampling step becomes two moves: form a complete clean hypothesis $\hat{x}_0 \sim p_\theta(\cdot \mid x_t)$, then renoise that hypothesis, $x_{t-1} \sim q(\cdot \mid \hat{x}_0)$, with no dependence on $x_t$ at all. Nothing is copied forward by decree, so any token can return to $m$ and be rewritten.

The formulation costs nothing to adopt. We prove that the variational bound of the star-shaped process simplifies to the same weighted cross-entropy objective as standard masked diffusion, only with different timestep weights, so off-the-shelf MDLM checkpoints drop in with zero fine-tuning. The sampler even turns out to be mathematically equivalent to ReMDM with its remasking parameter pinned to $1-\alpha_s$, which quietly deletes the per-schedule search over that hyperparameter.

There is a price, and we report it plainly: because $x_{t-1}$ no longer conditions on $x_t$, consecutive states can drift apart. Run the star sampler from the very first step and the generation loses the thread of its own draft.

Generation has two phases, and remasking belongs to one

Early in generation most positions are masked, so each full hypothesis $\hat{x}_0$ invents many tokens against almost no context. The inventions are individually plausible but mutually incoherent, and the independent renoising step then randomly keeps some of them while masking the context they depended on. Errors compound, step-to-step agreement between consecutive hypotheses collapses, and pure star sampling degenerates into near-zero MAUVE text.

Late in generation the same move becomes a repair tool: with a dense, coherent draft as context, remasking a few tokens and repredicting them from the global picture fixes mistakes instead of creating them. Sweeping the switch point $t_{\mathrm{on}}$ confirms the split cleanly, with quality peaking when standard MDLM drafts for roughly the first 60 to 80 percent of steps and the star-shaped sampler refines the rest. Both pure extremes lose.

Teach a small head where the mistakes are

Even in its proper phase, unguided renoising is wasteful, since $q(x_{t-1} \mid \hat{x}_0)$ is as happy to mask a correct token as a broken one. So we add a second, much smaller model: an error predictor $g_\phi$ whose only job is to guess which tokens of $\hat{x}_0$ the denoiser got wrong.

The training loop matters more than the architecture. We corrupt a clean text to a random noise level, let the frozen denoiser predict and sample a candidate, and mark every position where the candidate disagrees with the truth. Those disagreements, the model's own mistakes, become binary labels for $g_\phi$. A classification head on the frozen backbone is enough, and the learned detector generalizes: trained on web text, it keeps an AUC-ROC of at least 0.90 on unseen code, mathematics, and news without any fine-tuning.

At inference each guided step scores the candidate, scales the error logits by a temperature $\tau_{\mathrm{remask}}$, and draws $N = \lceil (1-\alpha_{t-1}) L \rceil$ positions without replacement by Gumbel-Top-K, remasking exactly those. The star-shaped structure is what makes the training honest: a refiner conditioned on $x_t$ would only ever see ground-truth context during training and would collapse into copying it, while our decoupled transition never lets it treat observed tokens as sacred.

Fewer steps, better text

On 512-token OpenWebText generation we trace full quality-diversity Pareto fronts by sweeping the denoiser temperature, and the guided sampler dominates them: G-Star with a 4x smaller budget of 128 steps reaches better perplexity and diversity than every 512-step baseline that works with a frozen backbone, including per-schedule-tuned ReMDM, confidence-based P2 and RDM, and the planner DDPD. The wall clock agrees, 3.43 seconds per sample against 9.16 for the 512-step samplers, while DDPD runs a second full backbone every step and doubles its passes.

Guidance is cheap by construction. With $B$ predictor blocks on a $D$-block backbone and a guided fraction $\Delta$ of the trajectory, the total cost is $(1 + \Delta B/D)\,T$ forward passes; the single-block variant lands at 133 passes for a 128-step run, about four percent over the baseline. The full-depth head-only predictor is the robust choice, and since only a classification head is trained, it adds essentially no parameters worth storing.

Code makes the few-step story concrete. At 32 steps, guided refinement reaches a conditional perplexity of 20.4 where the MDLM baseline needs four times the compute to reach a worse 26.7, and the ranking repeats at every budget we test.

The same head helps a 7B instruction model

To test the recipe where it matters, we bolt G-Star onto Dream-Instruct 7B: the backbone stays frozen, a classification head trains on Tulu 3, and ten percent of the unchanged step budget is spent as refinement steps. Every one of the seven benchmarks we evaluate moves in our favor.

The honest caveats: intensive refinement eventually trades diversity for quality, so MAUVE peaks and then declines if the loop runs too long; the full-depth predictor adds roughly fifty percent latency even though its memory cost is a rounding error; and the star phase must stay in the late stages where context is dense. Within those bounds, targeted correction beats stochastic correction everywhere we measured it.