Speculative Decoding Hierárquico: Medusa, EAGLE-2 e Multi-Token Prediction

A inferência de Large Language Models em lote unitário não esbarra na capacidade de cálculo da GPU. Ela é freada pela velocidade de leitura dos parâmetros na memória HBM.

A decodificação especulativa hierárquica supera esse limite físico. Ao propor múltiplos tokens e validá-los em lote paralelo, sistemas modernos aceleram a geração em até 3,4× sem alterar a distribuição matemática.

Por que a inferência autoregressiva tradicional sofre com gargalo de memória?

A decodificação comum é limitada pela largura de banda de memória porque cada token exige a leitura de todos os pesos da rede. Em Batch Size = 1, os núcleos de cálculo ficam ociosos aguardando transferências da HBM.

O modelo analítico Roofline formaliza esse limite pela intensidade aritmética:

I_A = FLOPs ÷ Bytes_transferidos

Em inferência unitária com precisão BF16, cada parâmetro de peso consome 2 bytes e realiza 2 operações aritméticas. A intensidade aritmética resultante na fase de decodificação densa é de apenas 1 FLOP/Byte:

I_A(decode, B=1) = (2 FLOPs × N_params) ÷ (2 Bytes × N_params) = 1 FLOP/Byte

Em um acelerador NVIDIA H100 SXM5 com 80 GB HBM3 e largura de banda de 3,35 TB/s, a capacidade atinge 989 TFLOPS em FP16/BF16 denso. O ponto de inflexão de hardware situa-se em:

I_roofline = 989 TFLOPS ÷ 3,35 TB/s ≈ 295,2 FLOP/Byte

Como a operação opera em 1 FLOP/Byte, o acelerador utiliza menos de 0,35% da sua computação de pico:

Utilização = 1 FLOP/Byte ÷ 295,2 FLOP/Byte ≈ 0,34%

Para um modelo denso de 70 bilhões de parâmetros em BF16 (140 GB de pesos), o tempo para ler a rede da memória é:

T_leitura = 140 GB ÷ 3.350 GB/s ≈ 41,8 ms

Esse teto físico impõe um limite de cerca de 23,9 tokens/s por sequência isolada. Para acelerar a geração, é preciso elevar a quantidade de trabalho útil por ciclo de transferência da memória.

Como funciona a mecânica clássica de Speculative Decoding e Modified Rejection Sampling?

A decodificação especulativa quebra o gargalo de memória ao separar a proposta de candidatos da validação do modelo principal. Um mecanismo ágil sugere múltiplos tokens, e o modelo principal verifica toda a cadeia em um único passo paralelo.

Formulado por Leviatan et al. (Google Research, ICML 2023) e Chen et al. (DeepMind, 2023), o método opera em dois estágios complementares:

1. Fase de Rascunho (Drafting): Um modelo rascunho compacto M_q gera sequencialmente γ tokens candidatos com baixo custo.
2. Fase de Validação (Verification): O modelo principal M_p avalia os γ candidatos em lote simultâneo através de atenção causal paralela.

Para manter a equivalência estatística exata com a amostragem direta do modelo alvo, o sistema aplica amostragem com rejeição modificada (Modified Rejection Sampling).

Para cada token candidato x̃_i proposto com probabilidade q(x̃_i) e avaliado pelo modelo alvo com probabilidade p(x̃_i):

• O token é aceito com probabilidade calculada por P(aceite) = min(1, p(x̃_i) ÷ q(x̃_i)).
• Se aceito, o token integra a sequência final e o processo avalia a posição seguinte.
• Se rejeitado no índice k ≤ γ, os palpites subsequentes são descartados. O sistema sorteia um token corretivo da distribuição residual normalizada:

p_residual(x) = max(0, p(x) - q(x)) ÷ ∑_y max(0, p(y) - q(y))

Essa formulação garante rigor matemático. A probabilidade composta de emissão preserva a distribuição original p(x):

P_final(x) = min(q(x), p(x)) + max(0, p(x) - q(x)) = p(x)

A decodificação especulativa clássica é totalmente sem perdas (lossless). Ela entrega exatamente o mesmo texto da amostragem padrão, reduzindo expressivamente o tempo total de inferência.

+-------------------------------------------------------------------------+
|                  PIPELINE DE VERIFICAÇÃO ESPECULATIVA                   |
+-------------------------------------------------------------------------+
|                                                                         |
|  [Draft Model / Cabeças] ---> Gera gamma tokens candidatos (x_1..gamma) |
|                                             |                           |
|                                             v                           |
|  [Target Model] ------------> Executa 1 Forward Pass com lote paralelo  |
|                                             |                           |
|                                             v                           |
|  [Critério de Aceite] ------> P(aceite) = min(1, p(x) / q(x))           |
|                               |                      |                  |
|                 (Se Aceito)   v                      v   (Se Rejeitado) |
|           Mantém token no prefixo           Descarta cauda de tokens    |
|           Avança para nó seguinte           Amostra de p_residual(x)    |
+-------------------------------------------------------------------------+

