Resumo Técnico: Localização Estocástica Discreta para Geração Não-Autoregressiva

1. Declaração do Problema

Os modelos de difusão contínuos tornaram-se um padrão para geração de alta dimensão em imagens, vídeo e áudio. No entanto, sua aplicação à geração de linguagem via difusão contínua sobre incorporações de tokens tem consistentemente apresentado desempenho inferior em comparação com os Modelos de Difusão Discreta Mascarada (MDMs). Os modelos de linguagem de difusão contínua existentes geralmente ficam atrás dos MDMs em fidelidade distribucional (medida por métricas como MAUVE).

Os autores argumentam que o gargalo não é a continuidade em si, mas a representação utilizada. Abordagens padrão (Preservação de Variância ou Explosão de Variância) representam o limite de zero informação como ruído gaussiano isotrópico, em vez de um estado de token semelhante a uma máscara. Além disso, nessas estruturas, o denoiser deve receber um rótulo de passo de tempo ou nível de ruído para calibrar a quantidade de evidência na incorporação ruidosa. Isso torna os denoisers altamente sensíveis à condicionamento temporal, exigindo tratamento distinto para estados mascarados versus incertos. Consequentemente, os modelos de difusão contínua falharam em aproveitar a supervisão mais rica fornecida pelo contínuo de estados de Razão Sinal-Ruído (SNR) finitos.

2. Metodologia: Localização Estocástica Discreta (DSL)

O artigo introduz a Localização Estocástica Discreta (DSL), uma estrutura de estado contínuo projetada para tornar o processo de remoção de ruído invariante ao SNR nominal.

Escolhas de Design Central

1. Incorporações de Token na Esfera Unitária: A DSL restringe todas as incorporações de token limpas a residir na esfera unitária ($\|x_i\|_2 = 1$). Esta é uma restrição estrutural na geometria, não um regularizador de perda.
2. Canal de Localização Estocástica: O processo de corrupção é definido como $z_t = t x + \sqrt{t} \epsilon$, onde $t$ representa o SNR. Sob a restrição da esfera unitária, o termo envolvendo a norma de $x$ no posterior de Bayes torna-se constante e cancela-se durante a normalização.
3. Denoiser Invariante ao SNR: Um resultado teórico chave é que, sob esta geometria, o denoiser ótimo de Bayes (a média posterior) depende apenas do estado de token ruidoso $z$, e não do SNR nominal $t$.
   $$\hat{x}(z, t) = \hat{x}(z) = \frac{\mathbb{E}_P[x e^{x \cdot z}]}{\mathbb{E}_P[e^{x \cdot z}]}$$
   Isso permite que uma única rede, agnóstica ao tempo, sirva como denoiser para todos os regimes de SNR (de totalmente mascarado a quase limpo).

Treinamento e Arquitetura

• Treinamento de Suporte Misto: Para aproveitar tanto a verossimilhança do integral de caminho contínuo quanto a verossimilhança do ponto final da difusão mascarada, a DSL emprega um objetivo de treinamento de suporte misto. Ela amostra SNRs de dois ramos:
  • Caminho Contínuo: SNRs extraídos de uma distribuição log-normal (estados de SNR finito).
  • Ponto Final ROAR (Autoregressivo de Ordem Aleatória): SNRs extraídos de valores extremos (próximo de 0 para mascarado, próximo de $\infty$ para revelado).
    Ambos os ramos são treinados usando uma única perda de entropia cruzada no posterior do token, que atua como um substituto para o Erro Quadrático Médio (MSE) do integral de caminho, evitando patologias de otimização no espaço de incorporação (por exemplo, colapso de incorporação).
• Conversor de Visão Posterior: Como o estado ruidoso $z$ vive em um espaço de incorporação contínuo, mas a base (Transformer/DiT) espera entradas do tipo token, um conversor mapeia $z$ para uma representação de "mistura de tokens". Isso mapeia o estado ruidoso para uma combinação limitada de incorporações de norma unitária, preservando a geometria posterior (a direção controla a preferência do token, a norma controla a confiança) enquanto garante atenção autoestável.

Flexibilidade de Inferência

Como o denoiser é agnóstico ao tempo, a DSL suporta caminhos de SNR arbitrários por token dentro de uma única execução de geração:

• Refinamento Mascarado: Remascaramento iterativo padrão (ReMDM).
• Autoregressivo de Ordem Aleatória (ROAR): Revelar tokens um por um em ordem aleatória sem remascaramento.
• Amostragem Híbrida: Um estágio de desruído contínuo seguido por um estágio de refinamento discreto mascarado.

