Resumo Técnico: DARTree

Problema

Os modelos de linguagem de grande escala (LLMs) autorregressivos modernos são inerentemente sequenciais, gerando apenas um token por etapa enquanto exigem acesso a quase todos os parâmetros do modelo. Isso cria um gargalo de limite de I/O que restringe a velocidade de inferência e desperdiça as capacidades de computação paralela do hardware moderno.

A decodificação especulativa aborda isso usando um "drafter" (rascunhador) leve para propor múltiplos tokens futuros, que são então verificados em paralelo pelo modelo alvo. No entanto, as abordagens existentes enfrentam um compromisso entre latência de proposta e condicionamento causal:

1. Drafters Autorregressivos: Produzem propostas sequencialmente. Embora mantenham dependências causais fortes, sua latência cresce linearmente com o comprimento do rascunho.
2. Drafters de Bloco Baseados em Difusão: Predizem um bloco inteiro de tokens em uma única passagem direta (forward pass), reduzindo drasticamente a latência de proposta. Contudo, eles predizem distribuições marginais por posição condicionadas apenas ao prefixo verificado, ignorando os tokens realizados dentro do mesmo bloco de rascunho. Isso cria um "descompasso causal" com a distribuição do modelo alvo.
3. Correção Causal: Cabeças de correção leves podem propagar tokens selecionados ao longo de uma única cadeia de rascunho para restaurar dependências intra-bloco. No entanto, quando combinadas com a verificação de árvore de rascunho (que expõe múltiplas continuações candidatas ao modelo alvo), surge um novo gargalo de eficiência. Como a pontuação corrigida dos filhos de um nó depende do caminho específico tomado para chegar a esse nó, a construção de árvore padrão de melhor busca (usando um max-heap) força a cabeça de correção a operar sequencialmente: um nó é removido, corrigido e seus filhos são inseridos antes que o próximo nó possa ser expandido. Esse entrelaçamento da inferência de correção e das operações de heap impede o loteamento (batching), tornando a construção da árvore uma fonte significativa de latência.

Metodologia: DARTree

Os autores introduzem o DARTree, um método de decodificação especulativa livre de treinamento que estende um drafter de bloco paralelo com correção causal pré-treinado de uma única cadeia para uma árvore especulativa. A inovação central é o desacoplamento da correção causal dependente do caminho das operações sequenciais de heap.

1. Expansão AR Paralela por Profundidade

Em vez de expandir a árvore nó a nó (o que força a correção sequencial), o DARTree expande a árvore por profundidade em lotes paralelos:

• Rascunho Paralelo por Bloco: O drafter de difusão primeiro gera representações compartilhadas para todo o bloco de rascunho.
• Seleção de Candidatos: Em cada profundidade $d$, o sistema seleciona os $K$ tokens candidatos do topo dos logits base.
• Correção em Lote: Todos os ramos ativos na profundidade atual são processados simultaneamente. A cabeça de correção pré-treinada avalia os candidatos para todos os ramos em um único lote, usando o estado de prefixo específico de cada ramo.
• Pontuação e Retenção: Cada extensão de candidato é pontuada usando sua probabilidade cumulativa corrigida. O sistema retém os $W$ ramos de maior pontuação para a próxima profundidade, carregando seus estados de correção adiante.
• Construção de Superárvore: Este processo se repete para todas as profundidades de rascunho, construindo uma "superárvore de candidatos de largura fixa" onde cada profundidade tem uma carga de trabalho regular e tensionada.

2. Poda de Melhor Busca Adiada

Uma vez que a superárvore completa é construída e todos os nós foram pontuados, o DARTree aplica uma estratégia de poda adiada:

• Seleção Global: Em vez de entrelaçar operações de heap com a correção, o DARTree realiza uma única seleção Top-$B$ global através de todos os prefixos materializados na superárvore.
• Equivalência Teórica: Os autores provam (Lema 1) que, se um bônus de profundidade não positivo ($\beta \le 0$) for aplicado para penalizar prefixos mais profundos, o conjunto de nós selecionados por uma operação Top-$B$ global é idêntico ao conjunto selecionado por um max-heap sequencial de melhor busca. Isso permite que o sistema ignore totalmente o gargalo sequencial.
• Verificação: A árvore resultante compacta é passada para o modelo alvo para uma única passagem direta usando atenção de árvore (tree attention).

Principais Contribuições

1. Identificação de um Gargalo Sequencial: O artigo identifica que o acoplamento da correção causal condicionada ao caminho com a construção de árvore de melhor busca nó a nó cria um gargalo de latência que impede o loteamento eficiente.
2. Construção em Lote por Profundidade Livre de Treinamento: O DARTree introduz um método que aplica a correção causal através de múltiplos ramos em paralelo em cada profundidade, alcançando um aumento de velocidade sem perda (lossless) sem exigir treinamento adicional da cabeça de correção.
3. Estratégia de Poda Adiada: Os autores propõem uma estratégia de poda de melhor busca simplificada que substitui as operações sequenciais de heap por uma única seleção Top-$B$ global após a superárvore ser totalmente pontuada, mantendo a estrutura de árvore ideal sob condições de pontuação específicas.

Resultados Experimentais

Os autores avaliaram o DARTree no Qwen3-4B e Qwen3-8B em sete benchmarks (Math: GSM8K, MATH-500, AIME25; Code: HumanEval, MBPP; Chat: MT-Bench, Alpaca) em temperaturas $T=0$ e $T=1$.

• Desempenho: O DARTree alcançou o maior comprimento de aceitação médio ($\tau$) e aceleração (speedup) em todas as quatro configurações de modelo-temperatura.
  • Na melhor configuração (Qwen3-4B, $T=0$, GSM8K), o DARTree aceitou 12,97 tokens por rodada de verificação, um aumento de 98,6% sobre o DFlash e 27,9% sobre o Domino.
  • Isso resultou em um aceleração sem perda de 9,73× sobre a decodificação autorregressiva medida localmente.
• Comparação:
  • O DARTree (podado) superou o método anterior de estado da arte baseado em difusão, Domino, em até 46,8% em comprimento de aceitação e 40,1% em aceleração.
  • Também superou o DDTree (que usa previsões marginais sem correção) e o Domino-chain + DDTree (que aplica correção apenas a uma única cadeia antes da construção da árvore).
• Ablação:
  • Latência: Uma abordagem de correção sequencial com um heap incorreu em ~70ms por rodada, enquanto o DARTree manteve ~28ms, demonstrando a eficiência da abordagem de lote por profundidade.
  • Generalização: O DARTree melhorou com sucesso o desempenho quando aplicado ao drafter DSpark (usando uma cabeça de correção de Markov), confirmando que é um método geral para drafters paralelos com correção causal.
  • Hiperparâmetros: Aumentar a largura da camada ($W$) melhorou a aceitação até certo ponto (retornos decrescentes após $W=12$), enquanto o orçamento de verificação ($B$) mostrou pico de aceleração em torno de $B=64-128$.

Significância

O artigo afirma que o DARTree resolve a tensão fundamental na decodificação especulativa entre fidelidade causal (que exige correção dependente do caminho) e eficiência computacional (que exige loteamento paralelo). Ao reestruturar o processo de construção da árvore para ser por profundidade em vez de por nó, o DARTree permite o uso de drafters de difusão de alta capacidade com correção causal sem incorrer nas penalidades de latência sequencial das buscas tradicionais baseadas em heap. O método é apresentado como um aprimoramento livre de treinamento que pode ser aplicado a cabeças de correção pré-treinadas existentes, oferecendo ganhos substanciais e consistentes na velocidade de inferência em diversos domínios e temperaturas de decodificação.