Какие trade-offs между разными архитектурами speculative decoding?

Краткий тезис

Speculative decoding — техника ускорения инференса LLM, при которой быстрая draft-модель генерирует несколько токенов, а большая target-модель верифицирует их параллельно. Основные архитектуры различаются по способу построения draft-модели: Independent Draft (отдельная маленькая модель), Self-Speculative (Medusa) (дополнительные головы на той же модели) и Speculative Speculative Feature-Aware (EAGLE) (использование скрытых состояний target-модели). Ключевые trade-offs: overhead по памяти, качество предсказаний (acceptance rate), сложность инференса и возможность параллелизации.

1. Что такое speculative decoding и зачем он нужен?

Speculative decoding (decoding|спекулятивное декодирование) — метод ускорения генерации текста в LLM без потери качества. Идея: маленькая и быстрая draft-модель (например, 1.5B параметров) генерирует последовательность из K токенов за один проход, а большая target-модель (например, 70B) проверяет их все сразу, используя tree attention и rejection sampling. Если токен принимается, он остаётся; если отвергается — target-модель генерирует свой токен, и процесс продолжается.

Зачем LLM с autoregressive decoding генерируют токены последовательно, что bottleneck для latency. Speculative decoding позволяет делать несколько шагов за один forward pass target-модели, достигая ускорения в 2–4x без изменения распределения выходных токенов (lossless).

Ключевые термины

• Draft-модель — маленькая модель, генерирующая кандидатов.
• Target-модель — основная модель, верифицирующая кандидатов.
• Acceptance rate — доля токенов, сгенерированных draft-моделью, которые были приняты target-моделью.
• Tree attention — механизм, позволяющий target-модели обрабатывать несколько последовательностей кандидатов одновременно (как дерево).
• Rejection sampling — процедура, при которой токен принимается, если его вероятность у target-модели не ниже, чем у draft-модели; иначе — семплируется из скорректированного распределения.

2. Архитектура Independent Draft

Independent Draft — классический подход: draft-модель — это отдельная, предварительно обученная маленькая модель (например, DistilGPT2, TinyLlama). Она работает независимо от target-модели.

Плюсы

• Простота можно использовать любую маленькую модель, не требуется дообучение target-модели.
• Гибкость легко заменить draft-модель на другую (например, более быструю или более точную).
• Изоляция draft-модель не зависит от внутреннего состояния target-модели, что упрощает инференс.

Минусы

• 2x memory overhead нужно держать в памяти две полные модели (draft + target). Для больших моделей (70B) это может быть критично.
• Overhead на переключение при каждом шаге нужно передавать скрытые состояния между моделями, что добавляет latency.
• Acceptance rate часто ниже маленькая модель может плохо предсказывать поведение большой, особенно на редких токенах.

Когда использовать когда память не является ограничением (например, несколько GPU), а простота реализации важнее максимального ускорения.

3. Архитектура Self-Speculative (Medusa)

Medusa (или Self-Speculative) — подход, при котором draft-модель не является отдельной, а представляет собой несколько дополнительных голов (heads), добавленных поверх последнего слоя target-модели. Каждая голова предсказывает следующий токен на своём шаге (например, голова 1 предсказывает токен t+1, голова 2 — t+2 и т.д.). Все головы обучаются совместно с target-моделью (или после заморозки основной модели).

Плюсы

• 0.5x memory overhead не нужно хранить отдельную модель; дополнительные головы добавляют лишь несколько процентов параметров.
• Единый forward pass target-модель и головы работают в одном вызове, нет overhead на переключение.
• Acceptance rate выше, чем у Independent Draft, так как головы используют те же скрытые состояния, что и target-модель.

Минусы

• Tree cache management при параллельной верификации нескольких последовательностей (дерево) нужно эффективно управлять кэшем KV-состояний, что усложняет инференс.
• Обучение голов требуется дополнительный этап обучения (fine-tuning) на данных, похожих на целевой домен.
• Ограниченное качество головы предсказывают только фиксированное число шагов вперёд (обычно 3–5), и их точность падает с каждым шагом.

Когда использовать когда память ограничена, и вы готовы дообучать модель ради ускорения.

4. Архитектура Feature-Aware (EAGLE)

EAGLE (Efficient AGgregation of Latent Embeddings) — более продвинутая архитектура, использующая скрытые состояния target-модели для генерации draft-токенов. Вместо отдельных голов, EAGLE обучает лёгкую feature-aware draft-модель, которая принимает на вход скрытые состояния target-модели (например, последний слой) и предсказывает следующий токен. Эта draft-модель может быть рекуррентной (LSTM) или трансформером малого размера.

Плюсы

• Лучшее качество acceptance rate достигает ~82% (против ~60–70% у Independent Draft и Medusa), так как draft-модель использует богатые скрытые представления target-модели.
• Меньший overhead draft-модель очень маленькая (например, 2 слоя трансформера), поэтому дополнительная память незначительна.
• Гибкость можно адаптировать draft-модель под конкретную target-модель без изменения последней.

Минусы

• Coupling между моделями draft-модель сильно завязана на внутреннее представление target-модели; при смене target-модели нужно переобучать draft.
• Сложность реализации требуется доступ к скрытым состояниям target-модели (не все фреймворки это поддерживают), а также дополнительный этап обучения.
• Latency на извлечение состояний хотя overhead мал, передача скрытых состояний из target в draft добавляет микро-задержки.

Когда использовать когда требуется максимальное ускорение и качество, и вы готовы инвестировать в обучение и интеграцию.

