Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Chao, Chen-Hao, Sun, Wei-Fang, Liang, Hanwen, Lee, Chun-Yi, Krishnan, Rahul G.
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866912664616173568
author Chao, Chen-Hao
Sun, Wei-Fang
Liang, Hanwen
Lee, Chun-Yi
Krishnan, Rahul G.
author_facet Chao, Chen-Hao
Sun, Wei-Fang
Liang, Hanwen
Lee, Chun-Yi
Krishnan, Rahul G.
contents Masked diffusion models (MDM) are powerful generative models for discrete data that generate samples by progressively unmasking tokens in a sequence. Each token can take one of two states: masked or unmasked. We observe that token sequences often remain unchanged between consecutive sampling steps; consequently, the model repeatedly processes identical inputs, leading to redundant computation. To address this inefficiency, we propose the Partial masking scheme (Prime), which augments MDM by allowing tokens to take intermediate states interpolated between the masked and unmasked states. This design enables the model to make predictions based on partially observed token information, and facilitates a fine-grained denoising process. We derive a variational training objective and introduce a simple architectural design to accommodate intermediate-state inputs. Our method demonstrates superior performance across a diverse set of generative modeling tasks. On text data, it achieves a perplexity of 15.36 on OpenWebText, outperforming previous MDM (21.52), autoregressive models (17.54), and their hybrid variants (17.58), without relying on an autoregressive formulation. On image data, it attains competitive FID scores of 3.26 on CIFAR-10 and 6.98 on ImageNet-32, comparable to leading continuous generative models.
format Preprint
id arxiv_https___arxiv_org_abs_2505_18495
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking
Chao, Chen-Hao
Sun, Wei-Fang
Liang, Hanwen
Lee, Chun-Yi
Krishnan, Rahul G.
Machine Learning
Masked diffusion models (MDM) are powerful generative models for discrete data that generate samples by progressively unmasking tokens in a sequence. Each token can take one of two states: masked or unmasked. We observe that token sequences often remain unchanged between consecutive sampling steps; consequently, the model repeatedly processes identical inputs, leading to redundant computation. To address this inefficiency, we propose the Partial masking scheme (Prime), which augments MDM by allowing tokens to take intermediate states interpolated between the masked and unmasked states. This design enables the model to make predictions based on partially observed token information, and facilitates a fine-grained denoising process. We derive a variational training objective and introduce a simple architectural design to accommodate intermediate-state inputs. Our method demonstrates superior performance across a diverse set of generative modeling tasks. On text data, it achieves a perplexity of 15.36 on OpenWebText, outperforming previous MDM (21.52), autoregressive models (17.54), and their hybrid variants (17.58), without relying on an autoregressive formulation. On image data, it attains competitive FID scores of 3.26 on CIFAR-10 and 6.98 on ImageNet-32, comparable to leading continuous generative models.
title Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking
topic Machine Learning
url https://arxiv.org/abs/2505.18495