DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling

Resumo Técnico: DiLaDiff: Difusão Latente Aumentada Destilada para Modelagem de Linguagem

Declaração do Problema
Os modelos de linguagem de difusão (DLMs) enfrentam um compromisso fundamental entre a qualidade da amostragem e a taxa de transferência. Embora os modelos de difusão discreta (especificamente Modelos de Linguagem de Difusão Mascarada, ou MDLMs) sejam eficazes para dados categóricos, eles falham intrinsecamente em capturar correlações entre tokens decodificados. Essa limitação surge porque o denoiser aproxima a posterior como um produto de marginais independentes token a token. Consequentemente, gerar múltiplos tokens em paralelo leva a frases incoerentes, forçando uma dependência de geração sequencial ou uma redução no paralelismo que limita severamente a taxa de transferência. Por outro lado, modelos de difusão contínua aplicados a representações de texto (por exemplo, embeddings) oferecem técnicas de amostragem e destilação fundamentadas, mas frequentemente lutam com o compromisso expressividade-decodabilidade, onde representações contextuais são difíceis de decodificar de volta para tokens discretos, ou falham em igualar o desempenho de seus equivalentes discretos.

Metodologia
Os autores propõem o DiLaDiff (Modelo de Linguagem de Difusão Latente Aumentada Destilada), um framework híbrido contínuo-discreto projetado para resolver o problema de correlação de tokens enquanto mantém alta taxa de transferência. A arquitetura consiste em três componentes principais:

1. Autoencoder Latente Aumentado (LaDiff):

   • Codificador: Um codificador baseado em transformer ($E_\phi$) aprende um espaço latente contínuo ($z$) a partir de representações contextuais (derivadas dos estados ocultos do BERT). Este espaço latente é projetado para capturar correlações globais de tokens e informações semânticas.
   • Decodificador: Um decodificador discreto ($x_\theta$) é inicializado a partir de um MDLM pré-treinado. Ele gera tokens de texto condicionados à variável latente $z$.
   • Treinamento: O sistema é treinado usando uma Limite Inferior de Evidência (ELBO) modificado que combina auto-codificação com difusão mascarada. Crucialmente, a distribuição posterior é fatorada sobre os tokens condicionados ao latente $z$. Como $z$ captura as correlações, a suposição de independência condicional se mantém, permitindo a geração paralela precisa de tokens.
   • Regularização: Para evitar que o espaço latente se torne muito esparsos para modelagem generativa, os autores empregam heurísticas de reforço de robustez, incluindo normalização coordenada, mascaramento aleatório de características e latentes, e injeção de ruído gaussiano.

2. Aprendizado de Priori Latente:

   • Como a sequência limpa não está disponível na inferência, um modelo de difusão contínua (Denoiser Latente $z_\psi$) é treinado para aprender a distribuição priori $p(z)$ sobre a saída do codificador. Isso permite que o modelo amostre um latente $z$ de alta qualidade via ODE de fluxo de probabilidade reverso antes da decodificação.

3. Destilação de Consistência (DiLaDiff):

   • Para reduzir ainda mais a sobrecarga computacional da etapa de difusão contínua, os autores aplicam a destilação MeanFlow. Um modelo estudante ($u_\eta$) é treinado para prever a velocidade média da trajetória ODE entre dois passos de tempo ($t$ e $r$), em vez da velocidade instantânea.
   • Esta destilação permite que o latente $z$ seja gerado em poucos passos (por exemplo, 5 passos) enquanto mantém o desempenho próximo ao do modelo professor (LaDiff), que requer muitos passos (por exemplo, 200 passos).

Principais Contribuições

• Arquitetura Híbrida: O artigo introduz uma receita para treinar um autoencoder de texto onde o decodificador é inicializado a partir de uma base de difusão discreta pré-treinada. Isso resulta em um espaço latente controlável e regularizado que captura informações semânticas e correlações de tokens.
• Paralelismo Guiado por Latente: Ao aprender uma priori latente com difusão contínua, o modelo (LaDiff) permite a amostragem de múltiplos tokens em paralelo sem sacrificar a qualidade do texto, pois o latente resolve as dependências de tokens que os modelos discretos ignoram.
• Primeira Destilação ODE Bem-Sucedida para Latentes de Texto: Os autores apresentam o DiLaDiff, a primeira tentativa bem-sucedida de destilar trajetórias ODE para latentes contextuais de texto. Isso reduz o custo de geração latente a níveis negligenciáveis em comparação com a decodificação discreta.

Resultados
Os experimentos foram conduzidos no conjunto de dados OpenWebText com um comprimento de sequência de 1024 tokens.

• Desempenho vs. Base: Mesmo sem destilação, o LaDiff supera a base de difusão mascarada (MDLM). Na geração incondicional, o LaDiff alcança uma melhoria absoluta de 28,5 na Perplexidade Generativa (GenPPL) e uma melhoria absoluta de 0,08 no MAUVE em comparação com o MDLM com 1024 passos discretos.
• Taxa de Transferência: O LaDiff acelera a inferência por um fator de 7× com um tamanho de lote de 32 em comparação com a configuração ótima do MDLM.
• Eficiência de Destilação: O DiLaDiff alcança desempenho próximo ao de seu professor LaDiff usando apenas 5 passos de difusão latente (em comparação com 200 para o professor). A sobrecarga computacional do componente de difusão latente cai para aproximadamente 5% do tempo total de parede, tornando a geração latente negligenciável em comparação com o custo da decodificação discreta.
• Geração em Poucos Passos: O DiLaDiff rivaliza com métodos de poucos passos mais avançados (como MDLM+SDTT e DUO+DCD) sem destilar explicitamente o decodificador discreto, demonstrando a eficácia do aumento latente.
• Diversidade Semântica: A análise mostra que o espaço latente preserva informações semânticas (verificado via BERTScore em latentes ruidosos) e permite geração diversificada, conforme evidenciado pela capacidade do decodificador de manter a entropia mesmo em temperaturas mais baixas.

Significado e Alegações
O artigo afirma que o DiLaDiff conecta com sucesso a lacuna entre a expressividade das representações contínuas e a robustez da decodificação discreta. Ao capturar correlações globais em um espaço latente, o modelo supera a limitação intrínseca dos modelos de difusão discreta em relação à geração paralela de tokens. Os autores enfatizam que sua abordagem permite a geração de texto de alta qualidade com latência de inferência significativamente reduzida.

O artigo permanece modesto quanto às suas limitações:

• O DiLaDiff ainda não iguala totalmente o desempenho de seu modelo professor (LaDiff) no regime de poucos passos, sugerindo que mais trabalho sobre destilação é necessário.
• O próprio decodificador discreto não é destilado; os autores observam que combinar sua destilação latente com a destilação de decodificador discreto é uma extensão futura natural para alcançar desempenho de poucos passos verdadeiramente state-of-the-art.
• As estratégias de regularização e os designs de agendamento de difusão são atualmente heurísticos, embora os autores apresentem uma tentativa mais fundamentada nos apêndices.

Em resumo, o DiLaDiff demonstra que uma abordagem híbrida contínuo-discreta, aumentada por destilação latente, pode acelerar significativamente a inferência de modelagem de linguagem enquanto mantém ou melhora a qualidade de geração, resolvendo efetivamente o compromisso taxa de transferência-qualidade inerente aos atuais modelos de linguagem de difusão.