Resumo Técnico: Revisar, Não Congelar: Treinamento Casado com o Amostrador para Modelos de Linguagem de Difusão com Máscara Autocorretivos

1. Definição do Problema

Modelos de Linguagem de Difusão com Máscara (MDLMs), como o LLaDA, geram texto iterativamente desmascarando sequências, re-prevendo cada posição em cada etapa. Esta arquitetura possui inerentemente uma "capacidade de revisão": mesmo após um token tornar-se visível, o desmascarador pode, teoricamente, atribuir uma previsão diferente para aquela posição em etapas posteriores.

No entanto, as estratégias de decodificação padrão para MDLMs utilizam o desmascaramento monotônico, que "congela" os tokens uma vez que eles são comprometidos. Uma vez que um token é revelado, ele é tratado como o token limpo, e as etapas subsequentes atualizam apenas as posições mascaradas. Isso leva a dois problemas críticos:

1. Erros Precoces Permanentes: Erros cometidos no início do processo de geração (por exemplo, um dígito errado em matemática ou um operador incorreto em código) tornam-se permanentes, impactando severamente o desempenho em tarefas intensivas de raciocínio.
2. O Hiato de Revisão: As abordagens existentes para revisão dependem de mecanismos heurísticos (limiares de confiança, passagens de verificação) ou exigem o re-mascaramento de tokens comprometidos de volta para [MASK] antes de re-predizê-los (ex: ReMDM). Um amostrador principiado que permita a revisão direta de visível-para-visível sem módulos auxiliares ou etapas de re-mascaramento permanece inexplorado.

Além disso, os autores identificam um obstáculo do lado do modelo denominado viés de preservação: mesmo quando o objetivo de treinamento não impõe a preservação do token, modelos treinados com objetivos de máscara padrão tendem a reproduzir seus próprios tokens visíveis incorretamente comprometidos, em vez de corrigi-los quando expostos ao contexto autogerado.

2. Metodologia

O artigo propõe uma solução de duas partes: D3IM (um novo amostrador) e SCOPE (um procedimento de treinamento), projetados para trabalhar em conjunto.

2.1 D3IM: Re-mascaramento de Folha em Branco

O D3IM é um amostrador livre de parâmetros derivado como uma atualização de correção reversa. Ao contrário dos amostradores padrão que preservam tokens comprometidos (definindo o peso de preservação $c_t > 0$), o D3IM define $c_t = 0$.

• Mecanismo: Em cada etapa de desmascaramento, o modelo prevê todas as posições ativas. O D3IM classifica essas previsões por confiança e retém apenas as top-$K$ posições.
• Canais de Revisão:
  • Token-para-Token: Um token previamente comprometido pode ser diretamente sobrescrito por uma nova previsão de maior confiança sem passar pelo estado [MASK].
  • Token-para-Máscara: Um token comprometido de baixa confiança é rebaixado para [MASK] para ser re-previsto em etapas subsequentes.
• Restrição: O próximo estado é determinado inteiramente pelas previsões atuais, em vez de copiar os tokens visíveis anteriores. Isso exige que o ranking de confiança do modelo seja confiável sob contexto autogerado.

.2 SCOPE: Autocondicionado sobre Erros de Predição

Aplicar diretamente o D3IM a um MDLM base degrada o desempenho porque o modelo carece da habilidade de corrigir seus próprios erros de alta confiança (viés de preservação). O SCOPE é um procedimento de pós-treinamento leve projetado para alinhar o comportamento de confiança do modelo com a dinâmica de amostragem do D3IM.

• Processo:
  1. O modelo realiza uma passagem direta (forward pass) em uma entrada parcialmente mascarada (autocondicionamento).
  2. Ele amostra previsões para as posições mascaradas.
  3. A fração top-$\rho$ dessas previsões (por confiança) são "comprometidas" como tokens visíveis, enquanto o restante permanece mascarado. Crucialmente, este processo expõe o modelo aos seus próprios erros de alta confiança.
  4. O modelo é então treinado para recuperar a sequência real (ground-truth) tanto para os tokens comprometidos (potencialmente errôneos) quanto para as posições restantes mascaradas.
• Objetivo: Treinar o modelo para tratar tokens visíveis autogerados como hipóteses, em vez de contexto fixo, reduzindo assim o viés de preservação e tornando o ranking de confiança confiável para a seleção top-$K$ do D3IM.

