Abstract:A natural desideratum for generative models is self-correction--detecting and revising low-quality tokens at inference. While Masked Diffusion Models (MDMs) have emerged as a promising approach for generative modeling in discrete spaces, their capacity for self-correction remains poorly understood. Prior attempts to incorporate self-correction into MDMs either require overhauling MDM architectures/training or rely on imprecise proxies for token quality, limiting their applicability. Motivated by this, we introduce PRISM--Plug-in Remasking for Inference-time Self-correction of Masked Diffusions--a lightweight, model-agnostic approach that applies to any pretrained MDM. Theoretically, PRISM defines a self-correction loss that provably learns per-token quality scores, without RL or a verifier. These quality scores are computed in the same forward pass with MDM and used to detect low-quality tokens. Empirically, PRISM advances MDM inference across domains and scales: Sudoku; unconditional text (170M); and code with LLaDA (8B).
| Comments: | Authorship statement: Jaeyeon Kim and Seunggeun Kim contributed equally, and Taekyun Lee is also a co first author |
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2510.01384 [cs.LG] |
| (or arXiv:2510.01384v4 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2510.01384 arXiv-issued DOI via DataCite |
From: Jaeyeon Kim [view email]
[v1]
Wed, 1 Oct 2025 19:15:25 UTC (757 KB)
[v2]
Fri, 7 Nov 2025 04:01:45 UTC (754 KB)
[v3]
Fri, 19 Dec 2025 06:19:54 UTC (1,149 KB)
[v4]
Fri, 22 May 2026 19:48:31 UTC (3,850 KB)
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。