Resumo Técnico: A Dualidade da Difusão, Capítulo II: Ψ-Samplers

Declaração do Problema

Os modelos de difusão discreta emergiram como uma alternativa poderosa aos modelos autoregressivos para geração de linguagem e imagem, oferecendo decodificação paralela e contexto bidirecional. Dentre estes, os Modelos de Difusão de Estado Uniforme (USDMs) destacam-se na geração em poucos passos e em tarefas guiadas devido à sua capacidade de autocorreção, revisando tokens múltiplas vezes. No entanto, os USDMs enfrentam duas limitações críticas:

1. Estagnação da Qualidade de Amostragem: Os samplers ancestrais padrão para USDMs falham em melhorar a qualidade de geração à medida que o número de passos de amostragem (NFEs) aumenta, ao contrário de seu desempenho em regimes de poucos passos.
2. Ineficiência no Treinamento: Tentativas anteriores de estreitar a lacuna de verossimilhança entre USDMs e Modelos de Difusão Mascarada (MDMs) por meio de aprendizado de currículo (por exemplo, Duo) são computacionalmente caras, exigindo alto uso de memória e longos tempos de treinamento devido à necessidade de materializar vetores latentes de alta dimensão durante a fase de treinamento.

Embora os Modelos de Difusão Mascarada (MDMs) geralmente alcancem verossimilhanças mais altas, eles sofrem com erros cumulativos durante a inferência, pois os tokens, uma vez desmascarados, não podem ser remascarados. Trabalhos recentes sobre métodos Preditores-Corretores (PC) (por exemplo, ReMDM) melhoraram os MDMs permitindo o remascaramento, mas avanços semelhantes para USDMs foram pouco explorados, com os métodos PC existentes para USDMs frequentemente tendo desempenho inferior aos samplers ancestrais.

Metodologia

O artigo apresenta o Duo++, um framework que aborda esses desafios através de duas inovações primárias: $\Psi$-samplers para inferência e um currículo escalável para treinamento.

1. $\Psi$-Posteriores e $\Psi$-Samplers

Os autores propõem uma família de posteriores não markovianos, denominados $\Psi$-posteriores, que generalizam os métodos PC anteriores para priores de ruído arbitrários (tanto Mascarados quanto Uniformes).

• Definição: O $\Psi$-posterior é definido como uma superposição linear do posterior reverso padrão ($q_{s|t}$) e a marginal do processo forward ($q_s$):
  $$\Psi_{s|t}(\cdot|x^\ell, z^\ell_t) = \kappa_t q_{s|t}(\cdot|z^\ell_t, x^\ell) + (1-\kappa_t) q_s(\cdot|x^\ell)$$
  onde $\kappa_t \in [0, 1]$ é um coeficiente dependente do tempo.
• Mecanismo: Esta formulação permite que o sampler injete ruído (via o termo $q_s$) durante o processo de geração.
  • Para MDMs, isso habilita o remascaramento: tokens que foram previamente decodificados podem retornar ao estado mascarado se o denoiser cometer um erro, corrigindo erros cumulativos.
  • Para USDMs, isso garante que cada token retenha uma probabilidade não nula de ser amostrado, mesmo que o denoiser lhe atribua probabilidade próxima de zero, impedindo que o sampler fique preso em ótimos locais.
• Garantia Teórica: Os $\Psi$-posteriores são construídos para compartilhar as mesmas marginais exatas do processo de difusão discreta markoviano padrão. Isso assegura que, dado um número suficiente de amostras, o processo convirja para a distribuição de dados correta.
• Generalização: O framework subsume métodos PC existentes (por exemplo, ReMDM, Campbell et al., Gat et al.) como casos especiais onde cronogramas específicos para $\kappa_t$ são escolhidos.

2. Currículo Escalável para Treinamento Mais Rápido

Para abordar o gargalo computacional da estratégia de aprendizado de currículo proposta em Duo (Sahoo et al., 2025a), os autores introduzem uma implementação eficiente em memória.

