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

2026-06-29

AI Summary

This work proposes a novel masked discrete diffusion model to address the lack of self-correction capability in standard approaches and the sparsity of training signals under large vocabularies. The method introduces a dynamic token editing mechanism that enables iterative refinement of previously unmasked tokens during inference and incorporates a grouped cross-entropy (GCE) loss to alleviate optimization challenges posed by large vocabularies. Combined with a custom fusion operator to reduce memory consumption, the model achieves significant improvements over existing methods, attaining superior performance on GenEval (0.90), DPG (86.9), and HPSv3 (10.76). These advances effectively enhance both fidelity and training efficiency in high-resolution text-to-image generation.

Abstract

We propose Nemotron-Labs-Diffusion-Image, a state-of-the-art masked discrete diffusion model (MDM) for high-resolution text-to-image synthesis. Compared with prior work on masked image generation, Nemotron-Labs-Diffusion-Image addresses two key challenges. First, unlike continuous diffusion models which progressively refine latent representations across the entire image, standard MDMs lack self-correcting capability because discrete tokens cannot be modified once they are unmasked. Second, although increasing the vocabulary size of discrete image tokenizers improves reconstruction fidelity, it introduces optimization difficulties for generative modeling as the per-token training signal becomes increasingly sparse. To address the first challenge, Nemotron-Labs-Diffusion-Image incorporates a token-editing mechanism that enables the model to dynamically revise already-unmasked tokens during inference, similar to how a sculptor iteratively refines their work. To tackle the second challenge, we propose a Grouped Cross-Entropy (GCE) objective that assigns positive learning signals to tokens neighboring the ground truth in embedding space, thereby alleviating signal sparsity. To further improve training efficiency, we implement a custom fused operator for GCE that significantly reduces VRAM usage in large-vocabulary settings. Experimental results demonstrate that these innovations substantially improve both training efficiency and image fidelity of masked discrete image generators, achieving a score of 0.90 on GenEval, 86.9 on DPG and 10.76 of HPSv3.

Problem

• masked discrete diffusion
• self-correcting capability
• vocabulary size
• signal sparsity
• high-resolution image synthesis

Innovation

• masked discrete diffusion
• token editing
• Grouped Cross-Entropy
• large-vocabulary tokenization
• high-resolution image synthesis