Guided Star-Shaped Masked Diffusion

The performance of pre-trained masked diffusion models is often constrained by their sampling procedure, which makes decisions irreversible and struggles in low-step generation regimes. We introduce a novel sampling algorithm that works with pre-trained models and, after a lightweight fine-tuning of a single layer, significantly improves sample quality and efficiency. Our method reformulates the generation process using a star-shaped paradigm, which inherently allows for error correction. To make this process effective, we augment it with a learnable remasking module that intelligently identifies and revises likely errors. This approach yields a substantial quality boost, particularly when using a small number of sampling steps. We extensively ablate key components of our approach and show its usability in different scenarios. In experiments on text, and code generation, our sampling algorithm outperforms or matches existing methods. Code is available at https://github.com/EgorShibaev/G-Star.

Masked diffusion models generate text by starting with masked tokens and gradually replacing them with words. This allows them to generate many tokens in parallel, but it also creates an important limitation: once a token is filled in, standard masked diffusion usually cannot revise it. As a result, early mistakes may remain in the final output, especially when generation must be done in only a small number of steps. We introduce G-Star, a new sampling method for masked diffusion models that adds a targeted correction mechanism. Instead of randomly choosing which tokens to revisit, G-Star uses a lightweight error predictor to estimate which generated tokens are most likely to be wrong. These likely errors are masked again and regenerated, while more reliable parts of the text are preserved. G-Star works with already trained masked diffusion models and only requires lightweight additional training. In experiments on text, instruction-following, and code generation, it improves sample quality and efficiency, especially in few-step generation settings. This makes masked diffusion a more practical approach for fast and parallel language generation.