Guided Star-Shaped Masked Diffusion

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Meshchaninov, Viacheslav, Shibaev, Egor, Makoian, Artem, Klimov, Ivan, Balagansky, Nikita, Gavrilov, Daniil, Alanov, Aibek, Vetrov, Dmitry
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866915885069893632
author Meshchaninov, Viacheslav
Shibaev, Egor
Makoian, Artem
Klimov, Ivan
Balagansky, Nikita
Gavrilov, Daniil
Alanov, Aibek
Vetrov, Dmitry
author_facet Meshchaninov, Viacheslav
Shibaev, Egor
Makoian, Artem
Klimov, Ivan
Balagansky, Nikita
Gavrilov, Daniil
Alanov, Aibek
Vetrov, Dmitry
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2510_08369
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Guided Star-Shaped Masked Diffusion
Meshchaninov, Viacheslav
Shibaev, Egor
Makoian, Artem
Klimov, Ivan
Balagansky, Nikita
Gavrilov, Daniil
Alanov, Aibek
Vetrov, Dmitry
Machine Learning
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.
title Guided Star-Shaped Masked Diffusion
topic Machine Learning
url https://arxiv.org/abs/2510.08369