Resumo Técnico: Nemotron-Labs-Diffusion-Image

Declaração do Problema

Embora os Modelos de Difusão Latente (LDMs) tenham alcançado um desempenho de estado da arte na síntese de texto para imagem, a geração de imagens discretas oferece vantagens distintas, incluindo compatibilidade nativa com Grandes Modelos de Linguagem (LLMs) e a capacidade de aproveitar técnicas de otimização estabelecidas para LLMs (ex: empacotamento de sequência, cache de KV). No entanto, os atuais Modelos de Difusão Discreta Mascarada (MDMs) enfrentam dois desafios fundamentais que impedem sua escalabilidade e fidelidade em comparação com os LDMs contínuos:

1. Falta de Autocorreção: Ao contrário dos LDMs contínuos, que refinam iterativamente as representações latentes, os MDMs padrão comprometem-se com um token uma vez que ele é desmascarado. Como os tokens discretos não podem ser modificados após o desmascaramento, os erros se acumulam durante o processo de inferência. A suposição padrão de amostragem de independência entre tokens falha em capturar as fortes dependências presentes nos tokens de imagem, levando a distribuições conjuntas subótimas.
2. Esparsidade do Codebook em Grandes Vocabulários: A geração de imagens de alta fidelidade requer codebooks (vocabulários) grandes para reduzir o erro de discretização. No entanto, aumentar o tamanho do vocabulário leva a uma severa esparsidade nos sinais de treinamento; tokens individuais aparecem infrequentemente no conjunto de dados. A perda de Entropia Cruzada (CE) One-Hot padrão trata todos os tokens não correspondentes ao ground-truth como igualmente negativos, falhando em capturar as semelhanças semânticas entre tokens que estão próximos no espaço de embedding, mas são distintos no codebook. Isso cria dificuldades de otimização e exige uma capacidade de modelo e dados massivos para convergir.

Metodologia

Os autores propõem o Nemotron-Labs-Diffusion-Image (NLD-Image), um modelo fundamental de 8 bilhões de parâmetros projetado para abordar esses desafios por meio de inovações arquitetônicas e de objetivo inovadoras.

1. Arquitetura: Transformer Decodificador Único (Decoder-Only)

Ao contrário de geradores de imagem discretos anteriores (ex: Meissonic) que utilizam um framework encoder-decoder com um codificador de texto congelado, o NLD-Image emprega um único transformer decodificador único.

• Inicialização: O modelo é inicializado a partir de um modelo de linguagem de difusão pré-treinado, aproveitando as capacidades existentes de compreensão de linguagem.
• Vantagens: Este design remove as restrições de comprimento de contexto dos encoders congelados (ex: o limite de 77 tokens do CLIP), suporta prompts mais longos e permite o empacotamento de sequências para eficiência de treinamento. A cabeça de modelagem de linguagem final é substituída por um MLP recém-inicializado para prever os tokens de imagem.

2. Mecanismo de Edição de Tokens (Autocorreção)

Para permitir o refinamento iterativo semelhante à difusão contínua, os autores introduzem um mecanismo de edição de tokens.

• Treinamento: O processo de difusão direta é aumentado para incluir a corrupção de tokens. Além de mascarar um subconjunto de tokens, o modelo corrompe uma parte dos tokens visíveis (limpos). O modelo é treinado para prever probabilidades para todas as posições: desmascarando tokens mascarados e corrigindo tokens visíveis corrompidos.
• Inferência: Durante a amostragem, o modelo não apenas desmascara tokens, mas também avalia tokens já desmascarados. Se o modelo prever um token diferente com confiança acima de um limiar $\tau$, o token existente é substituído. Isso permite que o modelo se autocorrija dinamicamente durante o processo de geração.

3. Objetivo de Entropia Cruzada Agrupada (GCE)

Para mitigar a esparsidade do codebook, os autores propõem a Entropia Cruzada Agrupada (GCE), uma função de perda hierárquica.

