💬 NLP

Improving Diffusion Language Model Decoding through Joint Search in Generation Order and Token Space

Este artigo introduz o Order-Token Search, um novo método de decodificação para Modelos de Linguagem de Difusão que explora conjuntamente a ordem de geração e o espaço de tokens para superar as linhas de base existentes em benchmarks de raciocínio matemático e codificação.

Autores originais: Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon

Publicado 2026-02-06

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está tentando resolver um quebra-cabeça complexo, como um problema de matemática ou um desafio de programação, mas você tem um assistente mágico (o Modelo de Linguagem de Difusão) que não apenas escreve a resposta da esquerda para a direita como um humano digitando uma letra. Em vez disso, esse assistente começa com uma página em branco cheia de espaços reservados "MASK" e tenta preenchê-los todos de uma vez, adivinhando quais palavras pertencem a cada lugar.

O problema é que o assistente é um pouco caótico. Ele pode preencher os espaços em qualquer ordem que desejar. Às vezes, ele acerta a palavra certa para o primeiro espaço, mas depois fica travado porque adivinhou a palavra errada para o segundo espaço. Outras vezes, ele adivinha as palavras certas, mas na ordem errada, levando a um beco sem saída.

O Jeito Antigo: Adivinhação e Poda

Anteriormente, as pessoas tentavam corrigir isso de duas maneiras principais, ambas com falhas:

A Abordagem "Confiante": O assistente observa os espaços sobre os quais ele se sente mais seguro e os preenche primeiro. É como um trilheiro que só caminha pelo caminho que parece mais sólido.
- O Bom: Geralmente obtém uma boa resposta rapidamente.
- O Ruim: Se o "caminho sólido" levar a um precipício (uma resposta errada), o trilheiro fica preso ali. Ele nunca explora outros caminhos que poderiam ter sido melhores.
A Abordagem "Aleatória": O assistente escolhe os espaços para preencher de forma completamente aleatória.
- O Bom: Ele explora uma enorme variedade de caminhos, então é muito provável que acabe encontrando a solução correta eventualmente.
- O Ruim: Ele é tão disperso que raramente escolhe o melhor caminho na primeira tentativa. É como um trilheiro vagando em círculos; eles podem encontrar o tesouro, mas também cavarão muitos buracos vazios antes.

A Nova Solução: Busca de Ordem-Token (Order-Token Search)

Os autores deste artigo introduziram um novo método chamado Order-Token Search. Pense nisso como uma Equipe de Exploradores trabalhando juntos.

Em vez de enviar um único trilheiro ou uma multidão caótica, o método envia uma pequena equipe (um "feixe" ou beam) de exploradores. Veja como eles trabalham:

Caminhos Divergentes (A Busca): Em intervalos regulares, a equipe se divide. Cada explorador tenta uma estratégia diferente:
- O Explorador A decide preencher a primeira palavra que falta.
- O Explorador B decide preencher a última palavra que falta.
- O Explorador C tenta uma palavra diferente para o meio.
- Analogia: Eles estão explorando tanto onde escrever a seguir (a ordem) quanto o que escrever (o token).
A Planilha de Pontuação (O Estimador de Verossimilhança): Esta é a parte mágica. A equipe tem um juiz especial (o estimador de verossimilhança) que não olha apenas para a resposta final. Em vez disso, o juiz observa cada passo que os exploradores deram.
- O explorador fez um movimento lógico?
- Esta frase parcial faz sentido dado o que foi escrito antes?
- Analogia: Imagine um treinador observando uma corrida de revezamento. Se um corredor tropeça cedo, o treinador não espera que ele termine a corrida para cortá-lo; ele o interrompe imediatamente porque o passo que ele deu estava errado.
Cortando os Becos sem Saída (Poda): O juiz pontua o progresso de cada explorador. Se um explorador está seguendo um caminho que parece improvável de ter sucesso (mesmo que ainda não tenha terminado), a equipe corta esse caminho e foca os recursos nos exploradores que estão nos melhores trajetos.

Por Que Isso Importa

O artigo testou isso em problemas matemáticos difíceis (como os conjuntos de dados GSM8K e MATH500) e tarefas de programação (HumanEval).

O Resultado: A "Equipe de Exploradores" (Order-Token Search) encontrou consistentemente as respostas corretas com mais frequência do que os antigos métodos "Confiante" ou "Aleatório".
A Comparação: Teve um desempenho tão bom que igualou ou até superou métodos que exigem o retreinamento caro e de meses de duração (como o diffu-GRPO). Isso significa que você pode tornar uma IA muito mais inteligente apenas mudando como ela pensa durante o teste, sem precisar retreinar o cérebro em si.

Uma Nota Especial sobre Sudoku