3. Principais Contribuições

1. Amostrador D3IM: Um amostrador de revisão direta, livre de parâmetros, que permite atualizações de token visível-para-visível sem módulos auxiliares ou re-mascaramento, derivado como uma atualização de correção válida para MDLMs pré-treinados.
2. Identificação do Viés de Preservação: Os autores demonstram que o treinamento padrão de MDLM leva a uma tendência comportamental onde os modelos reproduzem suas próprias previsões visíveis incorretas, mesmo quando o objetivo de treinamento permite a correção.
3. Treinamento SCOPE: Um procedimento de pós-treinamento leve que simula o processo de amostragem do D3IM. Ele reduz efetivamente o viés de preservação, permitindo que o D3IM funcione conforme pretendido.
4. Alinhamento Amostrador-Modelo: O trabalho estabelece que as capacidades de revisão não são genéricas; o objetivo de treinamento deve ser especificamente casado com as regras de decisão do amostrador de inferência (ex: a regra global top-$K$ do D3IM) para ser eficaz.

4. Resultados Experimentais

Experimentos foram conduzidos no LLaDA-8B-Instruct em quatro benchmarks: GSM8K (matemática), MATH-500 (matemática), HumanEval (código) e MBPP (código).

• Ganhos de Desempenho: Com 64 etapas de desmascaramento, a combinação de SCOPE + D3IM supera significativamente o LLaDA-8B original com desmascaramento padrão:
  • GSM8K: +13,0% (68,3% vs. 55,3%)
  • MATH-500: +4,8% (23,6% vs. 18,8%)
  • HumanEval: +15,3% (29,3% vs. 14,0%)
  • MBPP: +10,4% (30,8% vs. 20,4%)
• Escalabilidade: Os ganhos aumentam com o número de etapas de desmascaramento em tarefas de matemática e código (ex: GSM8K atinge 73,8% com 256 etapas).
• Necessidade de Alinhamento:
  • D3IM sem SCOPE: O desempenho cai significativamente (ex: GSM8K cai para 48,0%), confirmando que o modelo não adaptado não consegue decidir confiavelmente quais tokens revisar.
  • SCOPE sem D3IM: Usar modelos treinados com SCOPE com desmascaramento monotônico padrão produz resultados mistos, mostrando que os benefícios do treinamento são específicos para a mecânica de revisão do amostrador D3IM.
• Diagnóstico de Mecanismo:
  • Resolução do Viés de Preservação: Em um teste de estresse onde o modelo compromete um token errado, o modelo original preserva o erro 70,1% das vezes. O SCOPE reduz isso para 17,0% e recupera a sequência real 50,2% das vezes.
  • Análise de Trajetória: Sem o SCOPE, as correções do D3IM frequentemente resultam em "oscilação conteúdo-para-espaço em branco" (alternando entre tokens e quebras de linha). Com o SCOPE, as correções visam erros semânticos específicos (ex: conversões de unidades, correções de dígitos) e estabilizam rapidamente.

5. Significância e Alegações

O artigo argumenta que o gargalo para a autocorreção de inferência em MDLMs não é a capacidade de revisitar tokens, mas a confiabilidade do ranking de confiança sob contexto autogerado.

• Revisão Principiada: O D3IM fornece um método teoricamente fundamentado para a revisão direta de tokens, indo além de limiares heurísticos ou re-mascaramento indireto.
• Alinhamento Treinamento-Inferência: O trabalho destaca que estratégias de inferência avançadas (como o re-ranqueamento global top-$K$) requerem objetivos de treinamento específicos (SCOPE) para funcionar. O pré-treinamento contínuo padrão ou métodos genéricos de autocorreção não se transferem automaticamente para essas novas dinâmicas de amostragem.
• Eficiência: A solução proposta é computacionalmente eficiente, exigindo apenas um ajuste fino leve via LoRA (500 passos, ~20 minutos em uma A100) e sem mudanças arquiteturais ou módulos auxiliares durante a inferência.

Os autores concluem que, embora o D3IM+SCOPE melhore significativamente o raciocínio e a geração de código, ele é limitado a correções ao nível de token e não pode reestruturar trajetórias inteiras de raciocínio se a trajetória inicial for fundamentalmente falha. Trabalhos futuros podem explorar o uso das ricas trajetórias de revisão para modelagem de recompensa de processo ou aprendizado por reforço.