改进掩码扩散模型训练效率，让训练更贴近实际使用场景。
Stop Training for the Worst: Progressive Unmasking Accelerates Masked Diffusion Training
- 逐步解封掩码机制，使训练与推理时的掩码模式对齐。
- 在1.25亿参数规模上训练速度提升约2.5倍。
- 适合关注生成模型训练加速与高效推理的研究者。
掩码扩散模型（MDMs）在离散空间生成建模中展现出潜力，可通过任意顺序生成序列并实现并行解码，从而加快推理速度并在非因果任务上表现优异。然而，这种灵活性带来训练复杂度的权衡：MDMs需在指数级增长的掩码模式集合上训练，不仅计算成本高，还导致训练与推理时掩码模式不一致的问题。本文提出渐进式解封（PUMA），通过简化前向掩码过程，使训练阶段的掩码模式更贴近推理时的实际结构，从而将优化集中在与推理对齐的掩码上，显著提升训练效率。实验表明，在1.25亿参数规模下，PUMA可将预训练速度提升约2.5倍，并在常见训练策略（如自回归初始化）基础上提供额外优势。代码已开源。
原文摘要 · Abstract (English)
Masked Diffusion Models (MDMs) have emerged as a promising approach for generative modeling in discrete spaces. By generating sequences in any order and allowing for parallel decoding, they enable fast inference and strong performance on non-causal tasks. However, this flexibility comes with a training complexity trade-off: MDMs train on an exponentially large set of masking patterns, which is not only computationally expensive, but also creates a train--test mismatch between the random masks used in training and the highly structured masks induced by inference-time unmasking. In this work, we propose Progressive UnMAsking (PUMA), a simple modification of the forward masking process that aligns training-time and inference-time masking patterns, thereby focusing optimization on inference-aligned masks and speeding up training. Empirically, PUMA speeds up pretraining at the 125M scale by $\approx 2.5\times$ and offers complementary advantages on top of common recipes like autoregressive initialization. We open-source our codebase at https://github.com/JaeyeonKim01/PUMA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。