Guided Star-Shaped Masked Diffusion
This addresses sampling efficiency and quality issues in masked diffusion models for generative tasks like text and code, representing an incremental improvement over existing methods.
The paper tackles the irreversible decision-making and low-step generation limitations of pre-trained masked diffusion models by introducing a novel star-shaped sampling algorithm with learnable re-masking scheduler, which after lightweight fine-tuning yields substantial quality improvements, particularly with few sampling steps, outperforming or matching existing methods in text and code generation experiments.
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 re-masking scheduler 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.