Как работает speculative decoding? Как выбрать draft модель?

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

Speculative decoding — это техника ускорения инференса больших языковых моделей (LLM), при которой маленькая draft модель (1–3 млрд параметров) генерирует несколько токенов за один шаг, а большая target модель (например, 70B) параллельно верифицирует их корректность. Это позволяет получить ускорение в 2–3 раза без потери качества, так как распределение target модели сохраняется. Выбор draft модели зависит от доступных вычислительных ресурсов и требуемого ускорения: можно использовать отдельную маленькую модель (Independent Draft), квантованную версию target (Self-Speculative) или специализированную архитектуру, предсказывающую скрытые состояния (EAGLE-3).

1. Мотивация: проблема авторегрессивной генерации

Стандартный инференс LLM является авторегрессивным: каждый следующий токен генерируется на основе всех предыдущих. Это приводит к высокой задержке (latency) и низкой пропускной способности (throughput), особенно для больших моделей (70B+). Даже с использованием KV cache каждый шаг требует полного forward pass модели, что дорого.

Speculative decoding решает эту проблему, позволяя генерировать несколько токенов за один forward pass target модели, сохраняя при этом точное соответствие распределению target модели (без аппроксимации). Это достигается за счёт использования маленькой draft модели, которая быстро генерирует «черновик» из K токенов, а target модель затем проверяет их все параллельно.

2. Алгоритм speculative decoding

Алгоритм состоит из трёх этапов:

1. Draft generation: маленькая draft модель (например, 1.3B) авторегрессивно генерирует K токенов (обычно K=4–10) на основе текущего контекста. Это быстро, так как модель маленькая.

2. Parallel verification: большая target модель (например, 70B) получает на вход контекст + все K сгенерированных токенов и выполняет один forward pass, вычисляя логиты для каждого из K+1 шагов (включая исходный контекст). Это возможно благодаря маскированию внимания (causal mask) — target модель обрабатывает все позиции одновременно, но каждый токен «видит» только предыдущие.

3. Acceptance: для каждого из K токенов target модель вычисляет вероятность, с которой она сама бы его сгенерировала. Токен принимается, если он совпадает с наиболее вероятным токеном target модели. Если на каком-то шаге токен не совпадает, процесс останавливается, и принимается самый длинный префикс совпадающих токенов. Для первого несовпадающего токена применяется rejection sampling: он может быть принят с вероятностью min(1, p_target / p_draft), что гарантирует точное соответствие распределению target модели.

Псевдокод:

def speculative_decode(draft_model, target_model, context, K):
    # 1. Draft generation
    draft_tokens = []
    for _ in range(K):
        logits = draft_model(context + draft_tokens)
        token = sample(logits)
        draft_tokens.append(token)
    
    # 2. Parallel verification
    all_logits = target_model(context + draft_tokens)  # shape: (K+1, vocab)
    
    # 3. Acceptance
    accepted = []
    for i in range(K):
        p_target = softmax(all_logits[i])[draft_tokens[i]]
        p_draft = draft_model.probs[i][draft_tokens[i]]
        if p_target >= p_draft:
            accepted.append(draft_tokens[i])
        else:
            # Rejection sampling
            if random() < p_target / p_draft:
                accepted.append(draft_tokens[i])
            else:
                # Sample from target distribution
                accepted.append(sample(all_logits[i]))
                break
    return accepted

3. Acceptance rate и ускорение

Acceptance rate — доля токенов, сгенерированных draft моделью, которые были приняты target моделью. Типичные значения: 50–80% для хорошо подобранной draft модели.

Ускорение вычисляется по формуле:

Speedup = (K * t_draft + t_target) / (t_draft * K + t_target)

где t_draft — время одного шага draft модели, t_target — время одного шага target модели. Если t_draft << t_target, ускорение приближается к 1 / (1 - acceptance_rate). Например, при acceptance rate 80% и K=10 ускорение может достигать 3–4x.

На практике acceptance rate зависит от:

• Размера draft модели: чем больше, тем выше acceptance, но меньше ускорение от её использования.
• Сложности задачи: на простых задачах (например, продолжение текста) acceptance выше, чем на сложных (например, математические рассуждения).
• Степени «выученности» draft модели: если draft модель fine-tuned на данных target, acceptance растёт.

