Nemotron-Labs-Diffusion-Image：推进用于高分辨率图像合成的掩码离散扩散模型 Nemotron-Labs-Diffusion-Image: Advancing Masked Discrete Diffusion for High-Resolution Image Synthesis
提出带Token编辑和GCE损失的8B离散扩散模型，GenEval得分0.90
前置知识
Masked Discrete Diffusion Models (MDM)
MDM是一种生成模型，它将图像离散化为token序列，然后通过逐步unmask的方式生成图像。训练时，从完全masked的序列开始，每一步预测一些位置上的原始token。与自回归模型不同，MDM可以并行生成多个位置的token，推理速度更快。标准的MDM训练使用ELBO目标函数，只在masked位置计算loss。MDM的推理过程从完全masked的序列开始，每一步预测一些位置并unmask它们，直到所有位置都被预测完成。这种掩码建模方法的一个重要优势是可以处理固定长度的序列，并且可以原生支持图像inpainting等任务。
理解MDM是阅读本文的基础，因为本文的所有创新都是建立在标准MDM之上的。只有理解了标准MDM的工作原理（如masking策略、ELBO目标函数、推理过程），才能理解本文提出的token editing机制是如何扩展标准MDM的。特别是MDM缺乏自校正能力的根本原因——一旦unmask就固定——这是本文要解决的核心问题之一。如果读者不理解MDM的训练只在masked位置计算loss，就无法理解为什么引入corruption可以 enable token editing。此外，MDM与连续扩散模型的区别也是理解本文贡献的关键，因为本文的目标是让离散模型获得连续模型的自校正能力。
Vector Quantization (VQ) Tokenizer
VQ Tokenizer通过向量量化将连续的图像特征离散化为离散的token code。它包含一个codebook，每个code对应一个离散的token。较大的codebook可以减少量化误差，提高重建保真度，但也会导致训练信号稀疏问题——每个token出现的频率降低，难以获得充分的监督信号。VQ Tokenizer的工作流程是：首先使用encoder将图像编码为连续特征，然后将每个特征向量与codebook中的向量比较，选择最接近的code作为离散token。解码时，使用对应的codebook向量重建连续特征，再通过decoder重建图像。本文使用的tokenizer包含131K个codes，这是一个非常大的词表。
VQ Tokenizer是离散图像生成的基础，理解它的工作原理和挑战是理解本文GCE贡献的前提。特别是codebook sparsity问题——当词表增大时，每个token的频率降低，训练信号变得稀疏——这是本文GCE目标函数要解决的核心问题。只有理解了为什么大词表会导致训练困难，才能理解GCE的设计思路。此外，VQ Tokenizer的embedding空间特性（语义相似的codes在embedding空间中接近）也是GCE能够work的基础，因为GCE通过K-means聚类利用了这种语义相似性。如果不理解VQ Tokenizer的工作原理，就无法理解为什么语义相似的token应该获得正向监督信号。
Codebook Sparsity Problem
当使用大词表（如100K+）时，在固定的图像数据集上，每个token的出现频率会急剧下降。例如，100万张图像每张256个token，8192大小的词表平均每个token出现31250次，而200K大小的词表只有1280次——仅为前者的4%。这导致训练信号稀疏，模型难以充分学习每个token的特征。更严重的是，标准cross-entropy损失为所有非ground-truth token分配相同的负梯度，忽略了语义相似性。两个视觉上相似的特征在small codebook中可能映射到同一个code，但在large codebook中可能映射到不同codes，而标准损失无法捕捉这种语义相似关系，导致优化困难。
Codebook Sparsity Problem是本文要解决的核心问题之一，理解它的重要性是理解GCE贡献的关键。这个问题直接影响大词表生成模型的训练效果和稳定性。本文的GCE目标函数专门设计来缓解这个问题，通过层级聚类为语义相似的token提供正向监督信号。只有理解了sparsity问题如何影响训练（频率降低、梯度稀疏、优化困难），才能理解GCE的设计动机和效果。此外，sparsity问题也解释了为什么现有方法主要依赖暴力扩展（更大模型、更多数据）而不是算法改进，而本文提出了一个更优雅的算法解决方案。
研究动机
现有的掩码离散扩散模型存在两个根本性问题。第一个问题是缺乏自校正能力。与连续扩散模型（如LDM）通过迭代精炼latent embeddings来逐步改进输出不同，标准的MDM一旦unmask一个token就永久固定它，即使预测是错误的也无法后续修正。这种commitment问题导致误差累积，因为unmask的每个token都是从位置独立的logits中独立采样的，隐含假设了token之间的独立性，但实际图像token之间存在强依赖关系。第二个问题是训练大词表生成模型的困难。虽然更大的词表（如100K+）可以提高重建保真度，但它加剧了codebook sparsity问题，使得训练信号变得极其稀疏，优化困难。实验数据显示，从8192词表扩展到200K词表，平均每个token的频率从31250次降到1280次，仅为前者的4%。
本文的目标是本文的目标是设计一个能够生成高分辨率文本到图像的离散扩散模型，解决上述两个核心问题。具体来说，目标包括：第一，为MDM引入自校正机制，使其像连续扩散模型一样能够在推理过程中迭代修正已生成的token，从而减少误差累积，提高图像保真度。第二，设计有效的训练目标函数来缓解大词表带来的稀疏性问题，使模型能够高效地学习大词表，而不是依赖暴力扩展模型大小和数据量。第三，在保持高效推理的同时，在多个基准测试上达到最先进的性能，包括GenEval、DPG和HPSv3等指标。最终目标是推动离散图像生成向更大规模和更高保真度发展。
与已有工作不同的是，本文的独特切入角度是同时解决MDM的两个根本性限制。现有工作要么关注token editing（主要在语言领域，如LLaDa-2.1、EditFlow），要么关注大词表训练（如SNCE），但很少有工作同时解决这两个问题。更重要的是，现有的token editing方法依赖插入和删除操作，这对固定长度的图像token序列不适用，因为图像token数量由分辨率固定。本文是首个将token editing机制成功应用于离散图像生成的工作，并且提出了GCE目标函数作为SNCE的更有效替代方案，既减少了内存开销，又提高了优化稳定性。这种双重创新使得本文能够在不牺牲推理效率的前提下，显著提升图像生成质量。
核心方法
Nemotron-Labs-Diffusion-Image采用decoder-only transformer架构，同时处理文本prompt和图像token。整体思路是在标准MDM的基础上做两个关键扩展：一是引入token editing机制，使模型能够在推理过程中修正已unmask的token；二是设计Grouped Cross-Entropy（GCE）目标函数，通过层级聚类为语义相似的token提供正向监督信号，缓解大词表稀疏性问题。模型从8B参数的预训练diffusion语言模型初始化，使用131K大小的codebook，采用渐进式分辨率提升策略（从256乘以256逐步到1024乘以1024），在64块H100 GPU上训练300K步。这种设计使模型能够在保持高效推理的同时，达到最先进的生成质量。
核心创新点有两个：第一是Token Editing机制，通过在训练时对clean token进行corruption，而不仅是masking，使模型学习在所有位置预测token分布。推理时，如果模型对已unmask位置预测的不同token置信度超过阈值tau，就替换原有token，实现迭代精炼。corruption策略包括从同一图像中随机采样token，以及从tokenizer embedding空间中的最近邻采样。第二是GCE目标函数，使用K-means将131K code聚类成多个层级（如16384和8192），在标准token-level cross-entropy基础上，增加cluster-level的辅助损失。梯度分析显示，ground-truth token从所有项获得正向梯度，语义相似的token从部分层级获得正向梯度，而无关token获得负向梯度，从而有效缓解稀疏性问题。
方法步骤详情
方法步骤包括：第一，架构初始化，从8B预训练diffusion语言模型初始化，将语言建模head替换为预测图像token的MLP。使用decoder-only架构而非encoder-decoder，可以处理更长的prompt并利用sequence packing优化。第二，Token Editing训练，给定clean序列，采样timestep t，使用增强的前向过程生成。标准MDM的前向过程随机将clean token替换为[M] token，而本文的前向过程对未mask的clean token进行额外corruption：要么从同一图像随机采样，要么从embedding空间最近邻采样。第三，GCE损失计算，对每个位置，计算token-level的cross-entropy和cluster-level的cross-entropy。cluster概率通过将同一cluster内token的概率相加得到。第四，推理，每步进行标准unmasking，同时对已unmask位置，如果新预测token的置信度超过阈值tau，则替换原token。第五，优化实现，设计custom fused operator加速GCE计算，相比naive PyTorch实现，延迟从44ms降到20ms，最大VRAM从25GB降到16GB。
技术新颖性
技术新颖性体现在三个方面：第一，首次将token editing机制成功应用于离散图像生成。与语言领域的插入或删除操作不同，本文采用同图像随机采样和embedding空间最近邻采样两种corruption策略，更适合固定长度的图像token序列。第二，GCE目标函数相比SNCE有三个优势：动态调整置信度而非固定软标签；只存储离散cluster索引，减少内存开销；优化目标与预测一致性，完美预测时损失为零。第三，高效实现，custom fused operator在不引入额外内存开销的情况下，只增加8ms延迟，相比标准cross-entropy的12.71ms，额外成本可忽略不计。这使得在大规模训练中实用化大词表GCE成为可能，避免了纯理论设计的局限性。
实验结果
实验结果表明Nemotron-Labs-Diffusion-Image在多个基准上取得了最先进的性能。GenEval基准上达到0.90的overall得分，显著超过同为MDM的Meissonic（0.54）和统一多模态模型LaViDa-O（0.77），与20B参数的Qwen-Image-2507（0.87）竞争。DPG基准上得分86.9，超过GPT-4o的85.3。MJHQ-30k数据集上，FID达到6.46（未微调）和12.23（微调后），HPSv3达到9.61（未微调）和10.76（微调后）。消融实验显示token editing在所有NFE设置下一致提升图像质量，32步with editing的质量相当于64步without editing，有效实现2倍前向调用减少。GCE相比SNCE在ImageNet256上FID从3.62降到3.40（100 epoch）和从3.42降到3.00（300 epoch）。推理速度上，相比Emu3.5快42.4倍同时GenEval分数更高。Few-step生成实验显示，即使只4步也能生成合理质量的图像，而连续flow-matching模型在少步设置下生成模糊mean field。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GenEval整体评估 | Overall得分 | 0.90 | Meissonic: 0.54, LaViDa-O: 0.77 | 相比Meissonic提升66.7% |
| DPG基准测试 | DPG得分 | 86.9 | GPT-4o: 85.3 | 超过GPT-4o |
| MJHQ-30k图像保真度 | HPSv3 | 10.76（微调后） | SD3: 9.42, Flux-dev: 10.15 | 超过SD3 14.3% |
| ImageNet256类别条件生成 | FID（300 epoch） | 3.00 | SNCE: 3.42, CE: 5.44 | 相比SNCE降低12.3% |
| 推理速度 | 每张图像延迟（1024×1024） | 约2.5秒 | Emu3.5: 约106秒 | 快42.4倍 |
局限与改进
作者在论文中承认的局限性包括：FID指标可能不能完全反映高分辨率图像质量，因为它依赖在低分辨率图像上训练的过时特征提取网络，多项研究表明FID与人类感知的图像质量相关性较差。本文的观察显示，few-step生成虽然可以产生合理质量的图像，但性能仍有提升空间。消融实验中，随着NFE减少，图像质量持续下降，即使启用token editing也无法完全避免。另外，虽然本文使用了8B模型和131K词表，但相比某些前沿模型（如28.5B的Qwen-Image-2512）仍有差距，可能限制了某些复杂场景的生成能力。模型在极端长prompt下的表现也未充分评估，虽然decoder-only架构理论上支持更长context，但实际效果需要验证。
独立分析的弱点
独立分析的弱点包括：第一，计算成本，虽然本文相比自回归模型大幅提升了推理速度，但8B参数和64步采样仍需要显著的算力。对于实时应用或移动端部署可能不够高效。改进方向包括知识蒸馏到更小模型，或进一步减少采样步数。第二，corruption策略的启发式性，token editing的corruption策略（同图像随机采样和embedding最近邻采样）是经验选择的，缺乏理论指导。改进方向可以学习数据驱动的corruption分布，或根据训练动态自适应调整。第三，GCE聚类粒度的选择，本文使用了16384和8192两个粒度，但最优粒度可能因任务和数据而异。改进方向可以设计可学习的聚类结构或动态调整聚类粒度。第四，评估局限，主要依赖自动化指标，缺乏人类评估。改进方向应补充大规模人类偏好评估，特别是与商业产品如DALL-E 3的直接对比。
未来方向
作者提出的和基于成果可延伸的未来方向包括：第一，扩展到视频生成，GCE和token editing机制可以直接应用于视频token序列，但需要额外考虑时间一致性。可以探索时序感知的corruption策略和层级聚类。第二，统一理解-生成模型，本文的decoder-only架构天然适合统一模型。可以扩展到图像理解任务，如VQA和图像描述，利用相同的backbone。第三，更高效的训练，GCE的自定义算子已显著减少开销，但可以进一步探索混合精度训练、梯度累积优化等技术，降低训练成本。第四，多模态扩展，除了文本，可以加入音频、3D等其他模态，构建真正的统一多模态生成模型。第五，实时部署优化，针对实时应用，可以研究更激进的步数减少（如1-4步），或开发专用的推理加速芯片。作者提到的denoising dynamics分析显示MDM在少步设置下有优势，这为进一步优化提供了理论基础。
复现评估
复现评估：论文提供了相对详细的实现细节，但未明确宣布开源计划。训练使用了64块H100 GPU训练300K步，估算总算力需求约3.2秒每步乘以300K步等于960K秒约等于11.1 GPU-days，考虑64块并行约0.17天。但这是低估，因为实际包括数据加载、优化器更新、验证等开销，可能需要数天到一周。渐进式分辨率提升（从256到1024）进一步增加了训练复杂度。数据集组成未完全公开，但提到使用大规模文本-图像语料库。推理相对易于复现，只需要8B模型和131K词表。GCE自定义算子的实现细节在附录中，但可能需要CUDA编程经验。总体而言，完整复现训练过程需要显著算力资源，但推理和消融实验可以在较少资源上验证。代码和模型权重的开放程度将直接影响可复现性。
论文图表
这张图可视化展示了相邻token的语义相似性。使用预计算的K-means聚类，展示了top-1 token和同一cluster中随机选择的token的重建结果。结果表明这些token在语义上是相似的，但标准cross-entropy损失无法捕捉这种关系。图展示了原始图像、top-1 token重建、以及不同cluster粒度（16384、8192、4096）下随机token的重建结果。
这张图直观阐释了codebook sparsity问题和标准cross-entropy的局限性，为GCE目标函数的设计提供了动机。它帮助读者理解为什么语义相似的token应该获得正向监督信号，这是理解本文核心贡献的关键。