Cost-Aware Diffusion Draft Trees for Speculative Decoding

Resumo Técnico: Árvores de Rascunho de Difusão Conscientes de Custo para Decodificação Especulativa

Declaração do Problema
A decodificação especulativa acelera a inferência de modelos de linguagem autorregressivos usando um rascunhador (drafter) leve para propor tokens candidatos, que são então verificados em paralelo pelo modelo alvo. Embora a verificação baseada em árvores permita a verificação simultânea de múltiplas continuações candidatas, os métodos existentes para construir essas árvores de rascunho enfrentam um desalinhamento fundamental em seus objetivos de otimização.

Abordagens atuais, como o DDTree, constroem árvores de rascunho para maximizar o comprimento de aceitação esperado sob um orçamento fixo de nós. No entanto, os autores observam que o comprimento de aceitação é não decrescente em relação ao orçamento; árvores maiores sempre resultam em comprimentos de aceitação maiores, não oferecendo uma base principiológica para selecionar um tamanho de orçamento ideal. Além disso, o orçamento ideal não é estático; ele varia significamente entre as rodadas de decodificação, dependendo da confiança do rascunhador e do comprimento do contexto atual. Os métodos existentes que adaptam a forma da árvore por etapa ainda dependem de um orçamento total de nós fixo e selecionado offline, falhando em responder à variabilidade por rodada no equilíbrio entre cobertura e custo de verificação.

Metodologia: CaDDTree
O artigo apresenta o CaDDTree (Cost-aware Diffusion Draft Tree - Árvore de Rascunho de Difusão Consciente de Custo), um framework que otimiza diretamente o throughput de tokens (tokens gerados esperados por unidade de tempo) em vez do comprimento de aceitação. O método opera em rascunhadores de difusão em bloco (ex: DFlash), que geram distribuições marginais por posição em uma única passagem (forward pass).

1. Objetivo de Throughput: Os autores modelam o custo da rodada como $C(T; \ell) = C_d + C_v(|T|; \ell)$, onde $C_d$ é o custo de rascunho e $C_v$ é o custo de verificação dependente do tamanho da árvore $|T|$ e do comprimento do contexto $\ell$. O objetivo é maximizar a razão entre os tokens aceitos esperados ($1 + \Phi(T)$) e o custo total.
2. Decomposição: Os autores provam que maximizar o throughput sobre todas as árvores válidas reduz-se a uma busca unidimensional sobre o orçamento de nós $n$. Para qualquer $n$ fixo, a estrutura de árvore ideal é a árvore "greedy best-first" (selecionando nós com as maiores probabilidades de prefixo), que já é a construção utilizada pelo DDTree.
3. Unimodalidade e Parada Gananciosa (Greedy Stopping): Uma contribuição teórica fundamental é a prova de que, sob uma condição de convexidade suave na função de custo de verificação, o objetivo de throughput substituto é unimodal em relação ao tamanho da árvore. Essa propriedade permite um algoritmo ganancioso eficiente: o algoritmo adiciona nós incrementalmente à árvore e para assim que o ganho marginal no throughput torna-se não positivo.
4. Orçamento Adaptável: Ao contrário do trabalho anterior, o CaDDTree não requer uma busca em grade (grid search) offline para um orçamento fixo. Em vez disso, ele adapta o orçamento $n^*$ para cada rodada de decodificação com base na distribuição por posição atual e no custo de verificação estimado (que varia com o comprimento do contexto).
5. Perfil de Custo: A função de custo de verificação $C_v(n; \ell)$ é estimada uma única vez offline via regressão convexa em latências medidas através de vários comprimentos de contexto e contagens de nós. Isso garante que o modelo de custo permaneça convexo e monotônico, satisfazendo a garantia de unimodalidade.

Principais Contribuições

• Mudança de Objetivo: Propõe a otimização direta do throughput como o critério primário para a construção da árvore de rascunho, substituindo a maximização do comprimento de aceitação desalinhada.
• Garantia Teórica: Prova que o objetivo de throughput é unimodal sob custos de verificação convexos, permitindo uma regra de parada gananciosa eficiente e provadamente ótima sem treinamento adicional.
• Mecanismo Adaptável: Introduz um método que seleciona dinamicamente o tamanho ideal da árvore por rodada, eliminando a necessidade de ajuste de orçamento manual ou offline.
• Implementação: Instancia o framework usando rascunhadores de difusão em bloco, aproveitando suas distribuições marginais de passagem única para construir árvores de forma eficiente.

Resultados Experimentais
O método foi avaliado no Qwen3-4B e Qwen3-8B em oito benchmarks cobrindo raciocínio (MATH-500, GSM8K, AIME 2025), codificação (HumanEval, MBPP, LiveCodeBench) e seguimento de instruções (MT-Bench, Alpaca).

• Desempenho: O CaDDTree iguala ou supera o DDTree com seleção de orçamento órfão (oracle budget selection) (o melhor orçamento fixo encontrado via busca em grade) em quase todas as tarefas.
• Eficiência: O método alcança tempos de geração por token menores (maior throughput) em comparação com baselines autorregressivos e o DDTree de orçamento fixo. Por exemplo, no Qwen3-8B com MATH-500, o CaDDTree alcança um tempo por token de 4,60 ms comparado aos 5,83 ms do baseline autorregressivo.
• Adaptabilidade: A análise mostra que o orçamento ideal varia significativamente por rodada. Em tarefas de raciocínio como MATH-500, o orçamento selecionado flutua amplamente (desvio padrão $\approx 42$) dependendo da confiança do rascunhador, muitas vezes desviando-se significativamente do orçamento órfão fixo. Em tarefas de seguimento de instruções, o orçamento é mais estável, mas ainda assim distinto dos pontos de grade fixos.
• Robustez: Estudos de ablação indicam que o método é robusto a perturbações nas estimativas de parâmetros de custo, sugerindo que a calibração precisa offline não é estritamente necessária.

Significância
O artigo afirma que o CaDDTree fornece uma abordagem fundamentada e consciente de custo para a decodificação especulativa, resolvendo a tensão entre a cobertura da árvore e a latência de verificação. Ao provar a unimodalidade do objetivo de throughput, os autores permitem um algoritmo eficiente e livre de treinamento que se adapta automaticamente à dificuldade específica de cada etapa de decodificação. Isso elimina a necessidade de seleção heurística de orçamento e demonstra que a construção de árvores dinâmica e consciente de custo pode consistentemente superar baselines estáticos ajustados por oráculo em diversas tarefas e tamanhos de modelo.