O artigo também testou isso em quebra-cabeças de Sudoku. Curiosamente, não funcionou bem aí. Os autores explicam que o Sudoku exige regras globais estritas (como "não repetir números em uma linha") que a "planilha de pontuação" interna da IA simplesmente não consegue entender. É como dar a uma equipe de trilheiros um mapa que não mostra os precipícios; não importa o quão bem eles explorem, eles não conseguem evitar cair da borda se o mapa estiver errado. Isso sugere que, para algumas tarefas, a própria IA precisa ser treinada de forma diferente, não apenas o método de busca.

Em Resumo

O artigo mostra que, ao permitir que uma IA explore múltiplas formas diferentes de escrever (ordem) e múltiplas palavras diferentes (tokens) simultaneamente, e depois usar um sistema de pontuação inteligente para cortar as ideias ruins precocemente, podemos obter resultados muito melhores de Modelos de Linguagem de Difusão sem a necessidade de retreinar o modelo. Transforma um jogo de adivinhação caótico em uma busca estruturada e eficiente pela verdade.

Resumo Técnico: Melhorando a Decodificação de Modelos de Linguagem de Difusão através da Busca Conjunta no Espaço de Ordem e de Tokens

Declaração do Problema

Modelos de Linguagem de Difusão (DLMs), particularmente Modelos de Difusão Mascarados (MDMs), oferecem uma vantagem distinta sobre os modelos autorregressivos (AR): geração independente de ordem. Ao contrário dos modelos AR que geram tokens estritamente da esquerda para a direita, os MDMs podem realizar o denoising de múltiplos posicionamentos em paralelo, permitindo ordens de geração adaptativas. No entanto, os métodos de decodificação atuais não exploram plenamente essa flexibilidade.

As estratégias de decodificação existentes tipicamente se comprometem com uma única trajetória por execução, governadas por regras heurísticas de re-mascaramento:

Re-mascaramento de Baixa Confiança: Fixa greedily tokens de alta confiança precocemente. Embora isso frequentemente resulte em uma forte acurácia de amostra única (pass@1), restringe a busca a um conjunto estreito de trajetórias, levando a um crescimento lento na cobertura de múltiplas amostras (pass@k) e ao potencial aprisionamento em ótimos locais incorretos.
Re-mascaramento Aleatório: Amostra padrões de re-mascaramento uniformemente para maximizar a diversidade de trajetórias. Embora isso melhore o pass@k ao explorar um espaço de solução mais amplo, sofre de um pass@1 fraco por carecer de um mecanismo principiológico para selecionar a trajetória correta entre candidatos concorrentes.

Isso cria um trade-off entre exploração e explotação: os métodos atuais ou exploram excessivamente a confiança local (perdendo soluções válidas) ou exploram cegamente (falhando em selecionar a melhor solução). O artigo argumenta que o potencial dos DLMs de explorar diversas ordens de geração e valores de tokens permanece não realizado porque a decodificação padrão não realiza uma busca explícita neste espaço conjunto.

Metodologia: Busca Ordem-Token

Para abordar isso, os autores propõem a Busca Ordem-Token, um algoritmo de decodificação que realiza busca em feixe (beam search) no espaço conjunto de ordens de geração e valores de tokens. O algoritmo alterna entre expandir candidatos de trajetórias e podar aqueles pouco promissores.

1. Processo de Busca (Expansão)

A Busca Ordem-Token mantém um feixe de $K$ candidatos parcialmente denoised. Diferente da amostragem padrão que segue uma sequência fixa de decisões, o algoritmo ramifica explicitamente sobre duas fontes de incerteza em intervalos especificados pelo usuário:

Ordem de Geração: Quais posições (tokens) desmascarar a seguir.
Valores de Tokens: Quais tokens atribuir a essas posições.

Para gerenciar a complexidade computacional, a busca utiliza difusão em blocos. Em vez de ramificar em cada etapa individual de denoising (o que causaria um overhead proibitivo), o algoritmo expande candidatos nos limites de blocos contíguos de tokens. Dentro de cada bloco, ocorre o denoising independente. Isso reduz os pontos de decisão do comprimento total da sequência para o número de blocos, permitindo uma diversidade substancial tanto em atribuições de ordem quanto de tokens, mantendo os custos de inferência gerenciáveis.

2. Processo de Poda (Estimativa de Verossimilhança)

A inovação central reside no estimador de verossimilhança utilizado para pontuar e podar trajetórias parciais. Um objetivo de MDM padrão (verossimilhança de reconstruir a sequência limpa completa a partir de um estado ruidoso) é instável para estados intermediários com muitos máscaras restantes, pois força o modelo a resolver problemas de preenchimento (infilling) difíceis para os quais ele não foi otimizado durante o treinamento.

Os autores introduzem uma função de pontuação incremental de nível de ação:

