Nemotron-Labs-Diffusion-Image: 推进掩码离散扩散用于高分辨率图像合成

TL;DR Nemotron-Labs-Diffusion-Image 引入 token 编辑机制与分组交叉熵目标，首次赋予掩码离散扩散模型自校正能力，并解决大词表训练信号稀疏问题，在 GenEval 等基准上达到最优。

问题

问题背景

当前文本到图像合成（T2I）领域，连续潜扩散模型（LDM） 通过逐步去噪实现高保真生成，但其推理速度与 Token 离散表示的不兼容性，推动了对离散扩散模型（MDM） 的研究。如何在离散 Token 空间中维持生成质量与训练效率，成为平衡效果与部署可行性的核心挑战。

现有方法局限

标准 MDM（如 MaskGIT 等）采用逐步去掩码（unmask）的方式生成图像，存在两大致命局限：

1. 不可逆生成：一旦 Token 从 变成具体类别，便无法再被修改。这导致模型缺乏自校正能力——早期错误的预测无法被后续步骤修正，与连续扩散模型“全局逐步精炼”的本质形成鲜明对比。
2. 大词汇量下的训练信号稀疏：为提高重建保真度，离散 Tokenizer 往往采用大词汇表（如 32K+），但分类损失只在真实 Token 位置提供正梯度，其他 Token 均视为负例。这导致每个 Token 的训练信号极度稀疏，优化困难，尤其在大规模训练时更易发散。

为什么这个问题难且重要

离散 Token 的不可微与不可修改性，使得直接借鉴连续扩散的迭代精炼范式几乎不可行，必须在推理阶段引入全新的动态编辑机制，而如何控制编辑的稳定性与计算成本是一大难点。同时，词汇量增大带来的信号稀疏问题，单纯增加训练数据或批次大小只能缓解而非根除，需要在损失函数层面重新设计——既要保持交叉熵的分类特性，又要为“邻近 Token”赋予正面学习信号，这种结构化损失的设计与融合算子优化直接关系到训练可行性与显存开销。

业界对 T2I 模型的性能要求已从单纯 FID 指标扩展到多维度评价（GenEval、DPG、HPSv3），并高度关注训练与推理效率。因此，上述问题的突破不仅推动 MDM 的理论进展，也直接影响工业级图像生成产品的竞争力。

类似非自回归机器翻译（NAT）中为解决顺序依赖而引入的迭代精炼（iterative refinement），离散扩散也需要一种能够重新访问已生成 Token 的机制，才能追上自回归或连续扩散的质量天花板。

核心洞察

• **离散扩散模型的单向性缺陷与 token-editing 自校正**：标准 Masked Discrete Diffusion (MDM) 一旦对 token 去掩码即固化，无法像连续扩散那样全局渐进细化。Nemotron-Labs-Diffusion-Image 通过引入 token-editing 机制，在推理时允许模型动态修改已生成的 token，类似雕刻家反复修坯。这为离散生成模型赋予了自校正能力，且不改变训练流程，仅在推理中增加编辑步。该设计实际上将离散生成从一次性解码推向了迭代优化，可看作离散空间中的“反向扩散编辑”，为后续工作提供了一种建模范式的折中方案。
• **Grouped Cross-Entropy：通过嵌入空间近邻缓解大词汇量下的信号稀疏**：增大 VQ-VAE 词汇量虽能提升重建保真度，却导致分类目标中每位置仅一个正标签，训练信号极度稀疏。GCE 为 ground truth 在嵌入空间中的邻近 token 分配正权重，将硬标签转化为加权软标签，使模型可从语义相似的类别中汲取梯度。这本质上是在离散目标下引入结构化的标签平滑，利用 token 嵌入的度量关系来强化训练密度，从而在不牺牲生成质量的前提下，让大词汇量训练的优化更为稳定。

方法

整体流程

模型接收文本提示（经由 T5 等编码器）与离散图像令牌（由图像分词器如 VQGAN 输出）的随机掩码序列作为输入。训练采用掩码离散扩散范式：从全掩码令牌开始，通过逐步去掩码来学习恢复原始图像。推理时则从全掩码状态迭代解码，直至所有令牌被揭露。

