Décodage spéculatif guidé par la récompense pour un raisonnement efficace de LLM

Reward-Guided Speculative Decoding for Efficient LLM Reasoning

January 31, 2025

Auteurs: Baohao Liao, Yuhui Xu, Hanze Dong, Junnan Li, Christof Monz, Silvio Savarese, Doyen Sahoo, Caiming Xiong

cs.AI

Résumé

Nous introduisons le Décodage Spéculatif Guidé par la Récompense (RSD), un nouveau cadre visant à améliorer l'efficacité de l'inférence dans les grands modèles de langage (LLMs). RSD combine de manière synergique un modèle de brouillon léger avec un modèle cible plus puissant, incorporant un biais contrôlé pour prioriser les sorties à haute récompense, contrairement aux méthodes de décodage spéculatif existantes qui imposent une impartialité stricte. RSD utilise un modèle de récompense de processus pour évaluer les étapes de décodage intermédiaires et décider dynamiquement s'il faut invoquer le modèle cible, optimisant le compromis entre le coût computationnel et la qualité de la sortie. Nous démontrons théoriquement qu'une stratégie de mélange basée sur un seuil atteint un équilibre optimal entre l'utilisation des ressources et les performances. Des évaluations approfondies sur des benchmarks de raisonnement difficiles, y compris des tâches de niveau olympique, montrent que RSD offre des gains d'efficacité significatifs par rapport au décodage avec le seul modèle cible (jusqu'à 4,4 fois moins de FLOPs), tout en obtenant une précision significativement meilleure que la méthode de décodage parallèle en moyenne (jusqu'à +3,5). Ces résultats soulignent RSD comme une approche robuste et rentable pour le déploiement de LLMs dans des scénarios intensifs en ressources.

English

We introduce Reward-Guided Speculative Decoding (RSD), a novel framework aimed at improving the efficiency of inference in large language models (LLMs). RSD synergistically combines a lightweight draft model with a more powerful target model, incorporating a controlled bias to prioritize high-reward outputs, in contrast to existing speculative decoding methods that enforce strict unbiasedness. RSD employs a process reward model to evaluate intermediate decoding steps and dynamically decide whether to invoke the target model, optimizing the trade-off between computational cost and output quality. We theoretically demonstrate that a threshold-based mixture strategy achieves an optimal balance between resource utilization and performance. Extensive evaluations on challenging reasoning benchmarks, including Olympiad-level tasks, show that RSD delivers significant efficiency gains against decoding with the target model only (up to 4.4x fewer FLOPs), while achieving significant better accuracy than parallel decoding method on average (up to +3.5). These results highlight RSD as a robust and cost-effective approach for deploying LLMs in resource-intensive scenarios.