Resumo Técnico: Representações de Retransmissão Aprendidas para Modelos de Difusão Discreta Preditivos

1. Declaração do Problema: O "Reset Duro" em Modelos de Difusão Mascaramentada

Modelos de Difusão Mascaramentada (MDMs) geram sequências discretas por meio de um processo iterativo de remoção de ruído. Começando a partir de uma sequência totalmente mascarada, o modelo desmascara iterativamente uma fração de posições. No entanto, uma ineficiência crítica existe na inferência padrão de MDMs: após cada passagem direta, o modelo descarta os ricos cálculos contínuos internos (estados ocultos) gerados para posições mascaradas. O próximo passo começa exclusivamente a partir dos novos tokens discretos comprometidos.

Os autores denominam isso o problema do reset duro. Embora os MDMs realizem múltiplas passagens diretas por geração, a incapacidade de reutilizar o cálculo contínuo intermediário impede que o modelo amortize efetivamente o raciocínio entre os passos. Isso contrasta com o cálculo recorrente, onde o desenrolar de um modelo de parâmetros fixos ao longo dos passos demonstrou expandir a classe de funções que o modelo pode aproximar, auxiliando o desempenho em tarefas de raciocínio complexo. A questão central abordada é como permitir que os MDMs suportem cálculo recorrente que carregue informações mais ricas através dos passos sequenciais de desmascaramento, sem descartar estados latentes intermediários.

2. Metodologia: Representações de Retransmissão Aprendidas (Relay)

A solução proposta, Relay, introduz um estado latente contínuo e diferenciável que é passado adiante através dos passos de decodificação, permitindo que o modelo seja "preditivo".

Mecanismo Central

• Trajetória de Estado Aumentada: Em cada passo de inferência $k$, o modelo mantém um estado aumentado $s_k = (x_{t_k}, h_k)$, onde $x_{t_k}$ é a sequência parcialmente mascarada atual e $h_k$ é um estado de retransmissão aprendido (vetor latente contínuo).
• Arquitetura: O modelo consiste em uma base $f_\theta$, embeddings de tokens $Emb_\theta$, uma cabeça de desembedding $UnEmb_\theta$ e um módulo de retransmissão específico $R_\theta$.
  • O próximo estado oculto é calculado como: $h_{k+1} = f_\theta(Emb_\theta(x_{t_k}) + R_\theta(h_k))$.
  • O estado de retransmissão $h_k$ é projetado e adicionado aos tokens embutidos antes de entrar na base.
  • Os logits de saída $\ell_k$ são derivados de $h_{k+1}$ via $UnEmb_\theta$.
• Treinamento via BPTT Truncado: Ao contrário dos MDMs padrão treinados com passos independentes, o Relay é treinado usando Backpropagation Through Time (BPTT) Truncado.
  • O modelo realiza uma simulação de $K$ passos (por exemplo, $K=2$) começando de uma sequência totalmente mascarada.
  • Durante a simulação, as posições são desmascaradas com base em uma política (por exemplo, limiares de confiança), e os tokens comprometidos são "forçados pelo professor" a valores de verdade absoluta para evitar a acumulação de erros durante a simulação.
  • A perda é a soma das perdas de entropia cruzada sobre os $K$ passos.
  • Os gradientes são retropropagados através do caminho de retransmissão ($h_k \to h_{k+1}$), permitindo que o modelo aprenda não apenas a prever o passo atual, mas a codificar informações em $h_k$ que sejam maximamente úteis para passos futuros.

Inferência

No momento da inferência, o processo permanece largely inalterado em termos de cronograma de desmascaramento ou política de amostragem. A única adição é que o modelo carrega o estado de retransmissão $h_k$ adiante junto com os tokens comprometidos. O estado de retransmissão serve como um canal de memória diferenciável para informações ainda não comprometidas a tokens discretos.

3. Contribuições Principais