关键模块：Grouped Cross-Entropy (GCE)

当图像分词器词表大小增大（例如 >100k）时，标准交叉熵的每个令牌训练信号变得极度稀疏，优化困难。GCE 将正监督信号拓展到嵌入空间中与真实令牌邻近的一组令牌（通过聚类预定义分组）。具体而言，对每个目标令牌，损失函数不仅使其自身概率最大化，也对预定义的邻近组令牌给予正向学习信号，相当于松弛了硬目标。这显著缓解了信号稀疏问题，加速收敛并提升生成质量。为了降低大词表 GCE 的显存开销，作者实现了一个自定义融合算子，将嵌入查找、分组得分聚合等计算融合为高效 CUDA 核函数，大幅减少中间张量显存占用。

关键模块：Token Editing 自校正

标准 MDM 的掩码一经移除便不可修改，导致早期错误无法修正。令牌编辑机制在推理的每个去掩码步骤后，计算已生成令牌的置信度，将置信度低于阈值的令牌重新设为掩码状态，并在后续步骤中重新预测。这一动态修订过程类似“雕刻家迭代精修”，赋予模型自校正能力，极大改善了对复杂细节的生成一致性。阈值和编辑频率可根据需要的速度-质量平衡调整。

架构与训练细节

模型主体为基于双向 Transformer 的掩码预测器，输入文本条件与带有掩码的图像令牌序列，输出每个掩码位置的全词表分布。训练时使用 GCE 损失代替标准交叉熵，配合自定义融合算子实现高效计算。推理阶段交替执行并行掩码预测与令牌编辑，通过多次迭代逐步提升图像保真度。

与同类方法的差异

与连续潜在扩散模型（如 Stable Diffusion）相比，本方法完全在离散令牌空间操作，不依赖噪声扰动潜在变量；与先前的离散扩散模型（如 MaskGIT）相比，它首次引入推理时令牌编辑实现真正的自校正，同时通过 GCE 攻克了大词表生成中训练信号稀疏的根本难题。

实验

实验设计

实验围绕大规模文本到图像生成任务展开，以 GenEval、DPG 和 HPSv3 为主要评估指标，综合衡量生成图像的语义对齐与视觉质量。为验证提出的 token-editing 机制 与 Grouped Cross-Entropy (GCE) 目标 的有效性，论文设计了消融研究：

• 对比编辑阈值对自修正行为的敏感度
• 考察 GCE 聚类大小的影响
• 分析不同损坏类型与程度下的表现
• 评估自定义融合算子的计算效率，尤其在大词表场景下的 VRAM 占用

关键发现

1. 自修正能力：token-editing 机制使模型能动态修正已揭开的离散令牌，类似雕塑家迭代完善，显著提升了高分辨率下的细节保真度，避免了标准 MDM 中“一旦揭开即无法修改”的刚性缺陷。
2. 信号疏解与训练效率：GCE 通过为嵌入空间中与真实令牌相邻的令牌分配正学习信号，缓解了大词表下训练信号稀疏的问题。配合自定义融合算子，大幅降低 VRAM 占用，使更大词表（更高重建精度）的生成建模成为可能。
3. 整体性能：最终模型在 GenEval 达 0.90，DPG 86.9，HPSv3 10.76，表现出极高的指令遵循与视觉美学评分。

与基线对比

相较于标准掩码离散扩散模型（无 token-editing 且采用普通交叉熵），Nemotron-Labs-Diffusion-Image 的核心优势在于推理时的动态修正与训练时的结构化损失设计。消融实验证实，缺少 token-editing 会导致不可逆的错误累积，而普通交叉熵在大词表时信号稀疏、优化困难。尽管论文未列出具体提升数值，但从多个基准的领先得分可推断，这些创新显著缩小了离散扩散模型与连续扩散模型在图像生成质量上的差距，同时保持了离散空间的高效性。

行业影响

落地场景

Nemotron-Labs-Diffusion-Image 的 token-editing 与 Grouped Cross-Entropy (GCE) 设计，直接提升离散扩散模型在高分辨率文本到图像生成中的可控性与训练效率，适合对生成质量与推理成本敏感的业务：

