Medusa Heads

Головы Medusa — это лёгкие головы декодирования, добавленные поверх последнего скрытого слоя большой языковой модели; каждая обучается предсказывать токен, который расположен дальше, чем токен, предсказываемый собственным выходным слоем модели. За один прямой проход модель формирует обычное предсказание следующего токена и варианты для нескольких следующих позиций, а на следующем проходе шаг с древовидным вниманием проверяет эти варианты и сохраняет самую длинную последовательность, с которой согласна модель. Medusa — это разновидность спекулятивного декодирования, которой не нужна отдельная черновая модель, поэтому не требуется обучать, размещать или синхронизировать вторую сеть. Тяньле Цай, Три Дао и коллеги представили её в январе 2024 года в статье Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads.

Как головы предсказывают следующие токены#

При стандартном шаге генерации текста последнее скрытое состояние в позиции t подаётся на языковую голову, которая оценивает токен в позиции t + 1. Голова Medusa с номером k считывает то же скрытое состояние и оценивает токен в позиции t + k + 1. Каждая голова состоит из одного полносвязного слоя с активацией SiLU и остаточным соединением, за которым следует проекция на словарь. Проекция изначально копирует языковую голову, а полносвязный слой начинается с нулевых весов, поэтому до обучения каждая голова повторяет предсказание следующего токена базовой модели. Согласно статье, достаточно пяти голов. Этот пример на PyTorch создаёт три головы и сравнивает их выходы с выходом головы базовой модели для одного скрытого состояния:

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)  # # Выходной слой замороженной модели

# # Каждая голова Medusa: один остаточный слой SiLU с нулевыми начальными весами и копией головы 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)

# # Все головы считывают одно и то же последнее скрытое состояние; голова k нацелена на позицию 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])

Вывод [True, True, True] показывает, что остаточный слой с нулевой инициализацией заставляет каждую новую голову начинать с точной копии предсказания базовой модели. Во время обучения голова k начинает предсказывать собственную будущую позицию, а lm_head и базовая сеть остаются без изменений.

Внимание по дереву и принятие токенов#

Наилучшие варианты от голов объединяются в возможные продолжения. Если взять два лучших токена от головы 1 и три лучших токена от головы 2, декартово произведение даст шесть вариантов с общими префиксами, образующих дерево:

На следующем прямом проходе все токены вариантов сразу добавляются ко входу, а маска внимания позволяет каждому токену видеть только своих предков в дереве, поэтому модель параллельно оценивает все ветви. Сохраняется самая длинная ветвь, которую принимает модель, а её последнее скрытое состояние поступает на головы для следующего шага. Поскольку все головы считывают одно и то же скрытое состояние, голова 2 предлагает одни и те же токены для каждой ветви головы 1. Авторы выяснили, что разреженное дерево из 64 узлов, построенное из наиболее точных комбинаций, ускоряло работу сильнее, чем плотное дерево из 256 узлов.

При жадном декодировании токен-кандидат принимается, если он совпадает с наилучшим выбором модели. Для сэмплирования Medusa предлагает типичное принятие: сохраняются кандидаты, чья вероятность по базовой модели превышает порог, связанный с энтропией распределения. Так принимается больше токенов, но точное соответствие распределению не гарантируется.

Обучение Medusa-1 и Medusa-2#

Medusa-1 обучает только новые головы при замороженной базовой сети, поэтому качество выходов модели сохраняется; в статье сообщается об ускорении генерации более чем в 2.2 раза. Обучение голов Vicuna-7B заняло около пяти часов на одном GPU A100 с использованием 60,000 диалогов ShareGPT. Такое же разделение настраивает freeze=23 при дообучении Ultralytics YOLO26: слои 0–22 остаются замороженными, а только голова детекции слоя 23 обучается на новых данных.

Medusa-2 дообучает головы и базовую сеть совместно, что повышает точность голов и ускорение до 2.3–3.6 раза; рецепт обучения при этом сохраняет исходные возможности базовой сети. Если исходные обучающие данные недоступны, самодистилляция обучает головы на тексте, который модель генерирует самостоятельно.

Medusa — один из нескольких способов генерировать черновые токены с помощью самой целевой модели. EAGLE использует внутренние признаки модели с небольшим авторегрессионным слоем, а предсказание нескольких токенов, применяемое в DeepSeek-V3, добавляет во время предварительного обучения дополнительные модули предсказания, которые затем используют для черновой генерации. Все три метода сильнее всего снижают задержку инференса при небольших размерах пакета в интерактивных чат-ботах, где каждый прямой проход оставляет GPU в основном в ожидании данных из памяти.