Resumo Técnico: Modelos de Difusão Uniforme Revisitados

Declaração do Problema

Os modelos de difusão discreta emergiram como uma alternativa competitiva à geração autoregressiva para linguagem e domínios discretos estruturados. Existem dois paradigmas principais: Modelos de Difusão Mascarada (MDM), que substituem tokens por um símbolo de máscara dedicado, e Modelos de Difusão Uniforme (UDM), que substituem tokens por um token amostrado uniformemente do vocabulário. Embora os MDMs tenham demonstrado recentemente perplexidade superior em tamanhos de vocabulário grandes, as razões fundamentais para a lacuna de desempenho entre MDMs e UDMs permanecem obscuras. Especificamente, é incerto se essa lacuna decorre das diferenças intrínsecas em suas marginais de corrupção (mascaramento vs. ruído uniforme) ou das escolhas de parametrização e esquemas de amostragem tipicamente associados a cada processo.

Além disso, implementações padrão de UDM frequentemente utilizam uma parametrização de "plug-in de ponte", onde uma rede neural prevê um proxy de token limpo $\hat{x}_0(x_t, t)$ que é inserido diretamente na ponte de transição reversa. Surge uma questão teórica crítica: essa parametrização de plug-in otimiza a posterior de remoção de ruído padrão $p_0|t(x_0|x_t)$, ou otimiza um alvo diferente?

Metodologia e Principais Contribuições

O artigo revisita os UDMs através de duas lentes complementares: a otimização da parametrização do kernel reverso e a reformulação do processo forward.

1. A Optimalidade "Leave-One-Out" (LOO) dos UDMs de Plug-in

Os autores demonstram que, para UDMs, a parametrização padrão de plug-in de ponte não minimiza o Limite Inferior da Evidência (ELBO) na posterior de remoção de ruído padrão. Em vez disso, a previsão ótima para a parametrização de plug-in é a posterior leave-one-out (LOO).

• O Descompasso: Nos MDMs, a ponte é afim no token limpo, fazendo com que as parametrizações de plug-in e marginalização coincidam, com a posterior de remoção de ruído como o ótimo. Nos UDMs, a ponte é não linear devido ao termo de normalização. Consequentemente, a parametrização de plug-in visa uma quantidade onde a previsão para o token limpo na posição $\ell$ é condicionada a todos os tokens ruidosos exceto aquele na posição $\ell$ (ou seja, $p(x_0^\ell | x_t^{-\ell})$).
• Derivação Teórica: O artigo deriva fórmulas exatas de conversão entre o denoiser padrão, a posterior LOO e o score concreto. Essas fórmulas permitem:
  • Treinar um preditor LOO usando perda de entropia cruzada padrão convertendo a saída da rede.
  • Converter um denoiser treinado na representação LOO necessária para amostragem informada sem treinamento auxiliar.
• Implicação Arquitetural: A posterior LOO é estruturalmente invariante à observação local $x_t^\ell$. Embora os "Transformadores Ocos" (que bloqueiam explicitamente a atenção para a posição atual) imponham isso, os autores mostram que mecanismos de atenção padrão podem aprender essa invariância, e a sensibilidade da previsão ao token local serve como um diagnóstico para suboptimalidade.

2. Reformulação de Estado Absorvente (AUDM e MUDM)

Para fechar a lacuna entre UDM e MDM, os autores propõem reformular a difusão uniforme como um processo de estado absorvente condicionado a variáveis latentes.

• Difusão Uniforme de Estado Absorvente (AUDM): O processo forward é visto como uma mistura de transições de estado absorvente onde o estado absorvente em cada posição é um token aleatório amostrado da distribuição uniforme. Essa construção preserva a lei conjunta do UDM, mas a decompõe em operações que se assemelham ao MDM (desmascaramento de carry-over).
• Difusão Uniforme Mascarada (MUDM): Ao condicionar em tempos de transição latentes, os autores mostram que o problema de remoção de ruído do UDM pode ser mapeado exatamente para um problema de remoção de ruído do MDM. Isso permite reutilizar denoisers de MDM para definir um modelo com marginais de UDM.
• Mecanismo de Resampling: Os autores introduzem uma etapa de resampling para os estados absorventes (ou tempos de transição) durante a cadeia reversa. Esse amostrador "elevado" (ReAUDM) recupera a lei exata da cadeia reversa do UDM, enquanto herda as simplificações estruturais e as propriedades de carry-over da difusão mascarada.

3. Melhorias na Inferência

Aproveitando a representação LOO, o artigo introduz um amostrador preditor-corretor que não requer o treinamento de um modelo auxiliar.

• A posterior LOO fornece a distribuição condicional exata de uma coordenada necessária para uma etapa corretora de Gibbs.
• Esse corretor pode ser aplicado a um denoiser treinado padrão (após conversão) ou a um preditor LOO nativo, melhorando significativamente a qualidade da amostra com baixo custo computacional.

Resultados Experimentais

Os autores validam suas percepções teóricas em tarefas de modelagem de linguagem em grande escala (LM1B, OpenWebText) e em uma tarefa de Sudoku em menor escala.

• Parametrização: Em todos os conjuntos de dados e objetivos de treinamento (Entropia Cruzada e ELBO), modelos parametrizados com o alvo LOO superam consistentemente seus pares de denoiser padrão em termos de perplexidade de validação e fronteiras generativas (Gen-PPL vs. entropia).
• Eficiência de Treinamento: O uso do alvo LOO com perda de entropia cruzada padrão alcança desempenho comparável ao treinamento ELBO de tempo contínuo, mas com uma implementação mais simples.
• Amostragem: A aplicação de amostragem top-p ou por temperatura à representação LOO (em vez do denoiser bruto) melhora a geração. O amostrador preditor-corretor proposto, utilizando a condicional LOO, domina Pareto as fronteiras de amostragem ancestral.
• AUDM vs. MDM: A Difusão Uniforme de Estado Absorvente (AUDM) e sua variante com resampling (ReAUDM) igualam ou superam o desempenho dos MDMs padrão em tarefas de transferência zero-shot e fronteiras generativas, enquanto melhoram a verossimilhança dentro da família de difusão uniforme.
• Sudoku: Na tarefa de Sudoku, as parametrizações LOO alcançam precisão quase perfeita, superando os MDMs.

Significado e Afirmativas

O artigo conclui que a lacuna de desempenho empírico entre a difusão mascarada e a uniforme é impulsionada menos pela escolha das próprias marginais (mascaramento vs. ruído uniforme) e mais pela parametrização e design de amostragem associados a elas.

• Unificação: A reformulação de estado absorvente demonstra que as vantagens estruturais da difusão mascarada (carry-over, posteriores mais simples) podem ser transferidas para a difusão uniforme sem alterar a lei conjunta subjacente.
• Alvo de Otimização: A identificação da posterior LOO como o ótimo verdadeiro para UDMs de plug-in resolve um descompasso teórico e fornece um alvo de treinamento prático e de alto desempenho.
• Impacto Prático: O trabalho oferece melhorias imediatas na inferência sem treinamento (via conversão LOO e amostragem preditor-corretor) e uma nova estratégia de parametrização que melhora consistentemente a fronteira generativa dos UDMs.

Os autores observam que, embora a vantagem do LOO seja empiricamente robusta, uma explicação teórica para sua fronteira generativa superior permanece uma questão em aberto. Além disso, embora o AUDM melhore a verossimilhança, as melhores fronteiras generativas são ainda alcançadas por UDMs parametrizados com LOO, sugerindo mais espaço para melhoria nos modelos de estado absorvente.