Guided Star-Shaped Masked Diffusion

Viacheslav Meshchaninov, Egor Shibaev, Artem Makoian, Ivan Klimov, Nikita Balagansky, Daniil Gavrilov, Aibek Alanov, Dmitry Vetrov

Keywords: Discrete Diffusion, Text Diffusion Models, Masked Diffusion, Guided Sampling

TL;DR: We introduce a lightweight, plug-and-play remasking module that enables efficient error correction for frozen masked diffusion models, significantly boosting generation quality in low-step regimes.

Abstract: 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 comprehensive experiments on text, and code generation, our sampling algorithm outperforms or matches existing methods.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 46