💻 computer science

Nemotron-Labs-Diffusion-Image: Advancing Masked Discrete Diffusion for High-Resolution Image Synthesis

Nemotron-Labs-Diffusion-Image 是一种最先进的掩码离散扩散模型，它通过引入用于动态推理修正的标记编辑机制以及一种用于克服大词汇表设置下信号稀疏性的带有自定义融合算子的分组交叉熵目标函数，增强了高分辨率文本到图像的合成能力，并在 GenEval、DPG 和 HPSv3 基准测试中取得了卓越的性能。

原作者： Shufan Li, Greg Heinrich, Hanrong Ye, Yonggan Fu, Aditya Grover, Jan Kautz, Pavlo Molchanov

发布于 2026-06-30

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Shufan Li, Greg Heinrich, Hanrong Ye, Yonggan Fu, Aditya Grover, Jan Kautz, Pavlo Molchanov

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图教一个机器人根据你给出的描述来创作一幅杰作，比如“一只戴着帽子的猫在阳光明媚的日子里”。长期以来，最好的机器人使用一种类似于从一块粘土中雕刻的方法：它们从一个模糊的混乱状态开始，然后逐步、分阶段地进行精细化处理，在过程中不断修正错误。这被称为连续扩散（Continuous Diffusion）。

最近，一种被称为**离散扩散（Discrete Diffusion）**的新方法变得流行起来。这种方法不再使用粘土，而是使用一大袋乐高积木（即 Token/标记）。机器人从一堵隐藏的积木墙开始，然后逐个揭示它们。问题在于？一旦一块积木被揭示出来，它就固定死了。如果机器人在制作猫的耳朵时选错了积木，它就无法回头更换它。这就像一位雕塑家，一旦他凿掉了一块石头，就再也无法修复了。

这篇论文介绍了一个新的“雕塑家”——Nemotron-Labs-Diffusion-Image，它解决了这种“乐高式”方法中的两个主要问题。

1. “重来”按钮（Token 编辑）

问题： 在标准的乐高绘画中，如果你在应该放蓝色积木的地方放了一个红色积木，你就被困住了。机器人没有办法说：“等等，我改变主意了，”然后稍后再去修复它。这会导致微小的错误不断堆积，直到整幅画看起来很奇怪。

解决方案： 作者给了机器人一个“重来”按钮。他们不仅训练它去揭示隐藏的积木，还训练它去观察那些已经放置好的积木，并说：“实际上，那个看起来不对；让我们把它换掉。”

类比： 想象一位正在雕刻雕像的雕塑家。在旧的方法中，一旦他凿掉了一块石头，那块石头就永远消失了。而在这种新方法中，雕塑家可以轻轻地凿掉刚刚移除的一块碎片，将其磨平，并在看起来不太对劲时重新塑造它。这使得机器人能够迭代地精炼图像，在前进的过程中修正错误，就像旧有的“粘土”模型一样。

2. “大字典”问题（分组交叉熵）

问题： 为了制作高质量、细节丰富的图像，机器人需要一个庞大的积木词汇表（即 Codebook/码本）。词汇表里的积木越多，图像就越细腻。然而，对于一个拥有超过 10 万个独特积木的字典，机器人在训练期间很少能看到完全相同的特定积木。这就像是在学习一门语言，但每隔几年才能见到一次“大象”这个词。机器人会感到困惑，因为信号过于稀疏。

此外，标准的训练方法将每块错误的积木都视为同样糟糕。如果机器人需要“浅蓝色”积木却选了“深蓝色”，旧的方法会像选了“红色”一样大声尖叫：“错了！”它无法意识到“浅蓝色”和“深蓝色”在含义上其实是邻居。

解决方案： 作者引入了一种新的训练规则，称为分组交叉熵（Grouped Cross-Entropy, GCE）。

类比： 作者不再仅仅惩罚机器人选错了精确的积木，而是告诉它：“你选了一个深蓝色的积木，但我想要浅蓝色的。很接近了！给你部分分数。”
他们将 10 多万个积木组织成更小的组（例如“蓝色组”、“红色组”、“绿色组”）。在训练期间，如果机器人选了一个属于正确组别的积木（即使不是最完美的那个），它也会得到一个正向信号。这有助于机器人学习积木之间的关系，使它更容易学习庞大的词汇表而不至于迷失方向。

