PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Shufan Li, Greg Heinrich, Hanrong Ye, Yonggan Fu, Aditya Grover, Jan Kautz, Pavlo Molchanov

arXiv:2606.29814v2cs.CV

TL;DR

Masked discrete diffusion models struggle to self-correct unmasked tokens and to train large-vocabulary generators. Nemotron-Labs-Diffusion-Image introduces token editing and Grouped Cross-Entropy, with experiments showing strong performance across text-to-image benchmarks.

  • Problem

    Masked discrete image generators lack self-refinement and face sparse training signals when using large codebooks needed for high-fidelity synthesis.

  • Method

    The model combines inference-time token editing with Grouped Cross-Entropy supervision for large-vocabulary discrete image generation.

  • Results

    Experiments show strong performance across a wide range of text-to-image benchmarks, including GenEval, DPG, and MJHQ-30k.

  • Takeaways & Limitations

    The proposed mechanisms support iterative refinement, reduce training-signal sparsity, and improve the practical scalability of masked discrete image generation.

  • Takeaways & Limitations

    Uniformly replacing clean tokens with random vocabulary tokens performs poorly because their distribution mismatches incorrectly predicted tokens during inference.

Abstract

from arXiv · show

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.

1. Introduction

Nemotron-Labs-Diffusion-Image advances masked discrete image generation by addressing vanilla MDMs’ lack of self-refinement and the optimization difficulty caused by large codebooks. It introduces token editing and grouped cross-entropy (GCE), with evaluations showing strong performance across text-to-image benchmarks.

  • Introduction: Nemotron-Labs-Diffusion-Image targets two fundamental challenges: limited self-refinement and difficult training with large codebooks.The model is presented as an 8B foundational masked discrete diffusion model for text-to-image synthesis.
  • Self-refinement challenge: Vanilla MDMs permanently commit tokens after unmasking, preventing correction of incorrect predictions and effectively factorizing generation into position-wise marginal decisions.Tokens unmasked at each step are sampled independently from position-wise logits, creating an implicit inter-token independence assumption.
  • Token editing: Token editing enables MDMs to correct and refine previously unmasked tokens by training on both masked and corrupted visible tokens.The model predicts token probabilities at all positions, supporting refinement beyond the initial unmasking decision.
  • Large-codebook training: Larger codebooks improve reconstruction fidelity but make downstream generative modeling harder because token-level training signals become increasingly sparse.The difficulty can require greater model capacity and more training data.
  • Grouped cross-entropy: GCE supplies auxiliary supervision to non-top-1 tokens semantically close to the ground truth by grouping codebook entries with K-means.For example, a 100k-entry codebook can be clustered into 16k groups, with cluster probabilities formed by summing probabilities within each group.
  • Experimental evaluation: Experiments report strong performance across text-to-image benchmarks including GenEval, DPG, and MJHQ-30k.The evaluations are used to assess the effectiveness of the proposed token-editing and GCE approaches.

2. Background and Related Work

Discrete image generation progressed from slow autoregressive token prediction to masked discrete diffusion models that generate multiple tokens per step. However, MDMs cannot revise unmasked tokens, while larger codebooks improve fidelity but create codebook sparsity through sharply reduced per-token frequency.

  • Discrete Image Generation: Masked discrete diffusion models generate multiple image tokens at each inference step, addressing the slow inference speed of autoregressive generators.Discrete image tokenizers encode images as sequences of discrete codes for subsequent sequence modeling.
  • Masked Discrete Diffusion: MDMs start from a fully masked sequence and progressively unmask tokens until a clean sequence is obtained.Intermediate sequences contain both mask tokens and clean tokens.
  • Masked Discrete Diffusion: During training, MDMs randomly mask clean sequences and predict the original tokens at masked positions.The timestep is sampled uniformly from [0, 1], and masking follows the forward diffusion process q(y_t|y_0).
  • Limitations of MDMs: Once unmasked during inference, MDM tokens remain fixed, preventing revision of earlier decisions and causing error accumulation.Recent language-domain work introduced self-correction mechanisms to address this limitation.
  • Codebook Scaling: 31,250 average occurrences per code arise with an 8,192-entry codebook, whereas a 200K codebook sharply reduces per-token frequency through codebook sparsity.Larger vocabularies improve tokenizer expressiveness and image fidelity, especially for large-scale, high-resolution text-to-image models.

