Medusa Heads

As cabeças Medusa são cabeças de descodificação leves, adicionadas à última camada oculta de um modelo de linguagem grande; cada uma é treinada para prever um token mais à frente do que a própria camada de saída do modelo. Numa passagem direta, o modelo produz a sua previsão habitual do próximo token e previsões para as posições seguintes; uma etapa de atenção em árvore verifica essas previsões na passagem seguinte e mantém a sequência mais longa com a qual o modelo concorda. Medusa é uma forma de descodificação especulativa que não precisa de um modelo de rascunho separado, por isso não há uma segunda rede para treinar, hospedar ou manter sincronizada. Tianle Cai, Tri Dao e colegas apresentaram-na em Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads, em janeiro de 2024.

Como as cabeças fazem previsões antecipadas#

Uma etapa padrão de geração de texto alimenta o último estado oculto na posição t na cabeça do modelo de linguagem, que calcula a pontuação do token na posição t + 1. A cabeça Medusa k lê o mesmo estado oculto e calcula a pontuação do token na posição t + k + 1. Cada cabeça é uma única camada feed-forward com ativação SiLU e uma ligação residual, seguida de uma projeção para o vocabulário. A projeção começa como uma cópia da cabeça do modelo de linguagem e a camada feed-forward começa a zero; assim, antes do treino, cada cabeça reproduz a previsão do próximo token do modelo base. O artigo conclui que cinco cabeças são suficientes. Este exemplo em PyTorch cria três cabeças e compara as suas saídas com as da cabeça do modelo base num estado oculto:

import torch
import torch.nn.functional as F

hidden_size, vocab_size, num_heads = 256, 1000, 3
lm_head = torch.nn.Linear(hidden_size, vocab_size, bias=False)  # a camada de saída do modelo congelado

# Cada cabeça Medusa: uma camada SiLU residual (inicializada a zero) e uma cópia da cabeça do LM
w1 = [torch.nn.Linear(hidden_size, hidden_size) for _ in range(num_heads)]
w2 = [torch.nn.Linear(hidden_size, vocab_size, bias=False) for _ in range(num_heads)]
for layer, out in zip(w1, w2):
    torch.nn.init.zeros_(layer.weight)
    torch.nn.init.zeros_(layer.bias)
    out.weight.data.copy_(lm_head.weight.data)

# Todas as cabeças leem o mesmo último estado oculto; a cabeça k tem como alvo a posição t + k + 1
h = torch.randn(1, hidden_size)
head_logits = [out(h + F.silu(layer(h))) for layer, out in zip(w1, w2)]
print([torch.equal(logits, lm_head(h)) for logits in head_logits])

A saída [True, True, True] mostra que a camada residual inicializada a zero faz com que cada nova cabeça comece como uma cópia exata da previsão do modelo base. O treino desloca então a cabeça k para a sua própria posição futura, enquanto lm_head e a rede base permanecem inalterados.

Atenção em árvore e aceitação#

As principais previsões das cabeças combinam-se para formar continuações candidatas. Com os dois tokens mais prováveis da cabeça 1 e os três mais prováveis da cabeça 2, o produto cartesiano gera seis candidatos que partilham prefixos e formam uma árvore:

Na passagem direta seguinte, todos os tokens candidatos são acrescentados à entrada de uma só vez, e uma máscara de atenção permite que cada token veja apenas os seus antepassados na árvore, para que o modelo calcule em paralelo as pontuações de todos os ramos. Mantém-se o ramo mais longo aceite pelo modelo, e o seu último estado oculto alimenta as cabeças na ronda seguinte. Como todas as cabeças leem o mesmo estado oculto, a cabeça 2 propõe os mesmos tokens em todos os ramos da cabeça 1. Os autores concluíram que uma árvore esparsa de 64 nós, construída com as combinações mais precisas, acelerava mais do que uma árvore densa de 256 nós.

Com descodificação gulosa, um token candidato é aceite quando corresponde à escolha mais provável do modelo. Para amostragem, Medusa oferece a aceitação típica, que mantém qualquer candidato cuja probabilidade sob o modelo base ultrapasse um limiar associado à entropia da distribuição, aceitando mais tokens à custa de uma correspondência exata com a distribuição.

Treino de Medusa-1 e Medusa-2#

Medusa-1 treina apenas as novas cabeças sobre uma rede base congelada, preservando a qualidade das saídas do modelo; o artigo relata uma geração mais de 2,2 vezes mais rápida. O treino das cabeças do Vicuna-7B levou cerca de cinco horas numa única GPU A100, usando 60.000 conversas do ShareGPT. É a mesma divisão que freeze=23 configura ao ajustar finamente o Ultralytics YOLO26: as camadas 0-22 permanecem fixas e apenas a cabeça de deteção da camada 23 aprende com os novos dados.

Medusa-2 ajusta finamente as cabeças e a rede base em conjunto, o que torna as cabeças mais precisas e aumenta a aceleração para 2,3-3,6 vezes, usando uma receita de treino que protege as capacidades originais da rede base. Quando os dados de treino originais não estão disponíveis, a autodestilação treina as cabeças com texto gerado pelo próprio modelo.

Medusa está entre vários métodos de elaborar tokens a partir do próprio modelo-alvo. O EAGLE elabora tokens a partir das características internas do modelo com uma pequena camada autorregressiva, e a previsão de vários tokens, usada no DeepSeek-V3, treina módulos de previsão adicionais durante o pré-treino que os motores de inferência reutilizam depois para elaborar previsões. Os três métodos reduzem mais a latência de inferência nos tamanhos de lote pequenos dos chatbots interativos, em que cada passagem direta deixa a GPU sobretudo à espera da memória.