Consistent Diffusion Language Models

1. Declaração do Problema

Os Modelos de Difusão de Linguagem (DLMs) oferecem uma alternativa promissora aos modelos autoregressivos (AR), permitindo a geração paralela de tokens e inferência em tempo sublinear. No entanto, um gargalo crítico permanece: gerar texto de alta qualidade geralmente requer centenas de etapas de refinamento, anulando os ganhos de eficiência prometidos pelo paralelismo.

Em domínios contínuos (por exemplo, imagens), os Modelos de Consistência aceleraram com sucesso a difusão ao treinar modelos para serem invariantes ao longo de uma trajetória única e determinística definida pela Equação Diferencial Ordinária de Fluxo de Probabilidade (PF-ODE).

• A Lacuna: Em domínios discretos (como texto), não existe uma PF-ODE de espaço de amostra análoga. Processos de corrupção discreta (por exemplo, mascaramento ou ruído uniforme) não seguem um único caminho determinístico; em vez disso, envolvem transições estocásticas.
• A Consequência: Adaptar diretamente métodos de consistência contínua para dados discretos é mal definido. Soluções existentes frequentemente dependem de pipelines de destilação multiestágio (usando um modelo professor) ou substitutos de relaxação contínua, que são complexos, computacionalmente caros e frequentemente sofrem com colapso de modo ou perda de diversidade.

2. Metodologia: Consistência Discreta Multi-Caminho (MPDC)

Os autores propõem uma mudança fundamental: em vez de buscar um caminho determinístico inexistente, eles aproveitam as pontes posteriores exatas inerentes aos processos de difusão discreta.

Conceito Central: A Ponte Posterior

Para quaisquer dois níveis de ruído $s < t$ em um processo de difusão discreta, existe uma ponte posterior exata $q(x_s | x_t, x_0)$ de forma fechada. Esta ponte define um caminho estocástico válido conectando um estado ruidoso $x_t$ a um estado mais limpo $x_s$, condicionado aos dados verdadeiros $x_0$.

• Diferentemente da PF-ODE contínua, esta ponte é estocástica.
• Crucialmente, essas pontes formam um semigrupo (propriedade de Chapman-Kolmogorov), o que significa que os caminhos podem ser compostos transitivamente.

O Princípio MPDC

Os autores introduzem a Consistência Discreta Multi-Caminho (MPDC), que treina um denoiser para ser invariante de caminho em expectativa.

• Definição: Um modelo $f_\theta(x_t, t)$ é consistente se sua previsão em um estado ruidoso $x_t$ concordar (em expectativa) com a previsão feita após dar um "salto" intermediário para um estado mais limpo $x_s$ via a ponte posterior exata.
• Objetivo: O modelo minimiza a divergência entre sua previsão direta a partir de $x_t$ e a previsão esperada a partir de $x_s$ (onde $x_s$ é amostrado da ponte).
• Formulação Matemática:
  $$\mathcal{L}_{CDLM} = \mathbb{E}_{t, \delta, x_t} \left[ \sum_{i \in \mathcal{M}(x_t)} w(t, \delta) \cdot D\left( f_\theta(x_t, t)_i \parallel \tilde{f}_\theta(x_s, s)_i \right) \right]$$
  Onde $x_s \sim q(x_s | x_t, x_0)$, $\tilde{f}_\theta$ é uma rede alvo (EMA) e $D$ é uma métrica de divergência (tipicamente Divergência Jensen-Shannon).

Escolhas de Design Chave para Estabilidade

Para prevenir o colapso de modo (onde o modelo gera texto repetitivo para satisfazer trivialmente a consistência) e o desvio uniforme, os autores introduzem três mecanismos estabilizadores:

1. Tamanho do Passo como Currículo: Amostragem de tamanhos de passo $\delta$ de uma faixa específica (por exemplo, $1/8$ a $3/8$) e ponderação por $1/\delta$ para equilibrar sinais de aprendizado entre caminhos curtos e longos.
2. Âncora de Difusão de Passo Máximo: Uma pequena fração da perda de treinamento é ancorada aos dados verdadeiros ($s=0$) usando um objetivo de difusão padrão. Isso impede que o modelo convirja para soluções constantes degeneradas.
3. Assimetria de Otimização e Divergência: Uso de um stop-gradient na rede alvo e emprego da Divergência Jensen-Shannon (JSD) em vez de KL Direto/Inverso. A JSD fornece um compromisso estável entre busca de modo e cobertura de modo, evitando a instabilidade da divergência KL ao treinar do zero.