3. Method

Nemotron-Labs-Diffusion-Image uses a decoder-only transformer for joint text-and-image processing and adds token editing for iterative self-correction. Its training modifies corruption and introduces hierarchical grouped cross-entropy to provide positive supervision for semantically similar tokens in large codebooks.

  • Architecture: Nemotron-Labs-Diffusion-Image uses a single decoder-only transformer to process both text prompts and image tokens.Unlike prior encoder-decoder designs, the architecture is optimized end-to-end.
  • Architecture: The decoder-only architecture avoids a frozen text encoder’s context-length constraint and supports sequence packing optimization.Meissonic’s CLIP encoder supports a maximum of 77 tokens per prompt.
  • Token Editing: Token editing enables iterative refinement by predicting correction distributions for clean image tokens and replacing already-unmasked tokens when confidence exceeds threshold τ.Inference combines standard unmasking with confidence-based edits.
  • Token Editing: Training uses same-image replacements and tokenizer-embedding nearest-neighbor replacements instead of uniformly random vocabulary tokens, which mismatch inference.The modified objective computes loss at all token positions and includes corrupted, clean, and masked tokens.
  • Grouped Cross-Entropy: Hierarchical grouped cross-entropy addresses sparse large-codebook supervision by clustering semantically similar tokens and adding cluster-level cross-entropy losses.The implementation uses a 132K-token vocabulary with clustering granularities including 16,384 and 8,192 clusters.
  • Grouped Cross-Entropy: GCE gives every code in the target cluster positive gradients and provides supervision proportional to semantic proximity to the ground-truth token.Within a target cluster, gradient magnitude is proportional to each token’s post-softmax probability.

4. Experiments

Experiments evaluate Nemotron-Labs-Diffusion-Image on text-to-image benchmarks, token editing and GCE ablations, operator efficiency, inference latency, and few-step generation. Results show improved fidelity, smoother quality degradation with token editing, reduced operator cost, faster inference, and reasonable four-step outputs.

  • Text-to-Image Generation: Experiments evaluate text-to-image generation on GenEval, DPG, and MJHQ, using benchmark scores plus FID and HPSv3 for image fidelity.The paper notes that FID may poorly reflect high-resolution and human-perceived image quality, whereas HPSv3 provides an alternative fidelity metric.
  • Text-to-Image Generation: Nemotron-Labs-Diffusion-Image outperforms Meissonic and unified multimodal models such as LaViDa-O under the masked diffusion paradigm.Comparisons also include Flux-dev, SD3-Medium, DALLE-3, BAGEL, and MMaDa.
  • Token Editing: Token editing consistently improves image fidelity by correcting artifacts and refining texture details.With fixed random seeds, qualitative comparisons show better generated images when the token-editing pipeline is enabled.
  • Token Editing: Token editing improves HPSv3 quality across all sampling NFEs and makes quality degradation substantially smoother as NFEs decrease.Reducing NFEs lowers image quality in both settings, but the decline is less abrupt with token editing.
  • Optimized Operator: 44 ms to 20 ms: the optimized GCE operator reduces latency while lowering maximum VRAM usage from 25 GB to 16 GB.The benchmark processes 16,384 tokens on an H100 GPU and compares eager and optimized implementations.
  • Generation Latency: 42.4× faster: Nemotron-Labs-Diffusion-Image than Emu3.5, while also achieving a higher GenEval score.The inference-latency comparison includes Qwen-Image, Infinity-8B, and Emu3.5.
  • Few-Step Generation: 4 steps: Nemotron-Labs-Diffusion-Image generates images with reasonable quality without distillation, unlike continuous flow-matching models that produce blurry mean fields at few steps.Most experiments use 64 diffusion steps, while Figure 7 examines few-step generation.

5. Conclusion … A.2. Token Editing