Como a taxa de aceitação e a razão de custo determinam o ganho real de velocidade?

O ganho real de velocidade depende do equilíbrio entre a taxa de acerto dos palpites e o custo do rascunho. Se o rascunho for computacionalmente pesado ou errar previsões, a execução pode ficar mais lenta que a geração comum.

Considere α como a taxa média de aceitação (0 ≤ α ≤ 1), γ como a quantidade de palpites e c = T_draft ÷ T_target como a razão de custo. O número esperado de tokens válidos por ciclo é:

E[N_total] = (1 - α^(γ + 1)) ÷ (1 - α)

A razão de aceleração efetiva (Speedup Ratio S) compara o tempo do baseline com o ciclo especulativo:

S = [ (1 - α^(γ + 1)) ÷ (1 - α) ] ÷ (1 + c × γ)

Para garantir aceleração líquida (S > 1), a taxa média de aceitação deve superar a razão de custo do modelo proponente:

α_médio > c

Manter um modelo de rascunho separado na GPU consome memória adicional e exige alinhar vocabulários e sincronizar contextos. Essa sobrecarga impulsionou a evolução para arquiteturas autocontidas e sem modelo rascunho separado.

O que é Medusa e como múltiplas cabeças substituem o modelo de rascunho?

O Medusa substitui o modelo de rascunho separado por múltiplas cabeças acopladas à última camada do modelo principal. Pequenos blocos residuais aprendem a prever posições futuras a partir das ativações ocultas da própria rede.

Proposto por Cai et al. (Together AI e Princeton, 2024) com código no repositório oficial do Medusa, o método conecta K cabeças MLP à camada final:

• A cabeça original lm_head infere o próximo token imediato x_(t+1).
• A cabeça Medusa 1 projeta o estado latente para estimar x_(t+2).
• A cabeça Medusa 2 projeta o estado latente para estimar x_(t+3).
• Cada cabeça adicional estende a previsão para deslocamentos posteriores x_(t+k+1).

Em vez de avaliar apenas uma sequência linear, o Medusa gera uma árvore combinatória de candidatos. Uma seleção top-k escolhe as alternativas mais prováveis em cada bifurcação:

                  [ Token Base: x_t ]
                           |
             +-------------+-------------+
             |                           |
        [ Ramo A1 ]                 [ Ramo A2 ]
        (Cabeça 1)                  (Cabeça 1)
             |                           |
       +-----+-----+               +-----+-----+
       |           |               |           |
    [ B1_a ]    [ B1_b ]        [ B2_a ]    [ B2_b ]
   (Cabeça 2)  (Cabeça 2)      (Cabeça 2)  (Cabeça 2)

Essa árvore é processada sob uma máscara de atenção causal customizada (Tree Attention Mask). Os nós paralelos prestam atenção apenas aos seus ancestrais diretos, evitando contaminação entre caminhos alternativos.

O modelo principal avalia todos os ramos em um único forward pass paralelo. O algoritmo seleciona o caminho válido mais longo, avançando múltiplos tokens em uma única iteração.

Como o EAGLE-2 introduz decodificação em nível de ativações latentes e árvores dinâmicas?

O EAGLE-2 supera as limitações de cabeças estáticas decodificando ativações latentes contínuas em vez de tokens discretos. O método modela a dependência contextual entre previsões futuras usando uma camada compacta de decodificador Transformer.

Formulado por Li et al. (ICML 2024) e aprimorado no EAGLE-2 pela equipe do SafeAI Lab no GitHub, o sistema explora a geometria suave do espaço latente:

• No vocabulário discreto, variações pontuais causam saltos na distribuição categórica de probabilidades.
• No espaço vetorial intermediário (R^d), a transição contextual é suave e estável ao longo do tempo.

O EAGLE utiliza uma camada decodificadora única (< 1% dos pesos da rede principal) operando sobre os vetores de ativação:

h_(t+1) = TransformerLayer(h_t + Embedding(x_t))

As ativações estimadas são projetadas na cabeça do vocabulário compartilhado. A previsão de x_(t+2) é condicionada diretamente ao estado latente de x_(t+1).

O EAGLE-2 adiciona árvores de expansão dinâmica (Dynamic Draft Trees):

• Baixa incerteza: Em código ou sintaxe regular, a árvore expande em profundidade linear longa.
• Alta incerteza: Em respostas abertas e criativas, a árvore distribui nós em largura com ramos curtos.

