Medusa Heads

Les têtes Medusa sont des têtes de décodage légères ajoutées à la dernière couche cachée d’un grand modèle de langage ; chacune est entraînée à prédire un token de plus que la couche de sortie du modèle. En un passage avant, le modèle produit sa prédiction habituelle du token suivant ainsi que des hypothèses pour les quelques positions suivantes. Une étape d’attention arborescente vérifie ces hypothèses au passage suivant et conserve la plus longue séquence avec laquelle le modèle est d’accord. Medusa est une forme de décodage spéculatif qui ne nécessite pas de modèle brouillon distinct ; il n’y a donc pas de second réseau à entraîner, héberger ou synchroniser. Tianle Cai, Tri Dao et leurs collègues l’ont présentée dans Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads en janvier 2024.

Comment les têtes prédisent les tokens suivants#

À chaque étape standard de génération de texte, le dernier état caché à la position t est envoyé à la tête du modèle de langage, qui évalue le token à t + 1. La tête Medusa k lit le même état caché et évalue le token à t + k + 1. Chaque tête est une seule couche à propagation avant avec une activation SiLU et une connexion résiduelle, suivie d’une projection vers le vocabulaire. La projection commence par une copie de la tête du modèle de langage et la couche à propagation avant est initialisée à zéro ; avant l’entraînement, chaque tête reproduit donc la prédiction du token suivant du modèle de base. L’article montre que cinq têtes suffisent. Cet exemple de code PyTorch construit trois têtes et compare leurs sorties à celle de la tête du modèle de base sur un état caché :

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)  # la couche de sortie du modèle gelé

# Chaque tête Medusa : une couche résiduelle SiLU (initialisée à zéro) et une copie de la tête du modèle de langage
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)

# Toutes les têtes lisent le même dernier état caché ; la tête k cible la position 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])

La sortie [True, True, True] montre que la couche résiduelle initialisée à zéro fait commencer chaque nouvelle tête par une copie exacte de la prédiction du modèle de base. L’entraînement déplace ensuite la tête k vers sa propre position future, tandis que lm_head et le réseau de base restent inchangés.

Attention arborescente et acceptation#

Les meilleures hypothèses des têtes se combinent pour former des continuations candidates. Avec les deux tokens les plus probables de la tête 1 et les trois plus probables de la tête 2, le produit cartésien donne six candidats partageant des préfixes et formant un arbre :

Lors du passage avant suivant, chaque token candidat est ajouté simultanément à l’entrée, et un masque d’attention permet à chaque token de ne voir que ses ancêtres dans l’arbre ; le modèle évalue ainsi toutes les branches en parallèle. La branche la plus longue acceptée par le modèle est conservée, et son dernier état caché alimente les têtes au tour suivant. Comme toutes les têtes lisent le même état caché, la tête 2 propose les mêmes tokens pour chaque branche de la tête 1. Les auteurs ont constaté qu’un arbre clairsemé de 64 nœuds, construit à partir des combinaisons les plus précises, offrait une accélération supérieure à celle d’un arbre dense de 256 nœuds.

Avec le décodage glouton, un token candidat est accepté lorsqu’il correspond au premier choix du modèle. Pour l’échantillonnage, Medusa propose l’acceptation typique, qui conserve tout candidat dont la probabilité selon le modèle de base dépasse un seuil lié à l’entropie de la distribution. Cette méthode accepte plus de tokens, au prix d’une correspondance exacte avec la distribution.

Entraînement de Medusa-1 et Medusa-2#

Medusa-1 n’entraîne que les nouvelles têtes sur un réseau de base gelé ; les sorties du modèle conservent donc leur qualité, et l’article fait état d’une génération plus de 2,2 fois plus rapide. L’entraînement des têtes de Vicuna-7B a pris environ cinq heures sur un seul GPU A100 avec 60 000 conversations ShareGPT. Il s’agit de la même séparation que freeze=23 met en place lors de l’ajustement fin d’Ultralytics YOLO26 : les couches 0 à 22 restent figées et seule la tête de détection de la couche 23 apprend à partir des nouvelles données.

Medusa-2 affine conjointement les têtes et le réseau de base, ce qui améliore la précision des têtes et porte le facteur d’accélération à 2.3-3.6x, grâce à une recette d’entraînement qui préserve les capacités d’origine du réseau de base. Lorsque les données d’entraînement d’origine ne sont pas disponibles, l’auto-distillation entraîne les têtes à partir du texte que le modèle génère lui-même.

Medusa fait partie de plusieurs méthodes qui prédisent des tokens à partir du modèle cible lui-même. EAGLE prédit à partir des caractéristiques internes du modèle avec une petite couche autorégressive, et la prédiction multi-token, utilisée dans DeepSeek-V3, entraîne des modules de prédiction supplémentaires pendant le préentraînement, que les moteurs d’inférence réutilisent ensuite pour produire des prédictions. Ces trois méthodes réduisent le plus la latence d’inférence avec de petites tailles de lot dans les chatbots interactifs, où le GPU passe la majeure partie de chaque passage avant à attendre la mémoire.