Что такое Variational Speculative Decoding (VSD) и чем он революционен?

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

Speculative Decoding|Variational Speculative Decoding (VSD) — это метод ускорения инференса больших языковых моделей (LLM), который переформулирует обучение draft модели (быстрой вспомогательной модели). Вместо стандартной максимизации правдоподобия (MLE) токенов, VSD оптимизирует acceptance rate (долю принятых токенов) напрямую через вариационные методы. Это революционно, потому что меняет сам objective обучения draft модели: она учится предсказывать не «следующий токен», а «то, что target модель примет», что приводит к значительно более эффективному ускорению без потери качества.

1. Термин: Speculative Decoding (спекулятивное декодирование)

Speculative Decoding — это техника ускорения генерации текста в LLM без ухудшения качества. Идея: использовать быструю draft модель (например, модель меньшего размера) для генерации нескольких токенов, а затем target модель (большая, точная) проверяет их и принимает или отвергает. Если токен принят — он используется, если отвергнут — target модель генерирует свой токен, и процесс продолжается.

• Draft модель — маленькая, быстрая модель (например, 1B параметров).
• Target модель — большая, медленная, но более точная (например, 70B).
• Acceptance rate — доля токенов, сгенерированных draft моделью, которые были приняты target моделью. Чем выше acceptance rate, тем больше ускорение.

Стандартный speculative decoding использует draft модель, обученную на максимизацию правдоподобия (MLE) — предсказание следующего токена по контексту. Однако такой objective не гарантирует, что draft модель будет генерировать токены, которые target модель с высокой вероятностью примет. Это ключевая проблема, которую решает VSD.

2. Проблема стандартного обучения draft модели

В классическом speculative decoding draft модель обучается на корпусе текстов с помощью MLE (максимизация логарифма вероятности правильного следующего токена). Формально:

L_MLE = - E_{x ~ p_data} [ log p_draft(x_t | x_{<t}) ]

Где p_draft — распределение draft модели, x_t — истинный следующий токен.

Недостаток MLE учит draft модель воспроизводить распределение обучающих данных, а не распределение target модели. Даже если draft модель хорошо предсказывает «естественные» токены, она может генерировать токены, которые target модель отвергнет (например, из-за различий в стиле, знаниях или предпочтениях). Acceptance rate при этом может быть низким, и ускорение будет незначительным.

Пример: Draft модель, обученная на общих текстах, может предлагать токен «кошка», а target модель (обученная на медицинских данных) ожидает «пациент». MLE не учитывает это расхождение.

3. Основная идея VSD: оптимизация acceptance rate напрямую

VSD (Variational Speculative Decoding) предлагает новый objective для обучения draft модели: вместо MLE минимизировать расхождение между распределением draft модели и распределением target модели, но с учётом процесса принятия/отвержения. Ключевая интуиция: draft модель должна генерировать токены, которые target модель с высокой вероятностью примет, то есть токены, для которых p_target(token | context) велико.

VSD использует вариационные методы (отсюда название Variational) для вывода objective, который напрямую максимизирует acceptance rate. Вместо того чтобы учить draft модель предсказывать следующий токен из данных, VSD учит её предсказывать токены, которые максимизируют вероятность принятия target моделью.

Формально, VSD вводит вариационное распределение q(tokens | context) (draft модель) и оптимизирует его так, чтобы оно было как можно ближе к условному распределению target модели p_target(tokens | context), но с учётом того, что только часть токенов будет проверена. Это достигается через ELBO (Evidence Lower Bound) или минимизацию KL-дивергенции между q и p_target на каждом шаге.

4. Математическая формулировка VSD (упрощённо)

Пусть q_θ — draft модель с параметрами θ, p — target модель. В стандартном speculative decoding acceptance rate для одного токена равен:

α = E_{x ~ q_θ} [ min(1, p(x) / q_θ(x)) ]

VSD предлагает оптимизировать θ так, чтобы максимизировать α. Для этого используется вариационная нижняя граница (ELBO) на логарифм acceptance rate:

log α ≥ E_{x ~ q_θ} [ log p(x) - log q_θ(x) ] = - KL(q_θ || p)

Где KL(q_θ || p) — KL-дивергенция (расхождение Кульбака-Лейблера) между распределением draft модели и target модели. Минимизация KL-дивергенции эквивалентна максимизации нижней границы acceptance rate.

Таким образом, objective VSD:

L_VSD = KL(q_θ || p) = E_{x ~ q_θ} [ log q_θ(x) - log p(x) ]

