Beam search, speculative decoding e Medusa

1

Conceito

Geração autorregressiva tem dependência serial: o token $t+1$ só é finalizado depois de conhecer $t$. Vários métodos gastam trabalho paralelo ao redor dessa cadeia, mas otimizam coisas distintas. Beam search explora sequências para melhorar um objetivo de busca. Speculative decoding e Medusa propõem tokens futuros para o modelo caro validar várias posições numa passagem.

Greedy decoding mantém uma sequência parcial. Beam search mantém $b$ candidatas. Em cada etapa, expande cada beam com próximos tokens, soma log-probabilities aos scores acumulados e preserva as melhores por uma regra. Normalização de comprimento e stopping são essenciais porque a soma bruta favorece sequências curtas. Beams compartilham prefixos do KV cache e ramificam com copy-on-write.

Beam search ajuda quando saída concentrada e likelihood da sequência importam, como em certas transduções. Em diálogo aberto, beams costumam convergir para variações parecidas ou prosa previsível. Aumentar largura eleva compute e cache. Não é técnica geral de aceleração.

Speculative decoding usa um draft model menor ou barato. O draft propõe um bloco futuro de forma autorregressiva. O target model maior avalia essas posições em conjunto, explorando paralelismo semelhante a um prefill curto. Uma regra compara probabilidades do draft e target. Tokens aceitos avançam; após a primeira rejeição, o sufixo dependente é descartado e um token corrigido é amostrado.

Com regra correta de aceitação e correção, speculative sampling preserva a distribuição do target. Isso é mais forte que apenas conferir se o argmax coincide. Exatidão se refere ao sampler matemático sob premissas equivalentes, não a bytes idênticos em kernels de precisão diferente.

Velocidade depende da taxa de aceitação e do custo relativo. Um draft fraco sofre muitas rejeições. Um draft caro economiza pouco. Blocos longos oferecem avanço maior, mas desperdiçam mais trabalho após rejeição precoce. A configuração depende dos prompts, decoding, batch, hardware e carga.

Medusa evita um draft completo ao anexar várias heads leves ao modelo base. Elas preveem tokens em offsets futuros e criam uma árvore de continuações. Uma passagem verifica a árvore e aceita um caminho válido. Treinar essas heads precisa corresponder ao checkpoint; trocar a base pode invalidar a qualidade das propostas.

Tree attention e bookkeeping do cache são a complexidade discreta. Ramos compartilham prefixo e divergem depois. Máscaras impedem um ramo de ler outro. Estados aceitos são commitados em ordem; rejeitados liberam memória. Um bug pode produzir tokens plausíveis e incorretos, portanto testes de distribuição contra o baseline são importantes.

Constrained decoding pode ser combinado se proposta e verificação obedecerem ao mesmo estado gramatical. Tool calls e stops mudam fronteiras. Streaming pode chegar em rajadas quando vários tokens são aceitos; clientes não devem assumir um evento por etapa.

A distinção durável é a finalidade. Beam search paga por hipóteses para escolher uma sequência. Speculation paga um proponente barato para expor paralelismo mantendo o target como autoridade. Medusa move propostas para heads do alvo. Os três abrem futuros; objetivo, contrato de aceitação e custo decidem se isso ajuda.