Essa flexibilidade eleva o ganho de velocidade para a faixa de 2,5× a 3,4× na família Llama-3, superando o Medusa em mais de 40% em aceitação de tokens.

Como o Multi-Token Prediction nativo atua na arquitetura de modelos como DeepSeek-V3?

O Multi-Token Prediction incorpora o mecanismo especulativo diretamente na fase de pré-treinamento da rede. Cabeças de predição paralelas aprendem representações partilhadas desde a otimização inicial, dispensando módulos adicionais pós-treino.

A pesquisa fundamental da Meta AI por Gloeckle et al. (2024) comprovou que antecipar n tokens futuros melhora o planejamento sintático e acelera a inferência.

No relatório do DeepSeek-V3 (Dezembro de 2024), a arquitetura MoE de 671B integra nativamente um módulo MTP de profundidade unitária:

1. Um bloco Transformer dedicado compartilha as matrizes de embedding e a projeção lm_head com o modelo central.
2. A cada passo, o sistema emite o token corrente e sugere o próximo token especulativo com custo computacional mínimo.
3. Em tarefas técnicas e código, a taxa de aceitação supera 85%, assegurando ganho contínuo de 1,8× em produção.

Como runtimes modernos como vLLM e SGLang gerenciam memória e árvores de candidatos?

Motores de alta performance coordenam o KV Cache das árvores especulativas combinando paginação em blocos com grafos de prefixos. Essa gestão evita a fragmentação de VRAM e recicla nós rejeitados sem custo adicional.

No vLLM, o mecanismo PagedAttention divide o cache em páginas alocadas de forma dinâmica na memória HBM. Durante a validação de árvores:

• O subsistema aloca páginas temporárias compartilhadas para os nós ancestrais comuns.
• Concluído o passo paralelo, os blocos de ramos descartados sofrem desalocação imediata (Block Eviction).
• Somente o caminho aceito é integrado definitivamente à sequência do usuário.

No SGLang, a decodificação especulativa integra-se à árvore de prefixos RadixTree. Os candidatos compõem subárvores transitórias na memória:

• Validar múltiplos tokens resume-se a atualizar ponteiros no grafo da árvore Radix.
• Requisições encadeadas com chamadas de ferramentas e agentes reaproveitam o cache anterior sem recomputação.

O comando abaixo exemplifica a inicialização de um servidor vLLM servindo Llama-3-70B com aceleração EAGLE:

python3 -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Meta-Llama-3-70B-Instruct \
  --tensor-parallel-size 4 \
  --speculative-model yuhuili/EAGLE-llama3-70b-instruct \
  --num-speculative-tokens 5 \
  --gpu-memory-utilization 0.90 \
  --max-model-len 8192

Quais são os trade-offs de hardware entre latência interativa e alto throughput em datacenters?

O ganho prático da decodificação especulativa varia conforme o volume de requisições concorrentes. Em acessos interativos isolados, ela encurta a latência perceptível; sob lotes massivos saturados, pode reduzir o rendimento total.

A tabela resume os parâmetros observados nos principais aceleradores de datacenter:

Em cargas simultâneas elevadas (Batch Size ≥ 64), a GPU atinge a região limitada por computação do gráfico Roofline. Validar árvores de candidatos para múltiplos usuários concorre com o cálculo de tokens úteis.

Sob saturação contínua, o throughput agregado com especulação pode recuar entre 5% e 15%. Para neutralizar esse efeito, runtimes modernos utilizam desvio dinâmico (Dynamic Speculation Bypass): o mecanismo suspende a especulação automaticamente durante picos de fila.

Perguntas frequentes sobre Speculative Decoding hierárquico

A decodificação especulativa altera o vocabulário ou as respostas do modelo?

Não. A amostragem com rejeição modificada preserva com exatidão a distribuição estatística da rede principal. A saída gerada é matematicamente equivalente à decodificação pura.

Qual é a vantagem do EAGLE-2 sobre o Medusa?

O Medusa utiliza cabeças estáticas no espaço de tokens discretos, perdendo contexto em respostas abertas. O EAGLE-2 decodifica representações latentes contínuas com árvores dinâmicas, alcançando taxas de acerto superiores.

Como o sistema lida com a rejeição total de uma árvore de candidatos?

Quando todos os palpites são recusados, o algoritmo extrai um token corretivo da distribuição residual da rede principal. Esse token é emitido no ciclo, impedindo desperdício do passo computacional.

O que torna o Multi-Token Prediction vantajoso em modelos fundacionais?

O MTP treina cabeças especulativas de forma nativa junto com o modelo principal. Essa integração dispensa carregar módulos adicionais em produção, economizando memória VRAM e otimizando a inferência.