Что такое Wave Decoding и чем отличается от стандартного авторегрессивного?

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

Wave Decoding — это неавторегрессивный метод генерации текста, предложенный в 2025–2026 годах, который генерирует несколько токенов параллельно в разных «ветках» (кандидатах), а затем выбирает лучшую последовательность на основе скоринга. В отличие от стандартного авторегрессивного декодирования, где каждый следующий токен зависит от всех предыдущих и генерируется последовательно, Decoding|Wave Decoding жертвует небольшим качеством ради значительного прироста параллелизма и снижения latency.

1. Термин: Авторегрессивное декодирование (стандартный подход)

Авторегрессивное декодирование — это основной способ генерации текста в LLM. Модель предсказывает следующий токен, используя все ранее сгенерированные токены как контекст. Формально:

P(t_1, t_2, ..., t_T) = ∏_{i=[[1. Как бы вы спроектировали RAG-систему для 10 000 документов с разной структурой|1]]}^{T} P(t_i | t_1, ..., t_{i-[[1. Как бы вы спроектировали RAG-систему для 10 000 документов с разной структурой|1]]})

Каждый шаг требует полного forward pass модели, что делает генерацию последовательной и дорогой по времени, особенно для длинных последовательностей.

Пример (псевдокод):

def autoregressive_generate(model, prompt, max_len):
    tokens = tokenize(prompt)
    for _ in range(max_len):
        logits = model(tokens)  # forward pass
        next_token = sample(logits[-1])
        tokens.append(next_token)
    return tokens

2. Проблема авторегрессивного декодирования: latency и bottleneck

Главный недостаток — последовательная зависимость: каждый шаг ждёт завершения предыдущего. Это создаёт bottleneck при инференсе, особенно для real-time приложений (чат-боты, голосовые ассистенты). Даже с KV cache (кешированием ключей и значений внимания) latency растёт линейно с длиной генерации.

3. Wave Decoding: определение и интуиция

Wave Decoding — это метод, который генерирует сразу несколько токенов параллельно, используя несколько «веток» (wave) — независимых гипотез. Идея: на каждом шаге модель запускается один раз, но предсказывает распределение для нескольких будущих позиций одновременно, а затем выбирает лучший путь.

Интуитивно: вместо того чтобы идти по одному пути, мы «запускаем волну» из нескольких кандидатов, оцениваем их и выбираем наиболее вероятный.

4. Механизм Wave Decoding: параллельные ветки и выбор лучшей

Decoding|Wave Decoding работает в несколько этапов:

1. Инициализация: на основе текущего контекста модель предсказывает распределение для следующего токена.
2. Формирование веток: выбирается K наиболее вероятных токенов (top-K). Для каждого из них модель (или та же модель с модифицированным forward pass) предсказывает следующий токен, но уже в параллельном режиме — для всех K веток сразу.
3. Скоринг веток: каждая ветка (последовательность из 2+ токенов) оценивается по некоторой метрике (например, сумма логарифмов вероятностей или специальный score).
4. Выбор лучшей: ветка с наивысшим score принимается, и процесс повторяется с новым контекстом.

Псевдокод

def wave_decode(model, prompt, max_len, K=4, wave_len=2):
    tokens = tokenize(prompt)
    while len(tokens) < max_len:
        # 1. Получить распределение для следующего токена
        logits = model(tokens)
        probs = softmax(logits[-1])
        top_k_tokens = top_k(probs, K)
        
        # 2. Параллельно генерировать ветки длины wave_len
        branches = []
        for first_token in top_k_tokens:
            # Для каждой ветки предсказываем второй токен
            branch_tokens = tokens + [first_token]
            logits2 = model(branch_tokens)  # можно закешировать
            second_token = argmax(logits2[-1])
            branches.append((first_token, second_token))
        
        # 3. Оценить каждую ветку (например, сумма log-вероятностей)
        scores = []
        for first, second in branches:
            score = log(probs[first]) + log(softmax(logits2)[second])
            scores.append(score)
        
        # 4. Выбрать лучшую ветку и добавить оба токена
        best_idx = argmax(scores)
        tokens.append(branches[best_idx][0])
        tokens.append(branches[best_idx][1])
    return tokens