Этот градиент можно оценить с помощью REINFORCE или Gumbel-Softmax для дискретных токенов. На практике VSD использует вариационный вывод — итеративно обновляет q_θ, чтобы она приближалась к p в смысле KL, но только на тех токенах, которые draft модель реально генерирует (on-policy).

Псевдокод обучения draft модели в VSD:

# Псевдокод одного шага обучения VSD
for batch in dataloader:
    context = batch["context"]
    # Генерируем токены от draft модели (on-policy)
    draft_tokens, draft_logprobs = draft_model.sample(context)
    # Получаем логиты target модели для этих токенов
    target_logprobs = target_model.log_probs(context, draft_tokens)
    # Вычисляем KL(q || p) = E[log q - log p]
    kl = (draft_logprobs - target_logprobs).mean()
    # Минимизируем KL
    loss = kl
    loss.backward()
    optimizer.step()

5. Сравнение VSD и стандартного speculative decoding

6. Революционность VSD: смена парадигмы обучения draft модели

Революционность VSD заключается в том, что он меняет сам objective обучения draft модели. Вместо того чтобы учить draft модель быть «хорошей языковой моделью» (предсказывать следующий токен как в обучающих данных), VSD учит её быть «хорошим помощником для target модели» — генерировать токены, которые target модель с высокой вероятностью примет.

Это фундаментальный сдвиг:

• Старая парадигма draft модель — это уменьшенная копия target, обученная на тех же данных. Она пытается имитировать target, но не учитывает, что target может иметь другие предпочтения.
• Новая парадигма (VSD): draft модель оптимизируется специально для задачи ускорения. Она учится «подстраиваться» под target, минимизируя расхождение между их распределениями. Acceptance rate становится не просто метрикой, а целевой функцией.

Это особенно важно для Agentic RAG и других сценариев, где target модель может быть дообучена или иметь специфические знания (например, через RAG). VSD позволяет draft модели адаптироваться к этим изменениям без переобучения на больших корпусах.

7. Преимущества VSD

• Более высокий acceptance rate — draft модель генерирует токены, которые target модель почти всегда принимает, что даёт большее ускорение.
• Гарантированное качество — механизм accept/reject сохраняется, поэтому финальная генерация идентична target модели (без потери качества).
• Адаптивность — VSD может быть использован для адаптации draft модели к конкретной target модели, даже если target модель была дообучена или изменена.
• Эффективность — обучение VSD может быть быстрее, чем обучение draft модели с нуля на больших данных, так как target модель предоставляет «сигнал» о том, какие токены важны.

8. Возможные недостатки и ограничения

• Требуется доступ к target модели во время обучения — это может быть дорого (вызовы большой модели для каждого шага). Однако можно использовать кэширование или аппроксимации.
• Сложность реализации — вариационные градиенты (REINFORCE, Gumbel-Softmax) требуют аккуратной настройки.
• On-policy обучение — draft модель должна генерировать токены, чтобы получить градиенты, что может быть нестабильно.
• Зависимость от размера target — если target модель очень большая, вызов её для каждого токена во время обучения может быть узким местом.

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

VSD — это эволюция speculative decoding. Другие методы:

• Medusa — добавляет несколько головок для предсказания нескольких токенов сразу, но не меняет objective обучения.
• Self-speculative decoding — использует ту же модель в качестве draft (например, через ранний выход), но objective остаётся MLE.
• Lookahead decoding — использует n-граммы или кэш для ускорения.

VSD может быть комбинирован с этими методами: например, draft модель с несколькими головками (как Medusa) может обучаться с VSD objective.

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

Задача Реализовать VSD для ускорения инференса небольшой target модели (например, GPT-2) с помощью draft модели (например, дистиллированной версии или маленького трансформера).

Инструменты Python, PyTorch, Hugging Face Transformers, библиотека для speculative decoding (например, speculative-decoding или написать свою).

Шаги:

1. Выбрать target модель (например, gpt2-medium) и draft модель (например, gpt2).
2. Реализовать базовый speculative decoding (accept/reject по схеме из статьи).
3. Обучить draft модель с помощью VSD:
   • На каждом шаге генерировать токены от draft (on-policy).
   • Получать логиты target для этих токенов.
   • Вычислять KL-дивергенцию и обновлять draft через REINFORCE или Gumbel-Softmax.
4. Сравнить acceptance rate и ускорение до и после VSD-обучения.
5. Визуализировать результаты: график acceptance rate по шагам, сравнение времени генерации.

Ожидаемый результат После VSD-обучения acceptance rate должен вырасти (например, с 0.6 до 0.9), а скорость генерации увеличиться в 2-3 раза по сравнению с базовым speculative decoding.

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

12. Навигация

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

Навигация

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