Nemotron-Labs-Diffusion-Image combines token editing with Grouped Cross-Entropy to address irreversible token decisions and sparse supervision in high-resolution masked discrete diffusion. The appendix details masked diffusion formulation and an edit-aware inference process using embedding-space neighboring tokens.

  • 5. Conclusion: Nemotron-Labs-Diffusion-Image targets high-resolution text-to-image synthesis as a masked discrete diffusion model.
  • 5. Conclusion: Token editing enables iterative refinement during inference, addressing a limitation of standard masked discrete models.
  • 5. Conclusion: Grouped Cross-Entropy assigns positive supervision to semantically neighboring embedding-space tokens, alleviating sparse training signals in large-vocabulary discrete spaces.
  • A.1. Formulation of Masked Diffusion Models: The masked diffusion formulation gradually replaces original tokens with [M], producing a fully masked sequence at t = 1.
  • A.1. Formulation of Masked Diffusion Models: During standard sampling, unmasked tokens remain unchanged in subsequent steps, while the clean sequence is replaced by the model’s prediction during reverse inference.
  • A. Additional Technical Details: Token editing invalidates assumptions underlying the standard objective, motivating an edit-aware formulation.
  • A.2. Token Editing: The corruption distribution places probability mass on tokens from the same image and neighboring embedding-space tokens, reusing K-means clusters from GCE.
  • A.2. Token Editing: At each inference step, the method performs standard unmasking and computes model predictions for token editing; τ = 0.6 and α = 0.1 work best experimentally.

A.3. Grouped Cross-Entropy … B.2. Training Setup and Hyperparameters

The paper details a memory-efficient implementation of Grouped Cross-Entropy and reports the dataset, captioning pipeline, tokenizer, and multistage training configuration used for Nemotron-Labs-Diffusion-Image.

  • A.3. Grouped Cross-Entropy: Grouped Cross-Entropy handles logits shaped N×|V|, using torch.logsumexp for one term while accommodating target clusters with different sizes.Here, N = SeqLen×NumSeqs.
  • A.3. Grouped Cross-Entropy: 391 and 192 are the largest cluster sizes for 8,192 and 16,384 clusters, respectively, requiring less than 1% of naive memory in both settings.The optimized allocation uses N×|Cmax| rather than N×|V|.
  • A.3. Grouped Cross-Entropy: The custom backward pass computes dense and cluster-sparse gradients, combines them with in-place torch.scatter_add operations, and avoids allocating full-vocabulary intermediate tensors.The sparse term uses N×|Cmax| memory, while gradient updates are applied through scatter_add_.
  • A.3. Grouped Cross-Entropy: The implementation negates the GCE loss and gradients because J_GCE is a maximization objective, unlike vanilla PyTorch cross-entropy.This makes the objective compatible with minimization-based cross-entropy semantics.
  • B. Additional Experiment Details and Results: Additional experiments include qualitative samples and ablation studies alongside further training details and experimental results.This passage introduces the contents of the additional-experiments section.
  • B.1. Training Data: 137M text-image pairs from LAION-2B, COYO-700M, BLIP3o-60k, and ShareGPT4o-Image form the filtered training dataset.Filtering removes NSFW prompts, low CLIP-score samples, low aesthetic-score samples, and low-resolution images.
  • B.1. Training Data: Qwen3-VL recaptions most LAION-2B and COYO-700M images, while selected high-CLIP-score raw captions are retained to support keyword prompting.The pipeline randomly chooses between VLM-generated and retained raw captions, including keywords such as “high quality” and “4k”.
  • B.2. Training Setup and Hyperparameters: Training uses the 131,072-token Emu-3.5 tokenizer and proceeds from 256 × 256 pretraining to 512 × 512 and 1024 × 1024 stages on 64 H100 GPUs.The stages run for 200k steps at global batch size 1024, then 20k steps at 512 × 512 and 80k additional steps at 1024 × 1024, both with global batch size 256.

B.3. Ablation Studies of Editing Thresholds

The ablation studies examine editing thresholds across different numbers of inference steps using HPSv3 scores. Token editing consistently outperforms no editing, with τ = 0.6 achieving the best image quality for most NFEs.

  • Threshold Ablation: The study varies the editing threshold τ under different numbers of inference steps and evaluates the resulting HPSv3 scores.The results are visualized in Figure 9.
  • Threshold Ablation: Token editing consistently outperforms the no-editing baseline across the evaluated settings.This establishes a consistent benefit from enabling token editing.
  • Threshold Ablation: τ = 0.6 achieves the best image quality for most NFEs among the evaluated settings.The result holds across most tested numbers of inference steps.