5. Сравнительная таблица архитектур

6. Trade-off: скорость vs качество (acceptance rate)

Ускорение speculative decoding прямо пропорционально acceptance rate и длине спекулятивной последовательности (K). Формула ожидаемого ускорения (при идеальной параллелизации):

Speedup ≈ (K * acceptance_rate) / (1 + (K * (1 - acceptance_rate)))

• Высокий acceptance rate (EAGLE) позволяет использовать большее K (например, 5–6) и получать ускорение до 4x.
• Низкий acceptance rate (Independent Draft) ограничивает K (обычно 3–4), иначе слишком много токенов будет отвергнуто, и overhead перевесит выгоду.

Trade-off чем выше acceptance rate, тем больше ускорение, но тем сложнее архитектура и больше затраты на обучение.

7. Trade-off: overhead памяти vs overhead вычислений

• Independent Draft требует много памяти (две модели), но вычисления могут быть эффективно распараллелены (draft и target работают на разных GPU).
• Medusa экономит память, но требует сложного управления tree cache, что может увеличить latency на маленьких батчах.
• EAGLE балансирует: маленькая draft-модель почти не добавляет памяти, но извлечение скрытых состояний может создать узкое место при большом batch size.

Вывод для сценариев с ограниченной памятью (один GPU) лучше Medusa или EAGLE; для multi-GPU с достаточной памятью — Independent Draft может быть проще.

8. Trade-off: простота vs гибкость

• Independent Draft — самый простой в реализации: взяли DistilGPT2, подключили к GPT-3, и готово. Но acceptance rate может быть низким, и ускорение неоптимальным.
• Medusa требует дообучения голов, но после этого работает без отдельной модели. Однако головы фиксированы на определённое K, и их нельзя легко заменить.
• EAGLE даёт лучший acceptance rate, но требует глубокой интеграции с target-моделью (доступ к hidden states) и обучения. Смена target-модели — переобучение draft.

Trade-off если вам нужно быстро прототипировать — Independent Draft; если хотите максимальное ускорение и готовы потратить время — EAGLE.

9. Практические соображения при выборе архитектуры

1. Доступное оборудование сколько GPU, какая память? Если память в дефиците — Medusa или EAGLE.
2. Требования к latency для real-time приложений (чат) нужно минимальное время первого токена. Independent Draft может добавить задержку на переключение; Medusa и EAGLE работают быстрее.
3. Возможность обучения есть ли данные и ресурсы для fine-tuning? Если нет — только Independent Draft.
4. Совместимость с фреймворком поддерживает ли ваша библиотека (Hugging Face, vLLM, TensorRT-LLM) tree attention и доступ к hidden states? EAGLE требует кастомной реализации.
5. Качество генерации speculative decoding не должен менять распределение target-модели (lossless). Все три архитектуры это гарантируют при правильной реализации rejection sampling.

10. Дополнительные аспекты

• Tree attention — ключевой компонент для параллельной верификации. Позволяет target-модели обрабатывать все ветки дерева кандидатов за один forward pass. Реализация требует модификации attention mask.
• KV cache management — при спекулятивном декодировании кэш нужно обновлять нелинейно (ветвление). Medusa и EAGLE требуют специальных стратегий (например, prefix caching).
• Влияние на throughput speculative decoding увеличивает throughput (токенов в секунду) за счёт параллелизма, но может снизить throughput при маленьком batch size из-за overhead.
• Совместимость с другими оптимизациями speculative decoding можно комбинировать с quantization, pruning, flash attention. Однако некоторые комбинации (например, Independent Draft + INT8) могут снизить acceptance rate из-за ошибок квантования.

Пет-проект для закрепления

Задача реализовать и сравнить три архитектуры speculative decoding на небольшой модели (например, GPT-2 Medium как target, DistilGPT2 как draft для Independent Draft; добавить головы для Medusa; обучить лёгкую draft-модель на hidden states для EAGLE). Измерить acceptance rate и ускорение.

Инструменты

• Python, PyTorch, Hugging Face Transformers.
• Для Medusa: библиотека medusa-llm (или кастомная реализация).
• Для EAGLE: референсная реализация из официального репозитория (eagle-draft).

Шаги:

1. Загрузите target-модель (GPT-2 Medium) и draft-модель (DistilGPT2) для Independent Draft.
2. Реализуйте speculative decoding с rejection sampling (без tree attention для простоты).
3. Для Medusa: добавьте 3 линейные головы поверх последнего слоя GPT-2, обучите на датасете (например, OpenWebText) 1 эпоху.
4. Для EAGLE: обучите маленький трансформер (2 слоя, 4 головы) предсказывать следующий токен на основе hidden states GPT-2 (заморозьте target).
5. На тестовом датасете (100 промптов) замерьте:
   • Acceptance rate для каждого K (1..5).
   • Время генерации 100 токенов (с speculative decoding и без).
   • Совпадение распределения (проверьте, что speculative decoding не меняет качество — сравните perplexity на валидации).

Ожидаемый результат

• Independent Draft: acceptance rate ~60%, ускорение ~2x.
• Medusa: acceptance rate ~70%, ускорение ~2.5x.
• EAGLE: acceptance rate ~80%, ускорение ~3x.
• Вывод: EAGLE даёт лучшее ускорение, но требует обучения; Independent Draft — самый простой, но менее эффективный.

Связь с другими вопросами

Навигация

• Предыдущий: 163
• Следующий: 165
• Индекс: 00. Индекс разборов