Soft-Di[M]O：利用软嵌入改进单步离散图像生成
计算机视觉与模式识别
2026-03-20 v2 人工智能
机器学习
摘要
从掩码扩散模型（MDM）蒸馏出的单步生成器将多个采样步骤压缩为一次前向传播，从而实现高效的文本和图像合成。然而，它们存在两个关键局限：它们继承了教师模型的建模偏差，并且其离散的 token 输出阻断了梯度流，使得无法进行对抗训练、基于奖励的微调和测试时嵌入优化（TTEO）等蒸馏后优化。在这项工作中，我们引入了软嵌入，这是一种简单的松弛方法，用生成器输出分布下的期望嵌入来替代离散 token。软嵌入为单步离散生成器保留了表示保真度，同时提供了一个完全可微的连续替代方案，该方案与教师模型骨干网络和分词器解码器兼容。将软嵌入集成到 Di[M]O 蒸馏框架中（记为 Soft-Di[M]O）使得单步生成器能够进行端到端训练，并能够直接应用基于 GAN 的优化、可微奖励微调和 TTEO。实验表明，在多个 MDM 教师模型（例如 MaskBit、MaskGen）上，Soft-Di[M]O 取得了最先进的单步结果：改进了类别到图像的性能，在 ImageNet-256 上结合基于 GAN 的优化实现了 1.56 的单步 FID，通过奖励微调在文本到图像上获得了更高的 GenEval 和 HPS 分数，并通过 TTEO 获得了进一步的提升。
引用
@article{arxiv.2509.22925,
title = {Soft-Di[M]O: Improving One-Step Discrete Image Generation with Soft Embeddings},
author = {Yuanzhi Zhu and Xi Wang and Stéphane Lathuilière and Vicky Kalogeiton},
journal= {arXiv preprint arXiv:2509.22925},
year = {2026}
}
备注
ICLR 2026