给预训练扩散模型加个轻量插件，让生成时能自动纠错。
Fine-Tuning Masked Diffusion for Provable Self-Correction
- 用插件式重掩码法，在推理时计算每个词元质量得分。
- 在数独、文本和代码任务中，生成质量显著提升。
- 无需强化学习或验证器，适用于任意预训练掩码扩散模型。
生成模型的一个理想特性是自纠正——在推理时检测并修正低质量词元。尽管掩码扩散模型（MDMs）已成为离散空间生成建模的有前景方法，其自纠正能力仍不明确。此前尝试将自纠正引入MDMs要么需要重构架构或训练流程，要么依赖粗糙的词元质量代理，限制了实用性。为此，我们提出PRISM——一种轻量级、模型无关的推理时自纠正插件，适用于任何预训练的MDM。理论上，PRISM定义了一个可证明学习每个词元质量分数的自纠正损失，无需强化学习或验证器。这些质量分数与MDM前向传播同时计算，用于识别低质量词元。实验上，PRISM在多个领域和规模下均提升了MDM的生成性能：数独；无条件文本（170M参数）；以及使用LLaDA（8B参数）的代码生成。
原文摘要 · Abstract (English)
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).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。