Nemotron-Labs-Diffusion-Image: Advancing Masked Discrete Diffusion for High-Resolution Image Synthesis
Abstract Overview
This paper presents Nemotron-Labs-Diffusion-Image (NLD-Image), an 8B masked discrete diffusion model for high-resolution text-to-image synthesis. The work targets two limitations of standard masked discrete diffusion models: they cannot revise tokens once unmasked, and large tokenizer vocabularies make training signals sparse. To address these issues, the model adds a token-editing mechanism for inference-time self-correction and a grouped cross-entropy objective that gives auxiliary supervision to semantically related tokens in the codebook. The system also uses a decoder-only transformer initialized from a pretrained diffusion language model and introduces an optimized fused operator to make grouped cross-entropy more efficient in large-vocabulary settings.
Novelty
The paper's main novelty is extending masked discrete diffusion for images with explicit token editing, allowing previously unmasked image tokens to be revised during sampling. It also introduces grouped cross-entropy, a hierarchical cluster-based objective for large codebooks that differs from standard one-hot training and from prior soft-label approaches such as SNCE.
Results
On reported text-to-image benchmarks, NLD-Image reaches 0.90 on GenEval, 85.2 on DPG, and 9.61 HPSv3 on MJHQ-30k, with a synthetic-data-finetuned variant reaching 86.9 DPG and 10.76 HPSv3. Ablations indicate token editing improves image quality across sampling budgets and that 32-step generation with editing is comparable to 64-step generation without editing. In controlled ImageNet256 experiments, grouped cross-entropy outperforms both standard cross-entropy and SNCE, and the optimized implementation reduces GCE latency from 44 ms to 20 ms while lowering maximum VRAM from 25 GB to 16 GB.
Key Points
- NLD-Image augments masked discrete diffusion with token editing so the model can iteratively correct already unmasked image tokens during inference.
- The grouped cross-entropy objective provides cluster-level supervision for semantically neighboring tokens, addressing sparse learning signals in large codebooks.
- Empirically, the model reports strong 1024px text-to-image results and ablations showing gains in quality, step efficiency, and training/runtime efficiency.
References
- arXiv: https://arxiv.org/abs/2606.29814v1
- Fugu-MT: https://fugumt.com/fugumt/paper_check/2606.29814v1
- Hugging Face Papers: https://huggingface.co/papers/2606.29814