DSL-LLaDA: Scaling Continuous Denoising to 8B Masked Diffusion LMs

Resumo Técnico: DSL-LLaDA: Escalando a Denoising Contínua para 8B de Modelos de Linguagem de Difusão Mascarada

Definição do Problema
Modelos de Linguagem de Difusão Discreta (DLMs), como o LLaDA, geram texto via decodificação paralela iterativa. Embora esses modelos tenham escalado para bilhões de parâmetros, a qualidade de sua inferência degrada-se drasticamente sob orçamentos de passos baixos (geração de poucos passos). O problema central é um compromisso entre o comprimento da saída e a qualidade: com um número fixo de passagens para frente (forward passes), o desmascaramento iterativo padrão produz ou saídas curtas e de alta qualidade, ou textos longos, repetitivos e incoerentes. Essa falha decorre da fatoração da transição reversa, onde cada posição de token é decodificada independentemente de seu contexto, baseando-se na distribuição marginal condicionada, negligenciando correlações entre tokens. Quando os orçamentos de passos são baixos, cada passo deve comprometer muitas posições simultaneamente, amplificando esses erros e levando à terminação prematura ou loops de repetição.

Embora a denoising contínua no espaço de embeddings ofereça uma solução teórica ao evoluir todas as posições conjuntamente para evitar a fatoração no espaço de tokens, construir tal modelo do zero em escala permanece um problema de design aberto devido às complexidades da geometria de embeddings, agendamento de ruído (noise scheduling) e objetivos de treinamento.

Metodologia
Os autores propõem o DSL-LLaDA, um método para adaptar levemente um DLM pré-treinado existente (especificamente o LLaDA-8B-Instruct) para suportar a denoising contínua no espaço de embeddings, em vez de treinar um modelo contínuo do zero.

1. Adaptação por Localização Estocástica Discreta (DSL):

   • O método substitui a entrada binária de "limpo ou mascarado" dos DLMs padrão por estados contínuos ruidosos. Para um token limpo $x_i$, uma Razão Sinal-Ruído (SNR) $\gamma_i$ é amostrada, e um estado ruidoso $z_i$ é construído como uma mistura do embedding do token e ruído Gaussiano.
   • Um conversor softmax mapeia esses estados DSL contínuos de volta para o espaço de embedding de entrada do backbone. Este conversor produz uma "máscara suave" onde o baixo SNR corresponde a um estado de máscara, o SNR médio a tokens aleatórios não confiáveis, e o alto SNR a embeddings de tokens claros.
   • O backbone passa por um pré-treinamento contínuo de apenas 1.000 passos no FineWeb-Edu. O objetivo de treinamento permanece a entropia cruzada de token padrão, mas os estados de entrada são agora misturas contínuas em vez de máscaras binárias. Crucialmente, o conversor permite que o modelo processe esses inputs sem mudanças arquiteturais, pois a invariância de SNR do denoiser ótimo remove a necessidade de entradas explícitas de timestep.

2. Inferência Contínua (Sampler SDE):

   • Durante a inferência, o modelo mantém um estado contínuo para todas as posições. Ele utiliza um sampler baseado em Equações Diferenciais Estocásticas (SDE) (passos de Heun) para evoluir o estado de baixo SNR (alta incerteza) para alto SNR.
   • O modelo prevê uma distribuição de probabilidade sobre os tokens, que é convertida em uma média ponderada de embeddings de tokens (usando os top-$M$ tokens mais prováveis). Esta estimativa de embedding atualiza o estado contínuo.
   • O compromisso de token rígido (hard token commitment) é adiado até o passo final. Isso permite que o modelo refine conjuntamente todas as posições, preservando as correlações entre tokens durante todo o processo de geração.

3. Geração Condicional:

   • Para tarefas como sumarização, os tokens de contexto são fixados (clamped) em seus embeddings limpos, enquanto a região de resposta passa pela denoising contínua. O backbone atende bidirecionalmente, condicionando a denoising da resposta ao contexto limpo.

Principais Contribuições

1. Conversão Leve: Demonstrou que um DLM mascarado de 8B pré-treinado pode ser convertido para suportar a denoising contínua no espaço de embeddings via uma adaptação mínima de 1.000 passos usando DSL, evitando a necessidade de treinar um modelo contínuo do zero.
2. Evitação de Modos de Falha de Baixo NFE: Mostrou que a inferência SDE contínua evita efetivamente a terminação prematura e os loops de repetição característicos do desmascaramento iterativo em orçamentos de passos baixos (NFE $\le$ 16).
3. Robustez de Estado Ruidoso Seletivo: Identificou que a adaptação induz uma robustez onde o modelo pode corrigir tokens corrompidos enquanto preserva os tokens limpos (correção seletiva), um comportamento ausente em baselines discretos de mesmo computo.

Resultados

• Sumarização Zero-Shot: Em quatro benchmarks (XSum, CNN/DailyMail, arXiv, PubMed), o DSL-LLaDA-SDE alcançou as melhores pontuações ROUGE-1 em orçamentos de passos baixos ($\le$ 16 passagens para frente) sem heurísticas de controle de comprimento. Em contraste, os baselines discretos ou truncavam as saídas precocemente ou sofriam de altas taxas de repetição.
• Repetição e Comprimento: Na geração aberta, o DSL-LLaDA-SDE manteve as taxas de repetição abaixo de 9% em todos os orçamentos de passos, enquanto os baselines discretos excederam 60% de repetição em baixo NFE ou truncaram para comprimentos muito curtos.
• Correção de Erros: O modelo adaptado demonstrou "correção seletiva", recuperando tokens corrompidos enquanto preservava mais de 98% dos tokens limpos. Em comparação, baselines discretos com corrupção de token aleatório (XDLM) danificaram uma parte significativa dos tokens limpos, e a difusão mascarada padrão (MDM-CPT) mostrou capacidade de correção negligenciável.
• Experimentos de Controle: Experimentos de controle usando difusão mascarada padrão com o mesmo computo (MDM-CPT) e corrupção de token aleatório (XDLM) falharam em replicar os benefícios da denoising contínua ou a robustez seletiva, confirmando que as melhorias derivam da exposição ao ruído contínuo e não apenas de computo adicional.

Significância e Alegações
O artigo afirma que o DSL-LLaDA fornece um procedimento de conversão de propósito geral para DLMs mascarados, permitindo que eles aproveitem os benefícios da denoising contínua (evolução conjunta de tokens, adiamento de compromisso) com overhead de treinamento mínimo. Os resultados sugerem que a adaptação para ruído contínuo pode resolver o compromisso entre comprimento e qualidade inerente ao desmascaramento iterativo em orçamentos de passos baixos.

Os autores mantêm a modéstia quanto às implicações mais amplas:

• A avaliação é focada principalmente em geração ancorada (sumarização e seguimento de instruções). A escrita criativa aberta ainda mostra taxas mais altas de degeneração (24% vs. $\le$4% em sumarização).
• O método é validado atualmente apenas na escala LLaDA-8B; o desempenho em outras arquiteturas ou escalas maiores não foi testado.
• Embora o SDE supere os métodos discretos em baixo NFE, o LLaDA vanilla iguala ou excede o desempenho do SDE em tarefas de longa referência em alto NFE (onde a terminação prematura não é um problema).
• Os autores não alegam paridade com modelos autorregressivos (AR) de mesma escala, observando que modelos AR ainda lideram em tarefas de longa referência com altos orçamentos de passos.
• Avaliação humana é identificada como trabalho futuro.