Масштабирование спекулятивного декодирования с использованием опережающего рассуждения

Scaling Speculative Decoding with Lookahead Reasoning

June 24, 2025

Авторы: Yichao Fu, Rui Ge, Zelei Shao, Zhijie Deng, Hao Zhang

cs.AI

Аннотация

Модели рассуждений преуспевают в генерации длинных цепочек мыслей, но декодирование получающихся тысяч токенов происходит медленно. Спекулятивное декодирование на уровне токенов (SD) помогает, но его эффективность ограничена, поскольку вероятность того, что вся догадка из гамма-токенов будет верной, экспоненциально снижается с ростом гаммы. Это означает, что выделение большего объема вычислений для более длинных черновиков токенов сталкивается с алгоритмическим потолком, что делает ускорение скромным и независимым от оборудования. Мы поднимаем этот потолок с помощью метода Lookahead Reasoning, который использует второй уровень параллелизма на уровне шагов. Наше ключевое понимание заключается в том, что модели рассуждений генерируют пошагово, и каждый шаг должен быть семантически правильным, а не точным совпадением токенов. В Lookahead Reasoning легковесная модель-черновик предлагает несколько будущих шагов; целевая модель расширяет каждое предложение за один пакетный проход, а верификатор сохраняет семантически правильные шаги, позволяя целевой модели перегенерировать те, которые не прошли проверку. Спекулятивное декодирование на уровне токенов по-прежнему работает внутри каждого шага рассуждения, поэтому два уровня параллелизма умножаются. Мы показываем, что Lookahead Reasoning повышает пиковое ускорение SD как теоретически, так и эмпирически. На тестах GSM8K, AIME и других бенчмарках Lookahead Reasoning улучшает ускорение SD с 1.4x до 2.1x, сохраняя качество ответов, и его ускорение лучше масштабируется с увеличением пропускной способности GPU. Наш код доступен по адресу https://github.com/hao-ai-lab/LookaheadReasoning.

English

Reasoning models excel by generating long chain-of-thoughts, but decoding the resulting thousands of tokens is slow. Token-level speculative decoding (SD) helps, but its benefit is capped, because the chance that an entire gamma-token guess is correct falls exponentially as gamma grows. This means allocating more compute for longer token drafts faces an algorithmic ceiling -- making the speedup modest and hardware-agnostic. We raise this ceiling with Lookahead Reasoning, which exploits a second, step-level layer of parallelism. Our key insight is that reasoning models generate step-by-step, and each step needs only to be semantically correct, not exact token matching. In Lookahead Reasoning, a lightweight draft model proposes several future steps; the target model expands each proposal in one batched pass, and a verifier keeps semantically correct steps while letting the target regenerate any that fail. Token-level SD still operates within each reasoning step, so the two layers of parallelism multiply. We show Lookahead Reasoning lifts the peak speedup of SD both theoretically and empirically. Across GSM8K, AIME, and other benchmarks, Lookahead Reasoning improves the speedup of SD from 1.4x to 2.1x while preserving answer quality, and its speedup scales better with additional GPU throughput. Our code is available at https://github.com/hao-ai-lab/LookaheadReasoning