Who Should Lead Decoding Now? Tracking Reliable Trajectories for Ensembling Masked Diffusion Language Models

Resumo Técnico: Quem Deve Liderar a Decodificação Agora? Rastreando Trajetórias Confiáveis para o Ensamble de Modelos de Linguagem de Difusão Mascarada

Definição do Problema
Os Modelos de Linguagem de Difusão Mascarada (MDLMs) emergiram como um paradigma distinto para a geração de sequências, utilizando o remascaramento iterativo para desnortear sequências em paralelo, em vez de forma autorregressiva. À medida que a família MDLM se diversifica com modelos exibindo diferentes forças e dinâmicas de decodificação, surge um desafio crítico: como orquestrar ou fundir efetivamente o conhecimento de MDLMs heterogêneos. As estratégias tradicionais de ensemble para modelos autorregressivos — como a média das probabilidades do próximo token ou o roteamento para o modelo mais confiante — não são diretamente aplicáveis aos MDLMs. Nos MDLMs, as sequências são geradas em uma ordem flexível e não da esquerda para a direita, o que significa que diferentes modelos podem operar em diferentes sequências parcialmente desnorteadas em cada etapa. Essa falta de uma posição comum de "próximo token" torna as definições padrão de ensemble difíceis de aplicar, deixando o ensemble de MDLM amplamente subexplorado.

Metodologia: TIE (Ensemble Iterativo Baseado em Trajetória)
Para abordar isso, os autores propõem o TIE, um framework de fusão de conhecimento projetado especificamente para MDLMs. O TIE opera sobre duas observações principais derivadas da análise das dinâmicas de decodificação:

1. Estabilidade de Trajetórias Corretas: Amostras respondidas corretamente exibem dinâmicas mais estáveis e confiantes sobre tokens relevantes para a resposta (menor entropia, maior probabilidade top-1 e menos mudanças de tokens) em comparação com amostras incorretas, mesmo enquanto os tokens ainda estão mascarados.
2. Corrigibilidade via Injeção: Modelos que geram respostas incorretas podem frequentemente ser direcionados para respostas corretas se receberem estados intermediários parcialmente desnorteadas promissores de modelos que seguem trajetórias mais confiáveis.

Com base nessas percepções, o TIE implementa um ciclo recorrente de três etapas durante o processo de geração:

1. Geração de Trajetória: Cada modelo constituinte ($M$) decodifica independentemente por um número fixo de etapas (o intervalo de ensemble, $n$), produzindo uma trajetória de decodificação parcial.
2. Avaliação de Trajetória: O framework avalia a confiabilidade de cada trajetória pontuando as dinâmicas de confiança sobre posições de tokens de resposta (tokens que aparecem após um sufixo específico como "Resposta:"). A pontuação utiliza quatro métricas:
   • Contagem de Mudança de Token Negativa: Rastreia o histórico de mudanças de token top-1 sobre posições de resposta mascaradas. Uma contagem menor indica maior estabilidade.
   • Probabilidade Top-1: A probabilidade do token mais provável.
   • Entropia Negativa: Mede a certeza da distribuição.
   • Margem de Probabilidade: A diferença entre as probabilidades top-1 e top-2.
   • Pontuação Cross-Model: Para métricas baseadas em logits, o framework processa cada trajetória através de todos os modelos constituintes e calcula a média das pontuações. Isso mitiga as diferenças de calibração entre os modelos, garantindo que a trajetória selecionada seja consistentemente confiante em todo o ensemble, em vez de ser apenas favorecida pelo seu modelo de origem.
3. Relé de Trajetória: A trajetória com a maior pontuação agregada é selecionada. Esta trajetória "vencedora" é transmitida a todos os modelos constituintes, substituindo seus estados parciais atuais. Os modelos então retomam a decodificação independente a partir deste estado intermediário compartilhado e mais confiável.

Este ciclo se repete até que a geração termine. Finalmente, o sistema seleciona a resposta final de entre os $M$ candidatos com base na trajetória com a dinâmica de tokens de resposta mais estável (menor contagem de mudança de token normalizada).

Principais Contribuições

• Nova Análise de Decodificação: O artigo identifica que os tokens relevantes para a resposta em gerações corretas de MDLM seguem trajetórias de confiança estáveis, fornecendo um sinal para identificar caminhos confiáveis antes que a resposta completa seja gerada.
• Framework de Ensemble Específico para MDLM: O TIE introduz um mecanismo para fundir conhecimento ao transmitir iterativamente trajetórias parciais, permitindo que os modelos corrijam seus desvios durante o processo de desnortamento, em vez de esperar pelos outputs finais.
• Mecanismo de Pontuação Cross-Model: Os autores demonstram a necessidade de avaliar trajetórias através de todos os modelos constituintes para considerar as disparidades de calibração, evitando a seleção tendenciosa para modelos excessivamente confiantes.

Resultados Experimentais
Os autores avaliaram o TIE em oito benchmarks em quatro domínios: raciocínio geral (MMLU, ARC-Challenge, Winogrande), raciocínio matemático (GSM8K, MATH500), codificação (HumanEval, MBPP) e planejamento (Countdown).

• Ganhos de Desempenho: O TIE superou consistentemente os MDLMs individuais (LLaDA-1.5, Dream-7B, DreamCoder-7B, DiffuCoder-7B) e os baselines de ensemble pós-geração. Por exemplo, no GSM8K, o ensemble de LLaDA e Dream via TIE melhorou a precisão de 78,39% (Dream) para 83,47%.
• Métricas de Pontuação: Entre as funções de pontuação, a Contagem de Mudança de Token (especificamente a versão normalizada) provou ser a mais robusta, alcançando os melhores resultados em quatro de oito benchmarks.
• Sinergia de Modelos: O framework é mais eficaz quando os modelos constituintes possuem capacidades individuais comparáveis e fortes. Quando as lacunas de desempenho são grandes (ex: >15%), o ensemble torna-se menos eficaz, pois os modelos mais fracos introduzem sinais ruidosos.
• Liderança Dinâmica: A análise revela que o modelo que produz a trajetória de maior pontuação muda frequentemente ao longo do processo de decodificação (taxas de mudança de modelo de 12–44% dependendo da tarefa). Isso confirma que diferentes modelos contribuem com forças complementares em diferentes estágios da geração.
• Compatibilidade: O TIE permanece eficaz quando combinado com técnicas padrão de aceleração de MDLM, como thresholding e unmasking top-$k$.

Significância e Alegações
O artigo afirma que o TIE oferece uma abordagem prática para o problema pouco explorado de ensemble de MDLM sem exigir treinamento adicional. Ao alavancar as dinâmicas de confiança sobre os tokens relevantes para a resposta, o TIE permite que modelos heterogêneos refinem colaborativamente as respostas, permitindo que modelos mais fracos se recuperem de trajetórias subótimas. Os autores posicionam este trabalho como um passo significativo para uma orquestração mais eficaz de modelos de linguagem baseados em difusão, sugerindo que a inferência colaborativa guiada pela estabilidade da trajetória pode aumentar significativamente as capacidades de raciocínio em diversas tarefas. O artigo nota modestamente limitações, especificamente que a eficácia diminui com grandes lacunas de desempenho entre os modelos e que o estudo atual é limitado a ensembles de dois modelos.