Em vez de pontuar o estado inteiro, o algoritmo pontua as ações de denoising específicas tomadas para transitar de um estado mais corrompido $x_t$ para um estado menos corrompido $x_s$ .
A pontuação $s(x_t; x_s)$ mede com quanta confiança o modelo pode repredizer os novos tokens revelados em $x_s$ , condicionado à predição de sequência completa do modelo $x_0$ (que serve como contexto).
Esta abordagem avalia se as decisões passadas de uma trajetória (escolhas de ordem e token) são suportadas pelo modelo, fornecendo um sinal de menor variância e mais discriminativo para a poda.

A pontuação cumulativa de um candidato é a soma dessas pontuações incrementais sobre os intervalos de busca. O algoritmo retém os top- $K$ candidatos após cada expansão, filtrando efetivamente caminhos incorretos ou subótimos antes de continuar a decodificação.

Contribuições Principais

Identificação do Trade-off: O artigo demonstra empiricamente que os métodos de decodificação de MDM padrão ocupam extremos opostos do espectro de exploração–explotação. O re-mascaramento de baixa confiança explota excessivamente (alto pass@1, baixo crescimento de pass@k), enquanto o re-mascaramento aleatório explora cegamente (alto pass@k, baixo pass@1).
Algoritmo de Busca Ordem-Token: A proposta de um novo algoritmo de decodificação que busca no espaço conjunto de ordens de geração e valores de tokens. Ele utiliza um estimador de verossimilhança incremental estável para podar trajetórias parciais, convertendo a diversidade de trajetória em melhoria da acurácia de amostra única.
Validação Empírica: O método mostra-se consistentemente superior aos baselines em benchmarks de raciocínio matemático e codificação, igualando ou superando os ganhos tipicamente alcançados por métodos caros de pós-treinamento como o diffu-GRPO.

Resultados Experimentais

Os autores avaliaram a Busca Ordem-Token (OTS) em GSM8K, MATH500, Countdown e HumanEval utilizando LLaDA-8B-Instruct e LLaDA-1.5 (uma variante pós-treinada por RL).

Ganhos de Desempenho: A OTS alcançou consistentemente a maior acurácia média entre os benchmarks.
- No LLaDA-8B, a OTS melhorou o pass@1 em +3,1% no GSM8K, +3,8% no MATH500, +7,9% no Countdown e +6,8% no HumanEval comparado ao baseline de re-mascaramento de baixa confiança.
- No LLaDA-1.5, a OTS alcançou melhorias relativas semelhantes, atingindo 40,4% de acurácia média comparado aos 37,2% do melhor baseline (Low-conf + MV).
Comparação com Pós-Treinamento: A OTS igualou ou superou o desempenho de modelos submetidos a pós-treinamento por aprendizado por reforço (diffu-GRPO), sugerindo que a busca em tempo de inferência pode ser tão eficaz quanto o retreinamento do modelo para melhorar capacidades de raciocínio.
Estudos de Ablação:
- Estimador de Verossimilhança: A função de pontuação incremental específica (Eq. 2) foi crítica. Variantes usando a verossimilhança conjunta de todos os blocos revelados ou verossimilhanças do tipo AR performaram significativamente pior, confirmando que o estimador customizado é necessário para uma poda estável em MDMs.
- Espaço de Busca: A OTS superou "Busca de Ordem" (buscando apenas ordem) e "Busca de Token" (buscando apenas tokens) que utilizavam verossimilhanças do tipo AR, destacando a necessidade de busca conjunta e o descompasso das heurísticas AR com MDMs.
Eficiência: A OTS escala de forma mais eficaz com o computação de tempo de teste (NFE) do que a votação majoritária. Com orçamentos de computação equivalentes, a OTS alcançou maior acurácia (ex: 29,3% vs. 19,9% no Countdown) ao alavancar busca estruturada em vez de amostragem não estruturada.

Significância e Alegações

O artigo afirma que a busca conjunta ordem–token é um componente chave para avançar a decodificação de DLMs. Ao explorar explicitamente o espaço de ordens de geração e valores de tokens e realizar a poda baseada em um estimador de verossimilhança estável, o método converte o potencial inerente de múltiplas trajetórias dos modelos de difusão em um desempenho confiável de amostra única.

Os autores posicionam este trabalho como uma solução para o gargalo de "exploração–explotação" na decodificação atual de MDM. Eles afirmam que sua abordagem permite que os DLMs escapem dos ótimos locais que prendem métodos greedy e selecionem as soluções corretas que a amostragem aleatória perde, tudo sem exigir treinamento adicional do modelo. O trabalho estabelece que estratégias de busca em tempo de inferência, especificamente adaptadas para a natureza paralela e independente de ordem dos modelos de difusão, podem gerar melhorias sistemáticas em tarefas de raciocínio e codificação.