4. Выбор draft модели

Существует три основных подхода к выбору draft модели:

4.1 Independent Draft (отдельная маленькая модель)

Используется предобученная маленькая модель (например, LLaMA-1.3B или TinyLlama) в качестве draft. Преимущества: простота, не требует модификации target модели. Недостатки: acceptance rate может быть низким, если распределения моделей сильно различаются.

4.2 Self-Speculative (квантованная target)

В качестве draft используется квантованная версия самой target модели (например, INT8 или NF4). Преимущества: высокая согласованность распределений (acceptance rate > 80%), не нужно обучать отдельную модель. Недостатки: квантованная модель всё ещё может быть относительно медленной, требует дополнительной памяти для хранения квантованных весов.

4.3 Feature-aware (EAGLE-3)

EAGLE-3 — это специализированная архитектура draft модели, которая предсказывает не только следующий токен, но и скрытые состояния (hidden states) target модели. Она обучается на данных, сгенерированных target моделью, и использует feature-level информацию для более точного предсказания. Преимущества: очень высокий acceptance rate (до 90%), низкая вычислительная стоимость. Недостатки: требует обучения на большом количестве данных, сложность реализации.

Сравнительная таблица:

5. Практические рекомендации по выбору draft модели

1. Если у вас ограниченные ресурсы и нет возможности обучать модель — используйте Independent Draft с моделью, похожей по архитектуре на target (например, LLaMA-7B как draft для LLaMA-70B). Это даст ускорение 1.5–2x.

2. Если target модель уже квантована (например, для деплоя) — попробуйте Self-Speculative с более грубой квантизацией draft (например, target в FP16, draft в INT4). Acceptance rate будет высоким.

3. Если вы готовы инвестировать в обучение и хотите максимального ускорения (3–4x) — используйте EAGLE-3 или аналогичные feature-aware подходы. Это особенно актуально для продакшн-систем с высокой нагрузкой.

4. Метрики для оценки:

   • Acceptance rate (должен быть > 60% для ощутимого ускорения).
   • Latency (среднее время на токен).
   • Throughput (токенов в секунду).
   • Quality (perplexity или точность на бенчмарках — должна остаться неизменной).

6. Ограничения и подводные камни

• Speculative decoding не уменьшает количество вычислений — он перераспределяет их: draft модель делает много маленьких шагов, target — один большой. Если draft модель слишком медленная (например, 7B), ускорение может быть отрицательным.
• Параллельная верификация требует модификации кода инференса (поддержка causal masking и батчевой обработки).
• Rejection sampling может привести к небольшому увеличению дисперсии генерации, хотя математически распределение сохраняется.
• Для очень коротких последовательностей (например, 1–2 токена) overhead от draft генерации может перевесить выгоду.

7. Связь с другими техниками ускорения

Speculative decoding хорошо сочетается с:

• KV cache quantization — уменьшает память и ускоряет draft модель.
• FlashAttention — ускоряет параллельную верификацию.
• Continuous batching — позволяет обслуживать несколько запросов одновременно, увеличивая throughput.

Однако speculative decoding не заменяет квантизацию или pruning — это ортогональные методы.

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

Задача: Реализовать speculative decoding для пары моделей (draft: distilgpt2, target: gpt2) и измерить ускорение.

Инструменты: Python, PyTorch, Hugging Face Transformers.

Шаги:

1. Загрузите distilgpt2 (draft) и gpt2 (target) через AutoModelForCausalLM.
2. Напишите функцию speculative_generate, реализующую алгоритм выше (без rejection sampling для простоты).
3. Сравните время генерации 100 токенов с обычным авторегрессивным методом и speculative decoding.
4. Постройте график зависимости acceptance rate от K (количество черновых токенов).
5. Попробуйте разные draft модели (например, gpt2 как draft для gpt2-medium) и сравните acceptance rate.

Ожидаемый результат: Вы увидите ускорение в 1.5–2x для distilgpt2+gpt2 и более высокий acceptance rate для gpt2+gpt2-medium. Вы также заметите, что при K > 10 acceptance rate падает, так как draft модель начинает «ошибаться».

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

Навигация

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