Guided Star-Shaped Masked Diffusion

Guided Star-Shaped Masked Diffusion

In masked diffusion LLMs, there is a rather blunt problem: once a token is unmasked, it is frozen. If the model makes a mistake at an early step, it cannot correct that mistake and instead tries to build text around it. Such inconsistencies are particularly noticeable with a small number of steps. Existing remasking methods give the model a chance to reconsider, but they require hundreds of steps, which undermines the main advantage of diffusion — speed.

Authors from T-Tech, HSE, and Constructor University presented a method at ICML where all noisy states depend only on the clean text, not on each other. At each step, the model predicts the entire clean sequence, and the next state is resampled from it, inheriting nothing from the current one: any token can be returned to the mask and rewritten, and this does not require retraining. Generating from scratch is not possible — remasking at early steps destroys the unformed context, so the first ~70% of steps are written in the usual way, and only then the editing mode is activated. Where to edit is indicated by a separate lightweight error predictor trained on the model's mistakes during generation. The main downside of this approach is the inability to insert and delete tokens.

As a result, in 128 steps, the method surpasses the quality and diversity of text of the best previous approaches with a budget of 512 steps, allowing for nearly three times faster generation (3.4 vs. 9.2 seconds on H200). The error predictor trained on web texts transfers to code, mathematics, and news without retraining.