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

Declaração do Problema

A difusão contínua tornou-se um framework padrão para geração de alta dimensão em imagem, vídeo e áudio, mas consistentemente ficou atrás dos modelos de difusão discreta mascarada (MDMs) na geração de linguagem. Os modelos de linguagem de difusão contínua existentes tipicamente operam sobre embeddings de tokens ruidosos usando ruído Gaussiano Variance-Preserving (VP) ou Variance-Exploding (VE). Os autores argumentam que o gargalo não é a continuidade em si, mas a escolha da representação.

Na difusão contínua padrão, o processo de remoção de ruído depende de um regime de ruído indexado por passo de tempo. O modelo deve receber um rótulo de passo de tempo ou nível de ruído para calibrar quanto evidência o embedding ruidoso contém. Consequentemente, os denoisers nessa representação são altamente sensíveis à condicionamento temporal para distinguir entre estados mascarados (zero-informação) e incertos (SNR finita). Essa dependência de condicionamento temporal explícito impede que uma única rede suporte naturalmente caminhos de inferência diversos, como amostragem autoregressiva em ordem aleatória ou decodificação híbrida contínua-discreta, sem modelos específicos para amostradores separados ou retreinamento.

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

O artigo introduce Localização Estocástica Discreta (DSL), um framework de estado contínuo projetado para tornar o denoiser Bayes-ótimo invariante à relação sinal-ruído (SNR) nominal.

1. Geometria de Tokens na Esfera Unitária

A inovação central é restringir todos os embeddings de tokens limpos a residir na esfera unitária ($\|x_i\|_2 = 1$). Sob um canal de localização estocástica onde a observação ruidosa é $z_t = t x + \sqrt{t}\epsilon$, a fatorização do posterior Bayesiano inclui um termo proporcional a $t \|x\|_2^2$. Como os embeddings são de norma unitária, $\|x\|_2^2$ é constante para qualquer frase de um determinado comprimento. Este termo cancela-se durante a normalização, tornando o posterior — e, portanto, o denoiser ótimo — independente da SNR nominal $t$.

O denoiser ótimo torna-se uma função do estado ruidoso $z$ apenas:
$$\hat{x}(z) = \mathbb{E}[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 toda a família de caminhos de SNR por token.

2. Caminhos de SNR por Token

Como o denoiser não requer um passo de tempo global, a DSL permite que cada posição de token evolua sob seu próprio caminho de SNR arbitrário $\gamma_i(t)$. Isso unifica várias estratégias de inferência sob um único modelo:

• Autoregressiva em Ordem Aleatória (ROAR): Revelar tokens um por um corresponde a impulsionar $\gamma_i$ específicos de 0 a $\infty$ enquanto mantém os outros em 0.
• Refinamento Mascarado: Remascarar tokens corresponde a redefinir seus $\gamma_i$ para 0.
• Difusão Contínua: Aumentar conjuntamente todos os $\gamma_i$ corresponde à remoção de ruído contínua padrão.

3. Objetivo de Treinamento e Arquitetura

Para treinar um único modelo para esses caminhos diversos, a DSL emprega uma estratégia de amostragem de SNR de suporte misto:

• Ramo Contínuo: Amostra SNRs de uma distribuição log-normal para cobrir estados de SNR finita.
• Ramo de Extremidade: Amostra de uma distribuição estilo ROAR concentrada perto das extremidades mascarada ($\gamma=0$) e limpa ($\gamma=\infty$).

Ambos os ramos são treinados usando uma única perda de entropia cruzada (CE) em nível de token. Os autores justificam isso mostrando que minimizar a CE no posterior categórico de tokens é um substituto para o NLL exato de Erro Quadrático Médio (MSE) de integral de caminho, evitando patologias associadas à regressão direta no espaço de embeddings (como colapso de embeddings ou memorização).

A arquitetura inclui um conversor de visão de posterior que mapeia o embedding ruidoso $z_i$ para uma representação de mistura de tokens limitada $m_i$. Isso garante que a entrada para o backbone Transformer permaneça estável e interpretável como uma distribuição de probabilidade sobre o vocabulário, independentemente do nível de SNR.

Contribuições Principais

1. Um Framework de Estado Contínuo para Geração Discreta: A DSL introduz um canal de localização estocástica sobre embeddings de esfera unitária onde o denoiser Bayes-ótimo é uma função do estado ruidoso apenas, eliminando a necessidade de condicionamento temporal explícito.
2. Uma Receita Prática: O artigo propõe uma estratégia de amostragem de SNR de suporte misto que treina um único denoiser em estados de extremidade (máscara/revelação) e de SNR finita, acoplado a um conversor que faz a ponte entre a geometria do posterior contínuo e backbones Transformer padrão.
3. Evidência em Diversos Caminhos: Os autores demonstram que um único checkpoint DSL pode suportar refinamento mascarado, amostragem autoregressiva em ordem aleatória e amostragem híbrida contínua-então-discreta sem retreinamento.

Resultados Experimentais

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

• Geração OpenWebText: A DSL melhora substancialmente a fidelidade distribucional (medida por MAUVE) em todos os orçamentos de passos ($T=128$ a $T=1024$) em comparação com o MDLM de base.
  • Com um decodificador de refinamento mascarado (ReMDM), a DSL atinge um MAUVE de 0,722 em $T=1024$, superando métodos anteriores de última geração como PRISM-loop (0,527) e ADLM (0,699).
  • Mesmo sem refinamento iterativo (usando um ROAR ingênuo ou decodificador MDLM vanilla), o checkpoint ajustado finamente pela DSL supera significativamente a base, indicando que a melhoria decorre do posterior aprendido e não apenas do cronograma de amostragem.
• Amostragem Híbrida: Um amostrador híbrido que realiza uma curta etapa de remoção de ruído contínua seguida por refinamento mascarado discreto atinge geração coerente em tão poucos quanto 48 passos totais (MAUVE 0,662), superando bases sem refinamento em contagens de passos muito maiores.
• Verossimilhança Text8: No benchmark Text8, a DSL atinge Bits por Caractere (BPC) de $\le 1,45$. Embora este seja um limite superior (ELBO) e não diretamente comparável ao NLL exato de modelos autoregressivos, aproxima-se da faixa dos modelos de difusão discreta mascarada (MDLM $\le 1,40$) e melhora outras bases de difusão contínua (Plaid $\le 1,48$).

Significado e Alegações

O artigo alega que a lacuna entre difusão contínua e discreta em linguagem não se deve à natureza da continuidade, mas à representação do regime de ruído. Ao parametrizar a corrupção de modo que o estado ruidoso sozinho determine o posterior do token limpo, a DSL unifica o comportamento de "revelação de extremidade" da difusão mascarada com as capacidades de "refinamento contínuo" dos modelos de difusão.

A alegação central é que revelação de extremidade, refinamento mascarado e remoção de ruído contínuo/híbrido podem ser atendidos por um único estimador de posterior em vez de exigir modelos separados e específicos para amostradores. Isso sugere um novo paradigma para geração não-autoregressiva onde o modelo aprende um mapa de posterior agnóstico ao tempo, e o caminho de inferência (AR, mascarado ou híbrido) é determinado dinamicamente no momento da amostragem.

Os autores permanecem modestos quanto às limitações, notando que sua avaliação é limitada a OWT e Text8, e que os amostradores relatados (família ReMDM e híbrido) não são políticas de decodificação exaustivamente otimizadas, mas sim provas de conceito demonstrando a flexibilidade do posterior DSL. Direções de trabalho futuro incluem cronogramas de remascaramento otimizados, modelos maiores e geração condicional.