DBLAST: Dependent Block Drafting for Stochastic Speculative Decoding

Resumo Técnico: DBlast – Dependent Block Drafting para Decodificação Especulativa Estocástica

Problema

A decodificação especulativa acelera a inferência de Grandes Modelos de Linguagem (LLMs) ao utilizar um rascunhador (drafter) leve para propor múltiplos tokens futuros, que um modelo alvo então verifica em paralelo. Embora rascunhadores baseados em blocos recentes (ex: DFlash) possam prever várias posições em uma única passagem, eles tipicamente dependem da amostragem de blocos independentes. Essa abordagem assume que os tokens dentro de um bloco de rascunho são condicionalmente independentes dado o prefixo.

Essa suposição funciona bem para decodificação gananciosa (greedy) ou de baixa entropia, mas torna-se frágil em regimes de decodificação estocástica (não gananciosa), que são essenciais para tarefas de abertura, como diálogo, escrita criativa e exploração de reforço. Em configurações estocásticas, a distribuição alvo é deliberadamente diversa, e múltiplas continuações são plausíveis. O artigo identifica um descompasso fundamental: enquanto o rascunhador amostra tokens independentemente para corresponder às distribuições marginais do alvo, o modelo alvo os verifica sequencialmente, condicionando cada aceitação aos tokens previamente aceitos. Consequentemente, um bloco com marginais por posição precisas pode falhar em formar uma trajetória condicional coerente, levando a uma degradação significativa no comprimento de rascunho aceito conforme a entropia da distribuição de amostragem do alvo aumenta.

Metodologia

Os autores propõem o DBlast (Dependent Block Drafting), um framework que injeta dependências em rascunhadores de difusão de blocos sem sacrificar a eficiência da predição paralela. O método consiste em dois componentes principais:

1. Amostragem de Bloco Dependente via Mistura Latente:
   Em vez de fatorizar a distribuição de proposta através das posições ($q(x_{1:b}|y) = \prod q_i(x_i|y)$), o DBlast introduz uma variável latente categórica $z$ (um "ramo" ou "modo") para induzir correlações dentro do bloco. A distribuição conjunta é modelada como uma mistura canônica poliádica (CP) de baixo posto:
   $$q_{dep}(x_{1:b} | y) = \sum_{z=1}^{K} q(z | y) \prod_{i=1}^{b} q_i(x_i | y, z)$$

   • Arquitetura: O método modifica um rascunhador estilo DFlash adicionando um "expansor de ocultos de categoria" (category hidden expander) leve, que mapeia os estados ocultos do transformer para offsets específicos de categoria, e um "cabeçalho de prior de categoria" (category-prior head) para prever $q(z|y)$.
   • Inferência: Durante a inferência, uma categoria latente é amostrada (opcionalmente com uma temperatura $Z_T$), e o ramo correspondente é decodificado de forma gananciosa em uma única passagem paralela. Isso preserva a eficiência de uma única passagem do rascunho de blocos, permitendo que o modelo represente continuações de nível de bloco coerentes e alternativas.

2. Objetivo de Treinamento Orientado à Aceitação:
   O treinamento padrão minimiza a Log-Verossimilhança Negativa (NLL) dos blocos alvo, o que otimiza a verossimilhança da proposta, mas ignora a natureza sequencial da verificação especulativa. O DBlast introduz um substituto (surrogate) baseado no Comprimento Aceito Esperado (AL - Expected Accepted Length).

   • O objetivo é derivado da expectativa do comprimento do prefixo aceito, decomposto em eventos de aceitação de prefixo.
   • Para lidar com a instabilidade da ponderação de importância quando as distribuições do rascunhador e do alvo possuem baixa sobreposição (no início do treinamento), os autores propõem um substituto no domínio do log truncado por limiar. Este erro soma a razão de log das probabilidades e o log do comprimento aceito condicional sobre os prefixos retidos onde a razão de importância excede um limiar $\tau$.
   • Esta abordagem visa diretamente o comportamento de aceitação sequencial do verificador, complementando a modelagem de dependência quando a verossimilhança do bloco isoladamente é insuficiente.

Principais Contribuições

1. Diagnóstico do Descompasso: Os autores demonstram empiricamente que a amostragem de bloco independente falha cada vez mais à medida que a decodificação do alvo se torna menos determinística. Eles mostram que o comprimento aceito degrada significativamente em regimes de alta entropia porque as propostas independentes carecem da coerência necessária para a verificação sequencial.
2. Amostragem de Bloco Dependente: Eles introduzem a primeira aplicação de propostas de mistura latente do tipo CP para rascunhadores de difusão de subpalavras (subword block-diffusion drafters) para amostragem especulativa não gananciosa. Este método induz correlações entre as posições do bloco enquanto mantém a predição paralela de tokens.
3. Perda Orientada à Aceitação: Eles propõem um substituto de treinamento inovador motivado pelo comprimento aceito esperado. Este objetivo, combinado com a truncagem por limiar, fornece um sinal de treinamento estável que alinha o rascunhador com o processo de verificação sequencial.

Resultados Experimentais

Experimentos foram conduzidos utilizando modelos alvo Qwen3-4B e Qwen3-8B em quatro benchmarks: GSM8K (matemática), MT-Bench (chat), HumanEval (codificação) e um benchmark de escrita criativa. A avaliação cobriu três regimes de amostragem do alvo, variando de alta determinância (Temperatura 0.7) a baixa determinância (Temperatura 1.5).

• Ganhos de Desempenho: O DBlast superou consistentemente a amostragem de bloco independente (DFlash) e os baselines treinados com NLL.
  • No cenário de alta entropia (Temperatura do Alvo 1.5), o DBlast alcançou um ganho médio macro de 12,1% no comprimento aceito para o Qwen3-8B em comparação com o baseline independente.
  • Os ganhos foram mais pronunciados em regiões de baixa determinância, confirmando que a modelagem de dependência é mais valiosa quando múltiplas continuações são plausíveis.
• Estudos de Ablação:
  • Categorias Latentes ($K$): O aumento de $K$ melhorou o desempenho, com retornos decrescentes observados após $K=4$.
  • Objetivo de Treinamento: A perda de aceitação (AL) superou consistentemente o treinamento NLL, mesmo para rascunhadores independentes ($K=1$), embora a combinação de dependência e AL tenha gerado os melhores resultados.
  • Temperatura de Inferência: A calibração da temperatura de amostragem de categoria ($Z_T$) permitiu que o rascunhador melhor correspondesse à estocasticidade da distribuição do alvo, melhorando ainda mais o comprimento aceito.

Significância e Alegações

O artigo afirma que o DBlast aborda uma lacuna crítica na decodificação especulativa: a incapacidade dos atuais rascunhadores de blocos lidarem eficientemente com a verificação estocástica. Ao demonstrar que a modelagem de dependência intra-bloco e o treinamento alinhado à aceitação são complementares, os autores mostram que é possível alcançar altas taxas de aceitação em diversos cenários de decodificação não gananciosa sem sacrificar a eficiência paralela que torna a decodificação especulativa atraente.

Os autores observam modestamente que, embora seu substituto (surrogate loss) não seja um estimador imparcial ou um limite inferior formal, ele serve como um sinal de treinamento prático e eficaz. Eles também reconhecem que a relação entre a mistura latente suave usada durante o treinamento e a decodificação de ramo ganancioso usada na inferência permanece uma área aberta para caracterização formal, embora seus resultados empíricos demonstrem uma transferência eficaz. O trabalho sugere que futuros sistemas de decodificação especulativa para geração de texto aberto devem ir além da predição de tokens independentes para capturar a coerência estrutural de sequências futuras.