Nemotron-Labs-Diffusion-Image: Avançando a Difusão Discreta Mascarada para Síntese de Imagens de Alta Resolução

Resumo

Propomos o Nemotron-Labs-Diffusion-Image, um modelo de difusão discreta mascarada (MDM) de ponta para síntese de texto para imagem de alta resolução. Em comparação com trabalhos anteriores em geração de imagens mascaradas, o Nemotron-Labs-Diffusion-Image aborda dois desafios principais. Primeiro, ao contrário dos modelos de difusão contínua, que refinam progressivamente representações latentes em toda a imagem, os MDMs padrão carecem de capacidade de autocorreção, pois tokens discretos não podem ser modificados uma vez que são desmascarados. Segundo, embora aumentar o tamanho do vocabulário dos tokenizadores de imagem discretos melhore a fidelidade da reconstrução, isso introduz dificuldades de otimização para a modelagem generativa, pois o sinal de treinamento por token torna-se cada vez mais esparso. Para lidar com o primeiro desafio, o Nemotron-Labs-Diffusion-Image incorpora um mecanismo de edição de tokens que permite ao modelo revisar dinamicamente tokens já desmascarados durante a inferência, de forma similar a como um escultor refina iterativamente seu trabalho. Para enfrentar o segundo desafio, propomos um objetivo de Entropia Cruzada Agrupada (GCE) que atribui sinais de aprendizado positivos a tokens vizinhos ao valor real no espaço de embeddings, aliviando assim a esparsidade do sinal. Para melhorar ainda mais a eficiência do treinamento, implementamos um operador fundido personalizado para GCE que reduz significativamente o uso de VRAM em configurações de vocabulário grande. Resultados experimentais demonstram que essas inovações melhoram substancialmente tanto a eficiência do treinamento quanto a fidelidade da imagem de geradores de imagens discretas mascaradas, alcançando uma pontuação de 0,90 no GenEval, 86,9 no DPG e 10,76 no 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.