• Exploração de Esparsidade: O currículo original substitui entradas de tokens discretos por relaxações softmax de baixa temperatura de latentes gaussianos. Para vocabulários grandes ($K > 100.000$), materializar o vetor de pesos completo de $K$ dimensões é inviável.
• Aproximação Top-$k$: Os autores observam que, em baixas temperaturas, a distribuição softmax é esparsa, concentrando massa em apenas algumas entradas. Eles propõem aproximar a combinação linear completa usando apenas as top-$k$ entradas ($k \ll K$).
• Amostragem Eficiente: Em vez de gerar o vetor completo e ordenar, o método usa estatísticas de ordem para amostrar os valores top-$k$ e seus índices correspondentes diretamente das distribuições gaussianas subjacentes, sem materializar o vetor completo.
• Normalização: Uma expressão de forma fechada é derivada para aproximar o normalizador softmax ($\tilde{Z}$) contabilizando os termos top-$k$ amostrados e a contribuição esperada dos termos não amostrados.
• Eficiência: Esta abordagem reduz o uso de pico de memória em 33% e acelera o treinamento em 25%, mantendo o mesmo desempenho de verossimilhança.

Principais Contribuições

1. $\Psi$-Posteriores: Um framework unificado para difusão discreta que introduz posteriores de superposição não markovianos, habilitando a amostragem Preditores-Corretores para priores de ruído arbitrários (Mascarados e Uniformes) enquanto preserva a consistência marginal.
2. $\Psi$-Samplers: Uma família de samplers que generaliza métodos PC anteriores (incluindo ReMDM). Empiricamente, eles permitem que USDMs escalem melhor com NFEs, fechando a lacuna de desempenho com MDMs em geração de texto e superando-os em geração de imagem.
3. Currículo Eficiente: Uma estratégia de treinamento escalável que aproveita a esparsidade de relaxações softmax de baixa temperatura para reduzir significativamente os custos de memória e computação sem sacrificar a qualidade do modelo.

Resultados Experimentais

Modelagem de Linguagem (OpenWebText)

• Desempenho: Os $\Psi$-samplers melhoram significativamente a Perplexidade Generativa (Gen. PPL) para USDMs em comparação com a amostragem ancestral.
• Escala: Ao contrário da amostragem ancestral, que estagna à medida que os NFEs aumentam, os $\Psi$-samplers continuam a melhorar com mais passos. Em entropia unigrama correspondente (diversidade), o Duo++ com $\Psi$-samplers alcança menor Gen. PPL do que o MDLM com ReMDM.
• Tarefas de Descendência: O Duo++ iguala o desempenho do modelo Duo original em tarefas de resposta a perguntas de múltipla escolha (por exemplo, ARC, HellaSwag) enquanto usa 25% menos FLOPs.

Modelagem de Imagem (CIFAR-10)

• Qualidade: No CIFAR-10, o Duo++ com $\Psi$-samplers alcança FID e Pontuação de Inception (IS) superiores em comparação tanto com a amostragem ancestral quanto com o ReMDM.
• Escala: A melhoria na qualidade da imagem é consistente em variados números de passos de amostragem, com o Duo++ superando variantes do MDLM.

Eficiência de Treinamento

• Memória e Velocidade: O currículo eficiente proposto reduz o uso de pico de memória em 33% e reduz o tempo total de treinamento em 25% em comparação com o currículo Duo original.
• Verossimilhança: Os modelos treinados com o currículo eficiente alcançam perplexidade comparável no LM1B e OpenWebText aos modelos Duo originais.

Significado e Alegações

O artigo argumenta que a suposição de que a Difusão Mascarada é o futuro inevitável da modelagem de linguagem baseada em difusão é questionável. Ao introduzir $\Psi$-samplers, os autores demonstram que os Modelos de Difusão de Estado Uniforme podem alcançar:

1. Escala Superior: Desempenho que continua a melhorar com o aumento dos passos de amostragem, ao contrário dos samplers ancestrais.
2. Qualidade Competitiva ou Superior: Igualar ou exceder MDMs em geração de texto (quando emparelhados com amostragem PC) e superá-los em geração de imagem.
3. Eficiência Prática: Um pipeline de treinamento significativamente mais eficiente em memória e tempo do que abordagens anteriores de aprendizado de currículo.

O trabalho sugere que a "dualidade" entre difusão gaussiana e discreta, combinada com estratégias de amostragem não markovianas, oferece um caminho robusto e escalável para a modelagem generativa discreta, desafiando a dominância das abordagens mascaradas.