Nemotron-Labs-Diffusion-Image: Развитие маскированной дискретной диффузии для синтеза изображений высокого разрешения

Аннотация

Мы представляем Nemotron-Labs-Diffusion-Image — современную маскированную дискретную диффузионную модель (MDM) для высококачественного синтеза изображений по текстовому описанию. По сравнению с предыдущими работами по маскированной генерации изображений, Nemotron-Labs-Diffusion-Image решает две ключевые проблемы. Во‑первых, в отличие от непрерывных диффузионных моделей, которые постепенно уточняют латентные представления по всему изображению, стандартные MDM лишены способности к самокоррекции, поскольку дискретные токены невозможно изменить после их размаскирования. Во‑вторых, хотя увеличение размера словаря дискретных токенизаторов изображений повышает качество реконструкции, это создает оптимизационные трудности для генеративного моделирования из‑за того, что обучающий сигнал для каждого токена становится все более разреженным. Для решения первой проблемы Nemotron-Labs-Diffusion-Image включает механизм редактирования токенов, который позволяет модели динамически пересматривать уже размаскированные токены в ходе инференса — подобно тому как скульптор итеративно дорабатывает свою работу. Для решения второй проблемы мы предлагаем целевую функцию группированной кросс-энтропии (GCE), которая присваивает положительные обучающие сигналы токенам, соседствующим с истинными токенами в пространстве вложений, тем самым смягчая разреженность сигнала. Для дальнейшего повышения эффективности обучения мы реализовали пользовательский объединенный оператор для GCE, который значительно снижает потребление видеопамяти при больших размерах словаря. Экспериментальные результаты показывают, что эти инновации существенно улучшают как эффективность обучения, так и качество изображений маскированных дискретных генераторов, достигая оценки 0,90 на GenEval, 86,9 на DPG и 10,76 на HPSv3.

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.