Resumo Técnico: xPress: Refinamento Paralelo para Drafters de Difusão em Decodificação Especulativa

1. Definição do Problema

A decodificação especulativa (SD) acelera a geração autorregressiva usando um drafter (rascunhador) leve para propor tokens que serão verificados por um modelo alvo maior. Embora os drafters autorregressivos (AR) (ex: EAGLE) sejam eficazes, eles sofrem com o overhead sequencial ao gerar blocos longos. Os drafters de difusão por bloco (ex: dFlash) abordam isso gerando um bloco inteiro de tokens em uma única passagem direta paralela, reduzindo drasticamente a latência de rascunho.

No entanto, os drafters de difusão possuem uma limitação fundamental de precisão: eles preveem tokens baseados em distribuições marginais em vez de distribuições conjuntas. Em um processo de difusão de passagem única, o token $k$ é amostrado sem condicionamento aos valores realizados dos tokens $1$ através de $k-1$. Essa falta de causalidade intra-bloco leva a sequências onde tokens individuais são localmente prováveis, mas conjuntamente incoerentes (ex: discordância entre sujeito e verbo). Durante a verificação, o modelo alvo rejeita essas sequências incoerentes precocmente, limitando severamente o comprimento de aceitação ($\tau$) e anulando o potencial de aceleração do rascunho paralelo.

Soluções existentes, como o rascunho baseado em árvore (ex: PRESTO, DDTree), tentam mitigar isso expandindo o espaço de candidatos em uma árvore. No entanto, esses métodos são computacionalmente caros, complexos de servir e muitas vezes contornam o problema visando o recall em vez de corrigir a incoerência causal subjacente.

2. Metodologia: xPress

Os autores propõem o xPress, um refinador causal leve projetado para restaurar as dependências causais ausentes em drafters de difusão sem retornar a um loop autorregressivo lento, um por um.

2.1 Conceito Central

O xPress opera como um "refinador causal" que recebe a proposta inicial de bloco do drafter de difusão e a reconcilia através de refinamento paralelo. Em vez de regerar tokens sequencialmente, o xPress injeta informação causal nos logits do drafter de difusão para corrigir as distribuições marginais.

2.2 Arquitetura

O refinador é um módulo leve que adiciona um viés de logit aprendido ($\delta_k$) aos logits base ($s_k$) do drafter em cada posição $k$. A arquitetura consiste em quatro estágios:

1. Fuse (Fusão): Projeta três entradas em um espaço de baixo rank ($r=256$):
   • O estado oculto por posição do drafter ($h_k$).
   • Um resumo global do bloco ($g$) obtido via mean-pooling.
   • O embedding do token anterior presumido ($x_{k-1}$).
2. Mix (Mistura): Aplica um misturador linear estritamente triangular inferior ($L$) através das posições do bloco. Isso permite que a posição $k$ agregue informações de todos os predecessores ($j \le k$) em uma única operação, impondo uma estrutura causal sem computação sequencial.
3. Trans (Transformação): Um MLP residual no espaço de baixo rank.
4. Read (Leitura): Projeta o resultado de volta para o tamanho do vocabulário para produzir o viés $\delta_k$.

O logit refinado final é $\ell_k = s_k + \delta_k$.

2.3 Refinamento Paralelo via Decodificação de Jacobi

Um refinador causal direto exigiria uma passagem sequencial da esquerda para a direita, reintroduzindo a latência que o xPress visa evitar. Para resolver isso, o xPress utiliza a decodificação de Jacobi:

• Inicialização: O bloco é semeado com as previsões gananciosas (greedy) iniciais do drafter.
• Iteração: Todas as posições são atualizadas simultaneamente em paralelo usando o estado atual do bloco como o prefixo para a próxima iteração.
• Convergência: O processo se repete por $K$ iterações. Teoricamente, para um bloco de tamanho $B$, o ponto fixo sequencial é alcançado em no máximo $B$ iterações. Empiricamente, os autores descobrem que $K \approx 6$ iterações são suficientes para igualar o comprimento de aceitação de um decodificação sequencial completo, pois a semente do drafter já é um palpite de alta qualidade.

2.4 Estratégia de Treinamento

O xPress é co-treinado com o drafter de difusão usando uma função de perda composta:

