Mise à l'échelle du décodage spéculatif avec raisonnement prospectif
Scaling Speculative Decoding with Lookahead Reasoning
June 24, 2025
Auteurs: Yichao Fu, Rui Ge, Zelei Shao, Zhijie Deng, Hao Zhang
cs.AI
Résumé
Les modèles de raisonnement excellent en générant de longues chaînes de pensées, mais le décodage des milliers de tokens résultants est lent. Le décodage spéculatif au niveau des tokens (SD) aide, mais son bénéfice est limité, car la probabilité qu'une supposition entière de gamma-tokens soit correcte diminue de manière exponentielle à mesure que gamma augmente. Cela signifie qu'allouer plus de calcul pour des ébauches de tokens plus longues rencontre un plafond algorithmique, rendant l'accélération modeste et indépendante du matériel. Nous relevons ce plafond avec le Raisonnement Prospectif, qui exploite une deuxième couche de parallélisme au niveau des étapes. Notre idée clé est que les modèles de raisonnement génèrent étape par étape, et chaque étape n'a besoin d'être sémantiquement correcte, pas d'une correspondance exacte des tokens. Dans le Raisonnement Prospectif, un modèle d'ébauche léger propose plusieurs étapes futures ; le modèle cible développe chaque proposition en un seul passage groupé, et un vérificateur conserve les étapes sémantiquement correctes tout en permettant au modèle cible de régénérer celles qui échouent. Le SD au niveau des tokens opère toujours dans chaque étape de raisonnement, donc les deux couches de parallélisme se multiplient. Nous montrons que le Raisonnement Prospectif augmente l'accélération maximale du SD à la fois théoriquement et empiriquement. Sur GSM8K, AIME et d'autres benchmarks, le Raisonnement Prospectif améliore l'accélération du SD de 1,4x à 2,1x tout en préservant la qualité des réponses, et son accélération s'adapte mieux à un débit GPU supplémentaire. Notre code est disponible à l'adresse 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