1. Estrutura Relay: Um método geral para incorporar cálculo recorrente em MDMs, treinando o modelo para passar uma retransmissão latente aprendida adiante através dos passos de decodificação via BPTT truncado. É agnóstico à arquitetura e compatível com técnicas de aceleração de MDM existentes, como difusão em blocos e cache KV.
2. Escalabilidade e Adaptação: Validação do Relay na escala de Modelos de Linguagem de Grande Porte (LLM) através da adaptação de parâmetros completos do Fast-dLLM v2 (1.5B). O método adapta com sucesso um Modelo de Linguagem Discreta (DLM) de última geração com FLOPs de treinamento adicionais mínimos.
3. Melhorias de Desempenho-Latência: Demonstração empírica de que o Relay avança a fronteira de Pareto desempenho-latência, superando o ajuste fino supervisionado (SFT) padrão em tarefas de codificação enquanto reduz significativamente a latência de inferência.
4. Análise do Espaço de Design: Estudos de ablação extensivos em uma tarefa desafiadora de planejamento baseada em Sudoku que isolam as contribuições do procedimento de simulação, do canal de retransmissão e do objetivo de treinamento BPTT.

4. Resultados Experimentais

Tarefa de Planejamento Sudoku

Utilizando o conjunto de dados Sudoku-Extreme, os autores compararam quatro objetivos de treinamento:

• MLM: Difusão mascarada uniforme padrão (passo único, sem retransmissão).
• Rollout: Simulação de múltiplos passos com forçamento do professor, mas sem canal de retransmissão ($R_\theta \equiv 0$).
• Relay (sg): Inclui o caminho de retransmissão, mas usa stop-gradients (sem BPTT através da retransmissão).
• Relay: Método completo com $K=2$ BPTT através da retransmissão.

Descobertas:

• Fronteira Precisão-NFE: O Relay alcançou o melhor compromisso entre precisão e Número de Avaliações Diretas (NFE).
• Impacto dos Componentes: A transição de MLM para Rollout melhorou o desempenho. Adicionar o canal de retransmissão (Relay sg) forneceu um salto significativo. Finalmente, habilitar o BPTT (Relay) produziu os melhores resultados, demonstrando que o modelo aprende a manter o tabuleiro parcial autoconsistente sob desmascaramento mais agressivo.
• Eficiência: Em um limiar de confiança fixo, o Relay comprometeu mais células por passagem direta enquanto mantinha a legalidade (sem violações de linha/coluna/bloco), alcançando a mesma precisão em menos passagens diretas totais em comparação com as linhas de base.

Adaptação de Modelo Pré-treinado (Fast-dLLM v2)

Os autores adaptaram o Fast-dLLM v2 (1.5B) usando Relay em uma mistura de conjuntos de dados de código e matemática.

• Desempenho em Codificação: Nos benchmarks HumanEval e MBPP, o Relay superou o SFT padrão.
  • HumanEval: O Relay alcançou 42,1% de precisão (vs. 38,4% para SFT) com 88,3 NFE (vs. 130,7 para SFT).
  • MBPP: O Relay alcançou 46,6% de precisão (vs. 43,9% para SFT) com 78,8 NFE (vs. 84,8 para SFT).
• Latência: O Relay reduziu a latência de inferência em até 32% em comparação com o SFT padrão, enquanto melhorava a precisão.
• Sobrecarga de Memória: O perfilamento em uma GPU A100 mostrou que, embora o Relay aumente a memória de ativação ativa durante a segunda passagem direta, o uso máximo de memória permanece quase idêntico ao SFT puro (aprox. 20,1 GiB vs. 21,2 GiB), porque o gargalo é a passagem de volta da projeção de vocabulário, e não o armazenamento do estado de retransmissão.

5. Significado e Alegações

O artigo afirma que o Relay muda fundamentalmente a natureza da geração não-autorregressiva iterativa, permitindo que os MDMs acumulem cálculo contínuo intermediário.

• Preditivo: O método permite que os modelos otimizem explicitamente representações intermediárias para decisões futuras de desmascaramento, em vez de tratar cada passo como um problema de previsão novo.
• Avanço da Fronteira de Pareto: Ao reduzir o número de passagens diretas necessárias para atingir uma precisão-alvo, o Relay avança a fronteira desempenho-latência, tornando os modelos de difusão discreta mais eficientes para tarefas complexas de raciocínio e geração.
• Compatibilidade: A abordagem é apresentada como uma melhoria plug-and-play que funciona com DLMs de última geração existentes e técnicas de aceleração (por exemplo, difusão em blocos, cache KV) sem exigir reformas arquitetônicas.

Os autores concluem que o Relay representa um passo significativo em direção a uma geração não-gulosa e preditiva, oferecendo um substrato contínuo para raciocínio latente que pode ser explorado ainda mais para interpretabilidade e aplicação em outras modalidades (por exemplo, imagens, moléculas).