Medusa Heads

Las cabezas Medusa son cabezas de decodificación ligeras que se añaden a la última capa oculta de un modelo de lenguaje grande; cada una se entrena para predecir un token más adelante que la propia capa de salida del modelo. En una pasada hacia delante, el modelo produce su predicción habitual del siguiente token y, además, predicciones para las siguientes posiciones; un paso de atención en árbol comprueba esas predicciones en la pasada siguiente y conserva la secuencia más larga con la que coincide el modelo. Medusa es una forma de decodificación especulativa que no necesita un modelo borrador independiente, así que no hay una segunda red que entrenar, alojar ni mantener sincronizada. Tianle Cai, Tri Dao y sus colaboradores lo presentaron en Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads en enero de 2024.

Cómo predicen las cabezas los tokens futuros#

Un paso estándar de generación de texto introduce el último estado oculto en la posición t en la cabeza del modelo de lenguaje, que calcula la puntuación del token en t + 1. La cabeza Medusa k lee el mismo estado oculto y calcula la puntuación del token en t + k + 1. Cada cabeza es una única capa de propagación hacia delante con una activación SiLU y una conexión residual, seguida de una proyección al vocabulario. La proyección comienza como una copia de la cabeza del modelo de lenguaje y la capa de propagación hacia delante empieza a cero, así que antes del entrenamiento todas las cabezas reproducen la predicción del siguiente token del modelo base. El artículo concluye que cinco cabezas son suficientes. Este esquema en PyTorch crea tres cabezas y compara sus salidas con la cabeza del modelo base a partir de un 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)  # la capa de salida del modelo congelado

# Cada cabeza Medusa: una capa residual SiLU (inicializada a cero) y una copia de la cabeza 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 las cabezas leen el mismo último estado oculto; la cabeza k apunta a la posición 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 salida [True, True, True] muestra que la capa residual inicializada a cero hace que cada cabeza nueva comience como una copia exacta de la predicción del modelo base. El entrenamiento desplaza entonces la cabeza k hacia su propia posición futura, mientras que lm_head y el modelo base permanecen sin cambios.

Atención en árbol y aceptación#

Las predicciones principales de las cabezas se combinan para formar continuaciones candidatas. Con los dos tokens principales de la cabeza 1 y los tres principales de la cabeza 2, el producto cartesiano genera seis candidatos que comparten prefijos y forman un árbol:

En la siguiente pasada hacia delante, todos los tokens candidatos se añaden a la entrada a la vez, y una máscara de atención permite que cada token vea solo sus propios ancestros en el árbol, de modo que el modelo calcula en paralelo la puntuación de todas las ramas. Se conserva la rama más larga que acepta el modelo, y su último estado oculto alimenta las cabezas para la ronda siguiente. Como todas las cabezas leen el mismo estado oculto, la cabeza 2 propone los mismos tokens para cada rama de la cabeza 1. Los autores descubrieron que un árbol disperso de 64 nodos, creado con las combinaciones más precisas, ofrecía una tasa de aceleración superior a la de un árbol denso de 256 nodos.

Con la decodificación voraz, se acepta un token candidato cuando coincide con la opción principal del modelo. Para el muestreo, Medusa ofrece la aceptación típica, que conserva cualquier candidato cuya probabilidad según el modelo base supere un umbral ligado a la entropía de la distribución; así se aceptan más tokens a costa de no reproducir exactamente la distribución.

Entrenamiento de Medusa-1 y Medusa-2#

Medusa-1 entrena solo las cabezas nuevas sobre un modelo base congelado, de modo que las salidas del modelo mantienen su calidad; el artículo informa de una generación más de 2,2 veces rápida. Entrenar las cabezas de Vicuna-7B llevó unas cinco horas en una única GPU A100 con 60.000 conversaciones de ShareGPT. Esta es la misma división que configura freeze=23 al ajustar Ultralytics YOLO26: las capas 0-22 permanecen fijas y solo la cabeza de detección de la capa 23 aprende de los nuevos datos.

Medusa-2 ajusta finamente las cabezas y el modelo base conjuntamente, lo que aumenta la precisión de las cabezas y eleva la aceleración a 2,3-3,6 veces mediante una receta de entrenamiento que protege las capacidades originales del modelo base. Cuando los datos de entrenamiento originales no están disponibles, la autodestilación entrena las cabezas con texto generado por el propio modelo.

Medusa es una de varias formas de generar tokens preliminares a partir del propio modelo objetivo. EAGLE los genera a partir de las características internas del modelo mediante una pequeña capa autorregresiva, y la predicción de varios tokens, utilizada en DeepSeek-V3, entrena módulos de predicción adicionales durante el preentrenamiento que los motores de servicio reutilizan después como generadores preliminares. Los tres métodos reducen sobre todo la latencia de inferencia con tamaños de lote pequeños, habituales en los chatbots interactivos, donde cada pasada hacia delante deja la GPU casi siempre esperando a la memoria.