Nemotron-Labs-Diffusion-Image：推進遮罩離散擴散以實現高解析度影像合成

摘要

我們提出 Nemotron-Labs-Diffusion-Image，這是一個用於高解析度文字到影像合成的最先進遮罩離散擴散模型（MDM）。與先前的遮罩影像生成研究相比，Nemotron-Labs-Diffusion-Image 解決了兩項關鍵挑戰。首先，與在整個影像上逐步優化潛在表徵的連續擴散模型不同，標準的 MDM 缺乏自我校正能力，因為離散標記一旦被解除遮罩就無法修改。其次，雖然增加離散影像分詞器的詞彙量能提升重建保真度，但也會為生成模型帶來最佳化困難，因為每個標記的訓練訊號變得越來越稀疏。為了解決第一個挑戰，Nemotron-Labs-Diffusion-Image 引入了一個標記編輯機制，使模型能在推論過程中動態修改已解除遮罩的標記，類似於雕塑家反覆精修其作品的方式。為了應對第二個挑戰，我們提出了一個分組交叉熵（GCE）目標函數，該函數為嵌入空間中與真實值相鄰的標記分配正面的學習訊號，從而緩解訊號稀疏性的問題。為了進一步提升訓練效率，我們為 GCE 實作了一個自定義融合運算元，這能顯著減少大詞彙量設定下的 VRAM 使用量。實驗結果表明，這些創新顯著提升了遮罩離散影像生成器的訓練效率和影像保真度，在 GenEval 上達到 0.90 分，在 DPG 上達到 86.9 分，在 HPSv3 上達到 10.76 分。

English

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.