Medusa Heads

Medusa-Köpfe sind leichtgewichtige Dekodierköpfe, die auf der letzten verborgenen Schicht eines großen Sprachmodells aufsetzen. Jeder Kopf wird darauf trainiert, ein Token weiter vorauszusagen als die Ausgabeschicht des Modells selbst. In einem Vorwärtsdurchlauf erzeugt das Modell neben seiner üblichen Vorhersage des nächsten Tokens auch Schätzungen für die nächsten Positionen. Im darauffolgenden Durchlauf prüft ein Tree-Attention-Schritt diese Schätzungen und behält die längste Folge bei, der das Modell zustimmt. Medusa ist eine Form des spekulativen Dekodierens, für die kein separates Entwurfsmodell erforderlich ist. Es muss also kein zweites Netz trainiert, gehostet oder synchron gehalten werden. Tianle Cai, Tri Dao und Kollegen stellten das Verfahren im Januar 2024 in Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads vor.

So sagen die Köpfe voraus#

Bei einem üblichen Schritt der Textgenerierung erhält der Sprachmodellkopf den letzten verborgenen Zustand an Position t und bewertet das Token an Position t + 1. Medusa-Kopf k liest denselben verborgenen Zustand und bewertet das Token an Position t + k + 1. Jeder Kopf besteht aus einer einzelnen Feedforward-Schicht mit einer SiLU-Aktivierung und einer Residualverbindung, gefolgt von einer Projektion auf das Vokabular. Die Projektion beginnt als Kopie des Sprachmodellkopfs, die Feedforward-Schicht mit Nullwerten. Daher gibt jeder Kopf vor dem Training dieselbe Vorhersage des nächsten Tokens wie das Basismodell aus. Laut Paper reichen fünf Köpfe aus. Diese PyTorch-Skizze erstellt drei Köpfe und vergleicht ihre Ausgaben anhand eines verborgenen Zustands mit der Ausgabe des Basismodellkopfs:

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)  # die Ausgabeschicht des eingefrorenen Modells

# Jeder Medusa-Kopf: eine SiLU-Schicht mit Residualverbindung (mit Nullwerten initialisiert) und eine Kopie des LM-Kopfs
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)

# Alle Köpfe lesen denselben letzten verborgenen Zustand; Kopf k sagt die Position t + k + 1 voraus
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])

Die Ausgabe [True, True, True] zeigt, dass die mit Nullwerten initialisierte Residualschicht dafür sorgt, dass jeder neue Kopf zunächst eine exakte Kopie der Vorhersage des Basismodells ist. Beim Training wird Kopf k dann auf seine eigene zukünftige Position ausgerichtet, während lm_head und das Grundmodell unverändert bleiben.

Tree Attention und Akzeptanz#

Die besten Schätzungen der Köpfe werden zu möglichen Fortsetzungen kombiniert. Nimmt man die zwei besten Tokens von Kopf 1 und die drei besten von Kopf 2, ergeben sich durch das kartesische Produkt sechs Kandidaten mit gemeinsamen Präfixen, die einen Baum bilden:

Im nächsten Vorwärtsdurchlauf werden alle Kandidatentokens gleichzeitig an die Eingabe angehängt. Eine Aufmerksamkeitsmaske sorgt dafür, dass jedes Token nur seine eigenen Vorfahren im Baum sehen kann. So bewertet das Modell alle Zweige parallel. Der längste vom Modell akzeptierte Zweig wird beibehalten, und sein letzter verborgener Zustand dient den Köpfen als Eingabe für die nächste Runde. Da alle Köpfe denselben verborgenen Zustand lesen, schlägt Kopf 2 unter jedem Zweig von Kopf 1 dieselben Tokens vor. Die Autoren stellten fest, dass ein dünn besetzter Baum mit 64 Knoten, der aus den genauesten Kombinationen gebildet wurde, eine höhere Beschleunigung erzielte als ein dichter Baum mit 256 Knoten.

Beim Greedy-Decoding wird ein Kandidatentoken akzeptiert, wenn es der besten Wahl des Modells entspricht. Für das Sampling bietet Medusa eine typische Akzeptanzmethode: Sie behält jeden Kandidaten bei, dessen Wahrscheinlichkeit unter dem Basismodell einen an die Entropie der Verteilung geknüpften Schwellenwert überschreitet. Dadurch werden mehr Tokens akzeptiert, allerdings stimmt die Verteilung nicht mehr exakt überein.

Training mit Medusa-1 und Medusa-2#

Medusa-1 trainiert nur die neuen Köpfe auf einem eingefrorenen Grundmodell. Dadurch bleibt die Ausgabequalität des Modells erhalten, und das Paper berichtet von einer mehr als 2,2-fachen Beschleunigung der Generierung. Das Training der Vicuna-7B-Köpfe dauerte mit 60.000 ShareGPT-Konversationen etwa fünf Stunden auf einer einzelnen A100-GPU. Dasselbe Aufteilungsschema richtet freeze=23 beim Feinabstimmen von Ultralytics YOLO26 ein: Die Schichten 0–22 bleiben eingefroren, und nur der Erkennungskopf auf Schicht 23 lernt aus den neuen Daten.

Medusa-2 feinabstimmt die Köpfe und das Grundmodell gemeinsam. Dadurch werden die Köpfe genauer und die Beschleunigung steigt auf das 2,3- bis 3,6-Fache. Das Trainingsrezept bewahrt dabei die ursprünglichen Fähigkeiten des Grundmodells. Wenn die ursprünglichen Trainingsdaten nicht verfügbar sind, trainiert die Selbst-Destillation die Köpfe mit Text, den das Modell selbst generiert.

Medusa gehört zu mehreren Verfahren, die Tokens direkt aus dem Zielmodell entwerfen. EAGLE erstellt Entwürfe aus den internen Merkmalen des Modells mithilfe einer kleinen autoregressiven Schicht. Die in DeepSeek-V3 verwendete Multi-Token-Vorhersage trainiert während des Vortrainings zusätzliche Vorhersagemodule, die Serving-Engines später als Entwurfsmodelle wiederverwenden. Alle drei Verfahren senken die Inferenzlatenz vor allem bei den kleinen Batchgrößen interaktiver Chatbots, bei denen der GPU während jedes Vorwärtsdurchlaufs größtenteils auf den Speicher wartet.