Résumé Technique : ReCo (Reward-Coordinated Efficient Reasoning)

Problématique

Les grands modèles de raisonnement (LRM), tels qu'OpenAI-o1 et DeepSeek-R1, atteignent des performances élevées sur des tâches complexes grâce à de longues chaînes de pensée (Chain-of-Thought - CoT). Cependant, ce paradigme conduit fréquemment à une « sur-réflexion » (overthinking), où les modèles génèrent des chaînes de raisonnement excessivement longues même pour des entrées simples, ce qui entraîne une inflation des coûts d'inférence, une latence élevée et une surcharge significative de la mémoire due à l'accumulation du cache KV (Key-Value).

Bien que la compression du cache KV soit une solution standard pour réduire la mémoire et les coûts d'attention, les méthodes actuelles orientées vers le raisonnement souffrent de deux limitations critiques :

1. Politiques de compression uniformes : Les méthodes actuelles appliquent une stratégie de compression uniforme sur l'ensemble de la trajectoire de raisonnement, ignorant que différents états de raisonnement ont des tolérances variables à la perte de contexte.
2. Contrôle de génération découplé : Ces méthodes se concentrent uniquement sur ce qui est supprimé du cache, omettant de prendre en compte le fait que la compression peut induire des longueurs de génération plus importantes. L'article observe que la compression du cache pousse souvent les modèles à générer plus de jetons pour compenser la perte de contexte, annulant partiellement ou totalement les gains d'efficacité (un phénomène nommé « inflation de longueur »).

Méthodologie : Le cadre ReCo

Les auteurs proposent ReCo (Reward-Coordinated Compression), un cadre par étapes qui coordonne la compression du cache KV avec le contrôle de la génération sous un signal unique : une récompense de processus (process reward) légère.

1. Signal central : Estimation de la récompense de processus

Après chaque étape de raisonnement $c_i$ terminée, un estimateur léger de 30 millions de paramètres (distillé à partir de Skywork-o1-Open-PRM-7B) attribue une récompense scalaire $v_i \in [0, 1]$. Cette récompense reflète la probabilité que le chemin de raisonnement actuel soit sur la bonne voie et prometteur. Ce signal unique pilote trois composants coordonnés :

2. Compression du cache KV adaptative à la récompense

Contrairement à la compression statique ou uniforme, ReCo ajuste dynamiquement le ratio de rétention $\lambda_i$ en fonction de la récompense de l'étape $v_i$ :

• Étapes à haute récompense : Lorsque le modèle suit un chemin fiable ( $v_i$ élevé), le cadre élague agressivement le cache KV accumulé, car l'état de raisonnement tolère davantage de perte de contexte.
• Étapes à faible récompense : Lorsque le modèle explore ou est incertain ( $v_i$ faible), le cadre conserve plus de contexte pour éviter de briser la chaîne de raisonnement.
• Mécanisme : Le ratio de rétention est calculé comme $\lambda_i = \lambda - \delta \frac{v_i - \bar{v}}{v_{max} - v_{min}}$, où $\lambda$ est le ratio de base et $\delta$ limite l'ajustement. Dans le budget alloué, les jetons sont sélectionnés sur la base des scores d'attention des jetons récents.

3. Contrôle de réflexion par bandes de récompense

Pour prévenir l'« inflation de longueur » causée par la compression, ReCo supprime les jetons d'auto-réflexion redondants (ex: « Attendez », « Hmm », « Alternativement ») pendant le décodage :

• La même récompense $v_i$ détermine une pénalité de logit $\beta_i$ appliquée aux jetons de réflexion.
• Faible récompense ($v_i \le \tau_\ell$) : Aucune pénalité (encourage l'exploration).
• Récompense moyenne ($\tau_\ell < v_i < \tau_h$) : Demi-pénalité.
• Haute récompense ($v_i \ge \tau_h$) : Pleine pénalité (décourage fortement la réouverture de branches, orientant le modèle vers une conclusion).

4. Arrêt précoce basé sur la confiance

ReCo implémente un mécanisme d'arrêt précoce qui se déclenche lorsque le raisonnement est jugé fiable :

• Si la récompense reste dans la bande haute pendant deux étapes consécutives, un « prompt de clôture » est inséré pour solliciter une réponse provisoire.
• La confiance du modèle dans cette réponse est mesurée via la perplexité (PPL). Si $PPL(a) \le \tau_p$, la chaîne de raisonnement est interrompue ; sinon, le raisonnement continue. Cela évite un arrêt prématuré sur des réponses fluides mais incorrectes.

Principales contributions

1. Couplage de la compression et de la génération : L'article soutient que la compression du cache KV seule est insuffisante pour les LRM car elle peut être contre-productive en rallongeant le processus de raisonnement. ReCo est le premier à gouverner conjointement la compression du cache et la longueur de génération sous un signal de récompense unifié.
2. La récompense comme signal de tolérance à la compression : Les auteurs démontrent que les récompenses de processus suivent de manière fiable la tolérance d'un état de raisonnement à la perte de contexte. Les étapes à haute récompense sont robustes à un élagage agressif, tandis que les étapes à faible récompense nécessitent un contexte complet. Cela contraste avec les politiques uniformes existantes.
3. Cadre unifié par étape : ReCo intègre la compression du cache KV, le contrôle de la réflexion et l'arrêt précoce dans un cadre unique qui réduit le nombre de jetons et la latence tout en préservant la précision sur plusieurs modèles et benchmarks.

Résultats expérimentaux

Le cadre a été évalué sur trois modèles de raisonnement (DeepSeek-R1-Distill-Llama-8B, DeepSeek-R1-Distill-Qwen-7B, et Qwen3-8B) à travers six benchmarks (GSM8K, MATH-500, AMC2023, AIME24, AIME25, et GPQA).

• Réduction des jetons : ReCo a réduit les jetons générés de 37 % à 65 % par rapport au Full CoT. En revanche, les bases de référence de type « cache seul » (ex: SnapKV, R-KV) ont souvent augmenté le nombre de jetons en raison de l'inflation de longueur.
• Accélération de la latence : ReCo a obtenu des réductions de latence de bout en bout de 2,08× à 2,35× par rapport au Full CoT.
• Préservation de la précision : ReCo a maintenu la précision la plus proche du Full CoT parmi toutes les méthodes compressées. Par exemple, sur Llama-8B, ReCo a atteint une précision moyenne de 60,2 % contre 62,8 % pour le Full CoT, alors que les bases de référence de type « KV uniquement » ont chuté de manière significative (ex: SnapKV à 37,5 %).
• Efficacité de la mémoire : ReCo a réduit l'utilisation de la mémoire GPU de pointe de 37,7 % (17,92 Go contre 28,78 Go sur AIME25), surpassant à la fois les bases de référence de type « KV uniquement » et « longueur uniquement ».

Signification et affirmations

L'article affirme que la compression du cache KV doit être couplée au contrôle de la génération pour être efficace sur les modèles de raisonnement. Les auteurs soutiennent que compresser le cache sans contrôler la longueur de génération suivante peut éroder les gains d'efficacité en induisant des chaînes de raisonnement plus longues.

En utilisant une récompense de processus unique pour coordonner les deux aspects, ReCo atteint un équilibre où le calcul est dépensé uniquement lorsque le raisonnement est productif et retenu lorsque le modèle s'est stabilisé. Les auteurs soulignent que leur méthode ne nécessite aucun réentraînement du modèle de raisonnement ni d'ingénierie de prompt spécifique à la tâche, s'appuyant plutôt sur un estimateur de récompense léger pour s'adapter dynamiquement aux besoins de la trajectoire de raisonnement.