3. 速度提升（自定义算子）

问题： 进行这种“分组”数学运算需要极高的计算量。它通常需要大量的计算机显存（VRAM）并会降低速度，使得在大规模模型上进行训练变得困难。

解决方案： 团队构建了一个定制的、专门的工具（一个“融合算子”）来执行这些数学运算。

类比： 想象你在做复杂的会计工作。旧的方法是使用计算器，写下数字，拿起笔，把它写进账本，然后再使用计算器。新方法则是一个专门的机器，它能在一次极其快速的动作中同时完成计算并写入账本。这节省了大量的内存和时间。

结果

论文表明，这个新机器人 Nemotron-Labs-Diffusion-Image 是一个冠军：

质量更好： 与之前的基于乐高的方法相比，它能创造出更清晰、更准确的图像。
速度更快： 由于它可以在过程中自我修正，因此可以用更少的步骤生成图像。
更高效： 它可以处理巨大的细节词汇表而不会导致计算机内存崩溃。

简而言之，他们将一种僵化且易出错的方法，赋予了自我纠错的能力，教会了它理解词汇的“邻里关系”，并建造了一个更快的引擎来运行这一切。其结果是一个既更聪明、更高效的高分辨率图像生成器。

技术摘要：Nemotron-Labs-Diffusion-Image

问题陈述

虽然潜在扩散模型（LDMs）在文本到图像合成领域已取得了最先进的性能，但离散图像生成提供了独特的优势，包括与大语言模型（LLMs）的原生兼容性，以及利用成熟的 LLM 优化技术（如序列打包、KV 缓存）的能力。然而，现有的掩码离散扩散模型（MDMs）面临着两个根本性的挑战，阻碍了它们在可扩展性和保真度方面达到连续 LDM 的水平：

缺乏自我修正能力： 与连续 LDM 通过迭代细化潜在表示不同，标准的 MDM 一旦对某个 token 进行去掩码（unmasking）操作，就会对其进行定论。由于离散 token 在去掩码后无法被修改，错误会在推理过程中不断累积。标准的基于 token 间相互独立的采样假设无法捕捉图像 token 中存在的强依赖关系，导致联合分布并非最优。
大词表中的码本稀疏性： 高保真图像生成需要大型码本（词汇表）以减少离散化误差。然而，增加词汇量会导致训练信号严重稀疏；单个 token 在数据集中出现的频率极低。标准的单热交叉熵（One-hot Cross-Entropy, CE）损失将所有非真值 token 视为同等程度的负样本，未能捕捉到在嵌入空间中虽与真值不同但在语义上接近的 token 之间的语义相似性。这造成了优化困难，并需要庞大的模型容量和数据才能收敛。

方法论

作者提出了 Nemotron-Labs-Diffusion-Image (NLD-Image)，这是一个拥有 8B 参数的基础模型，旨在通过创新的架构和目标函数来解决上述挑战。

1. 架构：仅解码器 Transformer (Decoder-Only Transformer)

与以往使用带有冻结文本编码器的编码器-解码器框架（如 Meissonic）的离散图像生成器不同，NLD-Image 采用了单一的仅解码器 Transformer。

初始化： 模型从预训练的扩散语言模型初始化，从而利用现有的语言理解能力。
优势： 这种设计消除了冻结编码器（如 CLIP 的 77 个 token 限制）的上下文长度约束，支持更长的提示词，并支持通过序列打包提高训练效率。最终的语言模型头被替换为一个新初始化的 MLP，用于预测图像 token。

2. Token 编辑机制（自我修正）

为了实现类似于连续扩散的迭代细化，作者引入了一种 token 编辑机制。

训练： 前向扩散过程被增强，加入了 token 损坏（corruption） 过程。除了掩码掉一部分 token 外，模型还会损坏一部分可见（干净）的 token。模型被训练为能够预测所有位置的概率：既包括对掩码 token 的去掩码，也包括对受损可见 token 的修正。
推理： 在采样期间，模型不仅会对掩码 token 进行去掩码，还会评估已经去掩码的 token。如果模型预测出的不同 token 的置信度高于阈值 $\tau$ ，则会替换现有的 token。这使得模型能够在生成过程中动态地进行自我修正。

