Medusa Heads

Le teste Medusa sono leggere teste di decodifica aggiunte all'ultimo livello nascosto di un modello linguistico di grandi dimensioni; ciascuna è addestrata a prevedere un token più avanti rispetto al livello di output del modello. In un singolo passaggio in avanti, il modello produce la consueta predizione del token successivo e le stime per le posizioni successive; nel passaggio seguente, un passaggio di attenzione sull'albero verifica queste stime e conserva la sequenza più lunga concorde con il modello. Medusa è una forma di decodifica speculativa che non richiede un modello bozza separato, quindi non c'è una seconda rete da addestrare, ospitare o mantenere sincronizzata. Tianle Cai, Tri Dao e colleghi la introdussero nel gennaio 2024 con Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads.

Come le teste prevedono i token successivi#

Un passaggio standard di generazione del testo fornisce lo stato nascosto finale alla posizione t alla testa del modello linguistico, che assegna un punteggio al token in t + 1. La testa Medusa k legge lo stesso stato nascosto e assegna un punteggio al token in t + k + 1. Ogni testa è un singolo livello feed-forward con un'attivazione SiLU e una connessione residua, seguiti da una proiezione sul vocabolario. La proiezione inizia come copia della testa del modello linguistico e il livello feed-forward parte da zero; così, prima dell'addestramento, ogni testa riproduce la predizione del token successivo del modello base. Il documento mostra che cinque teste sono sufficienti. Questo esempio in PyTorch crea tre teste e confronta i loro output con quelli della testa del modello base su uno stato nascosto:

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)  # il livello di output del modello congelato

# Ogni testa Medusa: un livello residuo SiLU (inizializzato a zero) e una copia della testa 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)

# Tutte le teste leggono lo stesso stato nascosto finale; la testa k ha come obiettivo la posizione 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])

L'output [True, True, True] mostra che il livello residuo inizializzato a zero fa sì che ogni nuova testa inizi come copia esatta della predizione del modello base. L'addestramento sposta poi la testa k verso la propria posizione futura, mentre lm_head e il backbone restano invariati.

Attenzione sull'albero e accettazione#

Le ipotesi migliori delle teste si combinano in possibili continuazioni. Con i due token migliori della testa 1 e i tre migliori della testa 2, il prodotto cartesiano genera sei candidati che condividono prefissi e formano un albero:

Nel passaggio successivo in avanti, ogni token candidato viene aggiunto all'input in una sola volta e una maschera di attenzione consente a ciascun token di vedere solo i propri antenati nell'albero, così il modello valuta tutti i rami in parallelo. Viene mantenuto il ramo più lungo accettato dal modello e il suo stato nascosto finale alimenta le teste per il turno successivo. Poiché tutte le teste leggono lo stesso stato nascosto, la testa 2 propone gli stessi token sotto ogni ramo della testa 1. Gli autori hanno osservato che un albero sparso di 64 nodi, costruito con le combinazioni più accurate, garantiva un'accelerazione maggiore rispetto a un albero denso di 256 nodi.

Con la decodifica greedy, un token candidato viene accettato quando coincide con la scelta principale del modello. Per il campionamento, Medusa offre l'accettazione tipica: mantiene ogni candidato la cui probabilità secondo il modello base supera una soglia legata all'entropia della distribuzione, accettando più token a scapito della corrispondenza esatta alla distribuzione.

Addestramento di Medusa-1 e Medusa-2#

Medusa-1 addestra solo le nuove teste su un backbone congelato, preservando la qualità degli output del modello; il documento riporta una generazione oltre 2,2 volte più veloce. L'addestramento delle teste di Vicuna-7B ha richiesto circa cinque ore su una singola GPU A100, usando 60.000 conversazioni ShareGPT. È la stessa suddivisione che freeze=23 configura per il fine-tuning di Ultralytics YOLO26: i livelli 0-22 restano fissi e solo la testa di rilevamento del livello 23 apprende dai nuovi dati.

Medusa-2 esegue il fine-tuning congiunto delle teste e del backbone, rendendo le teste più accurate e portando l'accelerazione a 2,3-3,6 volte, grazie a una ricetta di addestramento che protegge le capacità originali del backbone. Quando i dati di addestramento originali non sono disponibili, l'auto-distillazione addestra le teste sul testo generato dal modello stesso.

Medusa rientra tra i diversi metodi per generare token bozza usando il modello obiettivo stesso. EAGLE genera bozze dalle caratteristiche interne del modello con un piccolo livello autoregressivo; la predizione multi-token, usata in DeepSeek-V3, addestra moduli di predizione aggiuntivi durante il preaddestramento, che i motori di serving riutilizzano poi per generare le bozze. Tutti e tre riducono maggiormente la latenza di inferenza con dimensioni di batch ridotte, tipiche dei chatbot interattivi, nei quali a ogni passaggio in avanti la GPU resta per lo più in attesa della memoria.