1. Cross-Entropy (Teacher Forcing): Maximiza a probabilidade do token real dado o prefixo real.
2. Total Variation (TV) Loss: Minimiza a distância entre a distribuição refinada e a distribuição do modelo alvo, otimizando diretamente para a taxa de aceitação.
3. Consistency Loss (Perda de Consistência): Uma perda auxiliar aplicada a uma segunda passagem direta onde as entradas são as próprias previsões do drafter (autocondicionamento). Isso alinha a distribuição de treinamento com o processo de Jacobi em tempo de inferência, evitando que o refinador sofra overfitting para prefixos de ouro (gold prefixes).
4. Drafter Anchor (Âncora do Drafter): Um termo de perda aplicado à saída do drafter base para garantir que o modelo de difusão subjacente permaneça de alta qualidade.

3. Principais Contribuições

• Injeção de Dependência Causal: O xPress injeta informação causal em drafters de difusão usando um mecanismo mínimo (projeção de baixo rank e um misturador triangular) que opera nos estados ocultos do drafter, evitando a necessidade de aumentos de parâmetros significativos.
• Refinamento Paralelo: Ao utilizar a decodificação de Jacobi, o xPress resolve dependências causais em paralelo. Ele alcança o efeito de $B$ passos autorregressivos dependentes usando apenas $K \ll B$ iterações paralelas, preservando os benefícios de velocidade do rascunho paralelo.
• Ganhos de Ponta a Ponta Consistentes: O método melhora significativamente tanto o comprimento de aceitação quanto o throughput sem modificar o processo de verificação ou exigir um drafter maior.

4. Resultados Experimentais

Experimentos foram conduzidos no Qwen3-8B usando um drafter dFlash de bloco-16 em sete benchmarks (Math: GSM8K, MATH-500, AIME25; Code: HumanEval, MBPP, LiveCodeBench; Chat: MT-Bench).

• Comprimento de Aceitação: O xPress aumenta o comprimento médio de aceitação em ~30% comparado ao drafter de difusão dFlash original.
  • Sob decodificação gananciosa ($T=0$), os ganhos variam de 18% a 56%.
  • Sob amostragem ($T=1$), os ganhos variam de 19% a 58%.
• Aceleração de Throughput:
  • Ganancioso ($T=0$): O xPress alcança uma aceleração média de 1.3× sobre o drafter dFlash, com picos de até 1.7× (ex: em GSM8K).
  • Amostragem ($T=1$): A aceleração melhora para uma média de 1.46×, atingindo até 1.6×.
• Comparação com Refinadores Sequenciais: Quando comparado ao Markov head (um refinador causal sequencial do DSpark), o xPress iguala ou supera seu comprimento de aceitação enquanto roda significativamente mais rápido.
  • Com $K=4$ iterações de Jacobi, o xPress iguala a precisão do Markov head, mas roda 1.6× mais rápido (379 µs vs 601 µs por bloco).
  • O Markov head requer um decodificação serial fixo de 15 passos, enquanto o xPress converge em menos passos paralelos.
• Escalabilidade: Em cenários de serviço de grande lote (vLLM, tamanhos de lote 16 e 32), o xPress mantém sua vantagem. Embora a aceleração absoluta sobre a linha de base AR diminua conforme os lotes crescem (devido à verificação limitada por computação), o xPress consistentemente supera o drafter dFlash puro em 1.2–1.4× em throughput, provando que os benefícios do refinador sobrevivem sob carga.

5. Significância e Alegações

O artigo afirma que o xPress aborda o "descompasso inerente" entre o rascunho de difusão em bloco e a verificação sequencial ao reparar a própria distribuição marginal, em vez de expandir o espaço de busca.

• Eficiência: O xPress demonstra que a informação causal pode ser injetada com custo quase zero em relação ao overhead de rascunho. O custo de refinamento é uma pequena fração do tempo total de rascunho, e a natureza paralela do refinamento evita as penalidades de latência da correção sequencial.
• Ortogonalidade: O método é ortogonal a outras melhorias; ele pode ser combinado com drafters maiores ou métodos de verificação baseados em árvore, pois melhora os logits de "espinha dorsal" nos quais esses métodos dependem.
• Generalização: Os autores sugerem que a filosofia de reconciliar previsões mutuamente dependentes via alguns passos paralelos pode ser aplicável além da decodificação especulativa, onde modelos emitem saídas interdependentes simultaneamente.

O artigo conclui que o xPress consegue restaurar com sucesso a causalidade perdida no rascunho de difusão de passagem única, permitos cadeias de aceitação mais longas e maior throughput sem sacrificar a eficiência paralela que torna os drafters de difusão atraentes.