3. 分组交叉熵 (Grouped Cross-Entropy, GCE) 目标函数

为了缓解码本稀疏性问题，作者提出了 分组交叉熵 (GCE)，这是一种层次化损失函数。

聚类： 利用基于嵌入空间距离的 K-means 算法，将大型码本（例如 131k 个 token）聚类为较小的组（例如 16k、8k）。
损失函数： 该目标函数结合了标准的 token 级单热 CE 损失与辅助的簇级（cluster-level）CE 损失。模型预测所有 token 的概率，并将这些概率求和以获得簇概率。损失函数同时应用于特定的真值 token 以及多个粒度下的真值簇。
收益： 这为那些在语义上与真值相似（处于同一簇内）的 token 提供了正向监督，即使它们不是精确匹配，从而缓解了信号稀疏问题。
实现： 实现了一个自定义的融合算子（fused operator）来高效计算 GCE，与朴素的 PyTorch 实现相比，显著降低了显存占用和延迟。

核心贡献

Nemotron-Labs-Diffusion-Image： 一个用于高分辨率（1024px）文本到图像合成的 8B 参数最先进掩码离散扩散模型。
Token 编辑： 一种新型机制，使 MDM 能够在推理过程中迭代地细化和自我修正输出，解决了标准离散模型缺乏自我细化能力的问题。
分组交叉熵 (GCE)： 一种训练目标，通过向语义相邻的 token 分配正向学习信号，缓解了码本稀疏性，从而提高了大规模词汇生成器的优化稳定性和效率。
高效实现： 用于 GCE 的自定义融合算子降低了内存占用和计算开销，便于训练大规模词汇量的模型。

实验结果

该模型在 64 张 H100 GPU 上通过渐进式缩放策略（从 256px 到 1024px）进行了 300k 步的训练。

基准测试表现：
- GenEval： 取得了 0.90 的得分，超越了专家模型 Meissonic (0.54) 和统一多模态模型 LaViDa-O (0.77)，并能与前沿模型 Qwen-Image-2507 (0.87) 相媲美。
- DPG： 取得了 86.9 的得分，超过了 SD3-Medium (83.5) 和 Flux-Dev。
- MJHQ-30k： 在微调后取得了 10.76 的 HPSv3 得分，表明具有极高的用户偏好对齐度。
消融研究：
- Token 编辑： 定性与定量结果（HPSv3）显示，token 编辑持续提升了图像保真度并减少了误差累积。它使模型能够在 32 个采样步数（NFE）下达到以往不使用编辑时需要 64 步才能达到的质量，从而在固定质量目标下将推理速度有效提升了一倍。
- GCE： 在 ImageNet 类条件生成任务中，在不同的模型规模和训练轮次下，GCE 在 FID 指标上始终优于传统的交叉熵（Vanilla CE）和 SNCE（软最近邻交叉熵）。
效率： 与 eager 实现相比，优化的 GCE 算子将延迟从 44ms 降低至 20ms，并将最大显存占用从 25GB 降至 16GB。NLD-Image 的推理速度比 Emu3.5 (AR) 快 42.4 倍，同时获得了更高的 GenEval 分数。
少步生成： 该模型在仅需 4 步的情况下即可生成质量尚可的图像，且无需蒸馏，其收敛速度快于连续流匹配（flow-matching）基线模型（后者在低步数下通常表现模糊）。

意义与主张

论文声称，NLD-Image 通过成功弥合离散 MDM 与连续 LDM 之间的性能差距，代表了离散图像生成领域的重大进展。通过引入 token 编辑，作者证明了离散模型可以克服其固有的缺乏自我修正能力的缺陷，从而实现迭代细化。此外，GCE 目标为码本稀疏性问题提供了可扩展的解决方案，允许使用具有表现力的大型词汇表而不会产生过高的训练成本。

作者将这项工作定位为迈向统一多模态模型的一步，利用仅解码器架构在单一框架内处理文本和图像。研究结果表明，配备了这些特定创新技术的离散扩散模型可以实现最先进的保真度和效率，为开发更强大、更具扩展性的图像合成系统铺平了道路。