На практике wave_len может быть больше 2, а выбор веток может использовать beam search-подобный подход.

5. Сравнение Wave Decoding с авторегрессивным

6. Преимущества Wave Decoding

• Снижение latency: количество forward passes уменьшается в wave_len раз (при wave_len=2 — в 2 раза, при wave_len=4 — в 4 раза). Это критично для интерактивных приложений.
• Возможность батчевой обработки: ветки можно обрабатывать на одном батче, используя параллельные вычисления GPU.
• Гибкость: wave_len и K можно настраивать под требования качества/скорости.

7. Недостатки и trade-off

• Снижение качества: из-за того, что ветки строятся независимо, модель может выбрать субоптимальную последовательность (например, первый токен хорош, но второй плох). Это приводит к режимному коллапсу (mode collapse) или менее разнообразным ответам.
• Увеличение FLOPs: на каждом шаге нужно делать K forward passes (для каждой ветки), что увеличивает общее количество вычислений. Однако за счёт параллелизации на GPU latency всё равно снижается.
• Сложность реализации: требуется модификация forward pass для параллельного скоринга веток, а также управление KV cache для каждой ветки.

8. Варианты и модификации Wave Decoding

• Wave Decoding с lookahead: ветки строятся не только на один шаг вперёд, а на несколько (например, 3–5 токенов). Это улучшает качество, но увеличивает вычислительную нагрузку.
• Adaptive Wave Decoding: K и wave_len динамически меняются в зависимости от уверенности модели (например, если распределение острое — используем меньше веток).
• Wave Decoding + Speculative Decoding: комбинация с спекулятивным декодированием, где «черновик» генерируется маленькой моделью, а большая модель проверяет ветки.

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

Wave Decoding относится к семейству неавторегрессивных методов (non-autoregressive generation). Другие подходы:

• Speculative Decoding: использует маленькую модель для генерации нескольких токенов, а большую — для верификации. Отличается тем, что не требует модификации большой модели.
• Parallel Decoding (например, Blockwise Parallel Decoding): генерирует блок токенов за один forward pass, используя специальные архитектуры (например, masked attention).
• Beam Search: тоже рассматривает несколько гипотез, но последовательно, а не параллельно.

Wave Decoding занимает промежуточное положение: он проще в реализации, чем Blockwise Parallel Decoding, и даёт больший прирост скорости, чем Speculative Decoding при высоком K.

10. Когда применять Wave Decoding

Wave Decoding оправдан, когда:

• Latency критична (чат-боты, голосовые ассистенты, real-time перевод).
• Допустимо небольшое снижение качества (например, для суммаризации новостей, где точность не так важна).
• Доступны мощные GPU с поддержкой параллельных вычислений (большие батчи веток).

Не рекомендуется для задач, где требуется высокая точность (медицинские, юридические ответы) или где каждый токен критичен (генерация кода с синтаксической корректностью).

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

Задача: Реализовать упрощённый Wave Decoding для небольшой LLM (например, GPT-2) и сравнить latency и качество с авторегрессивным декодированием.

Инструменты: Python, PyTorch, transformers (Hugging Face), timeit.

Шаги:

1. Загрузите предобученную модель GPT-2 (distilgpt2 для скорости).
2. Напишите функцию autoregressive_generate, которая генерирует текст пошагово.
3. Напишите функцию wave_decode с параметрами K=4, wave_len=2. Используйте model.generate с do_sample=False для веток? Нет, нужно вручную: для каждой ветки делайте forward pass и выбирайте argmax.
4. Замерьте время генерации для одинакового количества токенов (например, 50 токенов) на одном и том же промпте.
5. Сравните качество: используйте метрику perplexity (вычислите на сгенерированном тексте) или попросите другую модель оценить осмысленность (например, через BERTScore).
6. Поэкспериментируйте с разными K и wave_len.

Ожидаемый результат: Вы увидите, что Wave Decoding даёт прирост скорости в ~1.5–2 раза (в зависимости от wave_len) при незначительном росте perplexity (5–10%). Также можно заметить, что при малом K качество падает сильнее.

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

Навигация

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