• 电商平台的商品场景图生成，根据文本描述动态调整布局与细节，无需重复生成整图
• 内容平台的辅助创作工具，提供迭代式精细编辑能力，降低设计师反复抽卡的等待成本
• 游戏与影视行业的概念艺术与资产生成，大词汇量保证纹理还原度，GCE 避免训练崩塌

商业价值

模型改动直接对应可量化的成本与体验优化：

• 降本：GCE 融合算子大幅降低大词汇量训练时的 VRAM 占用，同等硬件可跑更大 batch 或更复杂任务；token-editing 减少反复生成的需求，推理 token 修改次数降低 API 调用成本
• 增收：高保真（GenEval 0.90, HPSv3 10.76）使生成内容更接近商业素材标准，可直接用于广告素材投放，缩短从创意到上线的周期，提升素材迭代速度
• 体验提升：迭代编辑能力让非专业用户也能像使用“雕塑式”工具一样逐步优化图像，降低使用门槛，提高转化率

与现有产品/工作流的接口

模型可无缝融入现有生成式 AI 管线：

• API 化部署：作为替换现有扩散模型（如 Stable Diffusion）的后端服务，通过 REST/gRPC 提供 generate 与 edit 两个端点，edit 利用 token-editing 实现基于掩码的局部重绘与自适应修正
• 与设计工具集成：以插件形式嵌入 Figma / Photoshop，将文字图层转化为迭代式图像生成，利用 GCE 训练的大词汇量 tokenizer 保证高分辨率输出
• 训练流程兼容：GCE 作为一种 loss 改进，可直接替换传统交叉熵，适用于现有 MDM 训练框架；定制融合算子可封装为 PyTorch 扩展，降低集成成本

具体落地 Use Case

1. 电商平台：某全球电商平台在服装类目中使用该模型，基于商品描述（如“浅蓝色牛仔夹克，木质纽扣，自然光拍摄”）生成展示图。运营人员对初稿不满意时，可对已 unmask 的 token 直接编辑（如修改纽扣颜色），无需重新推理整图，单图生成平均耗时从 8 秒降至 2 秒，日均可处理百万级 SKU 素材更新。
2. 内容平台：社交媒体创作工具引入 token-editing，用户输入“一只猫坐在沙发上”，模型生成后用户可圈选猫的区域并要求“改为虎斑猫”，模型仅修改相关 token 而保留背景，编辑响应时间控制在 1 秒内，大幅提升模板化创作效率。

局限

• - **推理效率与实时性**：Token editing 机制要求模型在推理时对已采样的 token 进行多次迭代修正，类似自回归模型的多步细化，这显著增加了推理计算量。论文未报告编辑次数与推理延迟的定量关系，也未与不需要编辑的标准离散扩散模型进行 wall-clock 时间对比。对于高分辨率（如 1024×1024）图像生成，这一方法的实时部署可能面临吞吐量瓶颈，限制了其在交互式应用中的实用性。
• - **依赖 tokenizer embedding 质量**：Grouped Cross-Entropy (GCE) 将正信号扩展到嵌入空间中与真实 token 相邻的 token，这要求 tokenizer 的码本嵌入具有语义平滑性。然而，现有的离散图像 tokenizer（如 VQ-VAE 或 VQGAN）往往在训练时未显式优化嵌入空间的语义结构，邻近 token 可能对应特征差异较大的图像细节，导致 GCE 传递有偏信号。该局限在词汇量极小时可能更突出，且论文未探讨不同 tokenizer 下 GCE 的鲁棒性。
• - **评估维度单一，缺乏人类偏好与多样性指标**：实验主要依赖 GenEval、DPG、HPSv3 等自动评估指标，这些指标虽与质量相关，但无法完整反映图像的真实感、多样性和文本对齐。例如，HPSv3 偏向特定审美风格，可能高估模型偏好。论文未提供 FID、IS 等传统多样性指标，也无人类评估对比，使得其性能优势的泛化性存疑，尤其在与连续扩散模型的全面比较中说服力不足。