• Agrupamento (Clustering): O grande codebook (ex: 131k tokens) é agrupado em grupos menores (ex: 16k, 8k) usando K-means baseado na proximidade do espaço de embedding.
• Função de Perda: O objetivo combina a perda CE one-hot padrão ao nível do token com perdas CE auxiliares ao nível do cluster. O modelo prevê probabilidades para todos os tokens, e estes são somados para obter as probabilidades do cluster. A perda é aplicada tanto ao token específico do ground-truth quanto ao cluster do ground-truth em múltiplas granularidades.
• Benefício: Isso fornece supervisão positiva para tokens que são semanticamente similares (no mesmo cluster) ao ground-truth, mesmo que não sejam a correspondência exata, aliviando assim a esparsidade do sinal.
• Implementação: Um operador fundido customizado é implementado para calcular o GCE de forma eficiente, reduzindo significativamente o uso de VRAM e a latência em comparação com implementações ingênuas em PyTorch.

Principais Contribuições

1. Nemotron-Labs-Diffusion-Image: Um modelo de difusão discreta mascarada de 8B de estado da arte para síntese de texto para imagem de alta resolução (1024px).
2. Edição de Tokens: Um mecanismo inovador que permite aos MDMs refinarem iterativamente e autocorrigirem saídas durante a inferência, abordando a falta de autorrefinamento em modelos discretos padrão.
3. Entropia Cruzada Agrupada (GCE): Um objetivo de treinamento que alivia a esparsidade do codebook ao atribuir sinais de aprendizado positivos a tokens semanticamente vizinhos, melhorando a estabilidade e a eficiência da otimização para geradores de grande vocabulário.
4. Implementação Eficiente: Um operador fundido customizado para o GCE que reduz a pegada de memória e o overhead computacional, facilitando o treinamento de modelos de grande vocabulário.

Resultados Experimentais

O modelo foi treinado por 300k passos em 64 GPUs H100 usando uma estratégia de upscaling progressivo (256px para 1024px).

• Desempenho em Benchmarks:
  • GenEval: Alcançou uma pontuação de 0.90, superando modelos especialistas como Meissonic (0.54) e modelos multimodais unificados como LaViDa-O (0.77), e competindo com modelos de fronteira como Qwen-Image-2507 (0.87).
  • DPG: Alcançou 86.9, superando o SD3-Medium (83.5) e o Flux-Dev.
  • MJHQ-30k: Alcançou uma pontuação HPSv3 de 10.76 (após fine-tuning), indicando alta conformidade com a preferência humana.
• Estudos de Ablação:
  • Edição de Tokens: Resultados qualitativos e quantitativos (HPSv3) mostram que a edição de tokens melhora consistentemente a fidelidade da imagem e reduz o acúmulo de erros. Ela permite que o modelo alcance qualidade comparável com 32 passos de amostragem (NFE) que anteriormente exigiam 64 passos sem edição, efetivamente dobrando a velocidade de inferência para um alvo de qualidade fixo.
  • GCE: No ImageNet com geração condicionada por classe, o GCE superou consistentemente tanto a Entropia Cruzada vanilla quanto o SNCE (Soft Nearest Neighbor Cross-Entropy) em pontuações FID através de diferentes tamanhos de modelo e épocas de treinamento.
• Eficiência: O operador GCE otimizado reduziu a latência de 44ms para 20ms e o uso máximo de VRAM de 25GB para 16GB em comparação com implementações eager. O NLD-Image demonstrou uma velocidade de inferência 42.4x mais rápida que o Emu3.5 (AR) enquanto alcançava pontuações GenEval mais altas.
• Geração de Poucos Passos: O modelo gera imagens de qualidade razoável em apenas 4 passos sem destilação, convergindo mais rápido do que baselines de flow-matching contínuo que frequentemente permanecem borrados em contagens baixas de passos.

Significância e Alegações

O artigo afirma que o NLD-Image representa um avanço significativo na geração de imagem discreta ao conseguir reduzir a lacuna de desempenho entre os LDMs contínuos e os MDMs discretos. Ao introduzir a edição de tokens, os autores demonstram que modelos discretos podem superar sua inerente falta de autocorreção, permitindo o refinamento iterativo. Além disso, o objetivo GCE fornece uma solução escalável para o problema da esparsidade do codebook, permitindo o uso de vocabulários grandes e expressivos sem custos de treinamento proibitivos.

Os autores posicionam este trabalho como um passo em direção a modelos multimodais unificados, aproveitando a arquitetura decoder-only para processar texto e imagens dentro de um único framework. Os resultados sugerem que modelos de difusão discreta, quando equipados com essas inovações específicas, podem alcançar fidelidade e eficiência de estado da arte, pavimentando o caminho para sistemas de síntese de imagem mais poderosos e escaláveis.