B.4. Ablation Studies of Cluster Sizes

The ablation varies GCE cluster counts on ImageNet-256 using FID scores. Combining 16,384- and 8,192-cluster supervision outperforms either clustering level alone, while the 16,384-only setting performs better than the coarser 8,192-only setting.

  • B.4. Ablation Studies of Cluster Sizes: The study evaluates how varying GCE cluster counts affects FID on ImageNet-256.The ablation compares different numbers of clusters used in the GCE objective.
  • B.4. Ablation Studies of Cluster Sizes: Combining 16,384- and 8,192-cluster supervision performs better than using either clustering level alone.The comparison is reported in Table 6’s ImageNet-256 ablation.
  • B.4. Ablation Studies of Cluster Sizes: With a single clustering level, 16,384 clusters outperform 8,192 clusters because the latter provides coarser, less refined supervision signals.The explanation is presented as a presumption about why the finer clustering level performs better.

B.5. Ablation Studies of Corruption Type and Scale · B.6. Additional Qualitative Results

The ablations show that combining neighboring-token and resampled-token corruption performs best, while excessive corruption harms optimization; qualitative comparisons show token editing improves fidelity by refining details and correcting artifacts.

  • B.5. Ablation Studies of Corruption Type and Scale: Combining neighboring-token and resampled-token corruption achieves the best HPSv3 performance after 10K steps of 1024-resolution Stage-2 training.The comparison includes random noise corruption, neighboring embedding-space tokens, and resampled tokens from the same input image.
  • B.5. Ablation Studies of Corruption Type and Scale: Reasonable corruption ratios produce no significant image-quality differences, but excessive corruption degrades quality.The final experiments use α = 0.1.
  • B.5. Ablation Studies of Corruption Type and Scale: At α = 0.5, the lower signal-to-noise ratio makes clean and corrupted tokens harder to distinguish, increasing optimization difficulty.The passage identifies high corruption as the cause of degraded image quality.
  • B.6. Additional Qualitative Results: Additional text-to-image samples further demonstrate the effectiveness of Nemotron-Labs-Diffusion-Image.Figure 10 presents additional text-to-image generation results.
  • B.6. Additional Qualitative Results: Token editing improves generated images by refining details and correcting artifacts.This observation is reported consistently across the qualitative comparisons.

C. Compute Resources

Training uses 64 H100 GPUs across 8 nodes and takes approximately 16 days in total.

  • C. Compute Resources: Training uses 64 H100 GPUs across 8 nodes and takes approximately 16 days.The total training duration is approximately 16 days.

D. Limitations · E. Broader Impact · F. Licenses

The paper identifies residual generation artifacts and possible optimization improvements, cautions that the model may enable harmful content or inherit biases, and states that it is intended primarily for research. It also lists the licensed models and datasets used.

  • D. Limitations: Token editing improves image quality but does not eliminate all artifacts, and the model may still generate erroneous outputs.The authors identify residual artifacts and erroneous generations as an unresolved limitation.
  • D. Limitations: Further improvements to GCE optimization may be possible through customized low-level CUDA kernels.The current fused operator improves efficiency, but the authors leave additional kernel-level optimization for future work.
  • E. Broader Impact: The model’s text-to-image capabilities may be misused to generate harmful or offensive content.The authors strongly caution against these use cases.
  • E. Broader Impact: The model may inherit biases from the base language model and training data, and is intended primarily for research purposes.The authors do not recommend using it for other purposes.
  • F. Licenses: The listed models are Emu-3.5-Tokenizer [17] (Apache-2.0), Qwen3-VL [60] (Apache-2.0), and Nemotron-Labs-Diffusion [41] (Nvidia Open Model License).These are the model assets identified in the paper’s licenses section.
  • F. Licenses: The listed datasets are LAION [55] (MIT), COYO (CC-BY-4.0), MJHQ (CC-BY-4.0), BLIP3o-60k [45] (Apache-2.0), and ShareGPT4o-Image [57] (CC-BY-4.0).These datasets and their licenses comprise the dataset assets reported by the authors.
Loading 2606.29814v2…