3. Contribuições Principais

1. Uma Estrutura de Estado Contínuo para Geração Discreta: A DSL introduz um canal de localização estocástica sobre incorporações na esfera unitária onde o denoiser ótimo de Bayes é uma função apenas do estado ruidoso. Isso permite que uma única rede agnóstica ao tempo suporte toda uma família de caminhos de SNR por token.
2. Uma Receita Prática de Treinamento: Os autores propõem uma estratégia de amostragem de SNR de suporte misto que treina o mesmo denoiser tanto em estados de ponto final (máscara/revelação) quanto em estados de SNR finito. Eles utilizam um conversor de mistura de tokens para alinhar a visão posterior contínua com arquiteturas Transformer padrão.
3. Evidência Através de Caminhos: Um único checkpoint DSL demonstra geração em poucos passos aprimorada no OpenWebText sob refinamento mascarado, suporta amostragem AR de ordem aleatória sem retreinamento e alcança verossimilhança competitiva no Text8.

4. Resultados Experimentais

Os autores avaliaram a DSL ajustando finamente um checkpoint pré-treinado de MDLM no OpenWebText (OWT) e no Text8.

• Geração OpenWebText (MAUVE):
  • Sem Refinamento: Mesmo sem remascaramento iterativo, o checkpoint ajustado com DSL superou significativamente o MDLM de base. Em $T=1024$, o MAUVE melhorou de 0,042 (MDLM) para 0,495 (DSL-FT + MDLM) e 0,551 (DSL-FT + ROAR).
  • Refinamento Iterativo (ReMDM): Quando combinado com decodificadores da família ReMDM, a DSL alcançou resultados state-of-the-art entre bases sem referência. A melhor configuração (loop ReMDM com $\eta_{cap}$ principista) atingiu um MAUVE de 0,722 em $T=1024$, superando métodos anteriores como PRISM-loop (0,527) e ADLM (0,699).
  • Eficiência em Poucos Passos: A DSL mostrou desempenho forte com orçamentos baixos de passos. Em $T=128$, DSL-FT + loop ReMDM alcançou 0,639 MAUVE, comparado a 0,015 para o MDLM de base.
• Amostragem Híbrida: Um amostrador híbrido (desruído contínuo seguido de refinamento discreto) alcançou um MAUVE de 0,662 em apenas 48 passos totais, excedendo o resultado do MDLM sem refinamento em $T=1024$.
• Verossimilhança Text8: A DSL alcançou Bits por Caractere (BPC) de $\le$ 1,45, aproximando-se da faixa dos modelos de difusão discreta mascarada (MDLM $\le$ 1,40, MD4 $\le$ 1,37) e melhorando sobre bases de difusão contínua anteriores (Plaid $\le$ 1,48).

5. Significado e Alegações

O artigo afirma que a vantagem primária da DSL é a unificação dos caminhos de inferência. Ao parametrizar a difusão contínua como desruído de estado posterior sobre tokens ruidosos (em vez de desruído indexado por passo de tempo), a DSL permite:

• Revelação de ponto final, refinamento mascarado e desruído contínuo/híbrido serem atendidos por um único estimador posterior.
• Compatibilidade de checkpoint: Um único modelo pode ser consultado ao longo de múltiplos caminhos de SNR (por exemplo, AR de ordem aleatória, refinamento mascarado, híbrido) sem retreinamento ou destilação.
• Fidelidade Distribucional Aprimorada: Os estados de SNR finito contínuos fornecem supervisão mais rica, levando a pontuações MAUVE substancialmente melhores em todos os orçamentos de passos em comparação com difusão contínua padrão ou bases de difusão mascarada.

Os autores concluem que o ingrediente faltante na difusão contínua para linguagem não era a continuidade, mas a escolha da representação. Ao adotar incorporações na esfera unitária e localização estocástica, eles fecham a lacuna entre difusão contínua e discreta, permitindo que um único modelo suporte diversas estratégias de geração com alta eficiência e qualidade.

Limitações: A avaliação está atualmente limitada ao OpenWebText e Text8. Os amostradores relatados (família ReMDM, híbrido) não são políticas de decodificação exaustivamente otimizadas, mas servem para demonstrar a flexibilidade do posterior subjacente da DSL. Direções de trabalho futuro incluem a otimização de cronogramas de remascaramento, regras de transferência híbrida e exploração de geração condicional.