3. Contribuições Chave

1. Um Novo Princípio para Modelagem Generativa Discreta: O artigo estabelece que as pontes posteriores exatas servem como um substituto rigoroso e de forma fechada para a PF-ODE faltante, transformando o obstáculo da "aceleração discreta" em um ativo analítico.
2. Um Framework de Treinamento Unificado: O CDLM fornece um objetivo de estágio único, sem professor. Os autores provam que a difusão mascarada padrão, modelos de consistência contínua, destilação progressiva e métodos recentes de destilação em dois estágios (como SDTT e DUO-DCD) são todos casos especiais, limites ou acoplamentos aproximados do framework CDLM.
3. Desempenho State-of-the-Art: O CDLM alcança resultados superiores tanto na geração de texto incondicional quanto condicional, superando modelos base fortes e igualando ou superando modelos destilados multiestágio complexos sem exigir um professor.

4. Resultados Experimentais

Os autores avaliaram o CDLM (especificamente MCDLM para difusão mascarada e UCDLM para difusão uniforme) contra baselines como MDLM, DUO, SDTT e DUO-DCD no conjunto de dados OpenWebText.

• Geração Incondicional:

  • Eficiência em Poucos Passos: O CDLM estabelece novos resultados state-of-the-art (SOTA) para modelos de estágio único. Supera significativamente os modelos base (MDLM, DUO) em todos os orçamentos de amostragem.
  • Comparação com Destilação: Remarkavelmente, o CDLM de estágio único frequentemente supera modelos destilados multiestágio (SDTT, DUO-DCD) em termos de Perplexidade (PPL) enquanto mantém entropia muito maior (diversidade).
  • Colapso de Modo: Modelos destilados usando amostradores gananciosos (por exemplo, DUO-DCD) mostraram baixa entropia (3,9), indicando colapso de modo severo (caracteres repetitivos). O CDLM manteve entropia saudável (5,0+), provando sua capacidade de preservar a diversidade.
  • Velocidade: O CDLM alcança acelerações de 32x–64x sobre baselines autoregressivos e 64x–128x sobre o MDLM padrão, mantendo qualidade comparável.

• Geração Condicional:

  • O CDLM demonstrou desempenho robusto na recuperação de sentenças originais a partir de entradas com 50% de corrupção em múltiplos conjuntos de dados (OpenWebText, Lambada, Wikitext-103, PTB).
  • Alcançou melhores pontuações BLEU e menor perplexidade do que SDTT e DUO, particularmente na preservação de condições de entrada onde outros modelos falharam.

5. Significado e Impacto

• Unificação Teórica: O artigo fornece uma "lente unificadora" para a modelagem generativa discreta, mostrando que técnicas de aceleração diversas estão fundamentalmente conectadas através do conceito de consistência de caminho.
• Escalabilidade: Ao eliminar a necessidade de destilação multiestágio e modelos professores, o CDLM oferece uma base principiada, escalável e sem professor para modelos generativos discretos rápidos de próxima geração.
• Generalizabilidade: Embora demonstrado em texto, o framework aplica-se a qualquer domínio discreto com uma ponte posterior tratável (por exemplo, sequências biológicas, grafos, síntese de programas).
• Mudança de Paradigma: Reenquadra a difusão discreta rápida não como uma aproximação de EDOs contínuas, mas como o aprendizado de um denoiser independente de caminho sobre pontes estocásticas, resolvendo com sucesso o problema de geração em poucos passos em espaços discretos.

Em conclusão, o CDLM preenche a lacuna entre a eficiência dos modelos de consistência contínua e as realidades dos dados discretos, oferecendo uma solução robusta de estágio único que alcança geração de texto rápida e de alta fidelidade sem a complexidade dos pipelines de destilação.