Résumé Technique : OS-Pruner

Énoncé du Problème

Les grands modèles de langage (LLM) utilisant le prompting par Chaîne de Pensée (Chain-of-Thought - CoT) ont connu un succès significatif dans les tâches de raisonnement complexe. Cependant, ces modèles font souvent preuve d'une « sur-réflexion computationnelle » (computational overthinking), générant des étapes de raisonnement redondantes qui augmentent la latence et les coûts d'inférence sans améliorer la précision de la réponse finale. Ce comportement inclut la reformulation d'arguments, la vérification de conclusions déjà établies ou la poursuite de calculs à faible valeur ajoutée. Bien que des études récentes suggèrent que les trajectoires de CoT peuvent être élaguées, les solutions existantes souffrent de limitations spécifiques :

• Budgets Statiques : Les méthodes comme l'imposition de budget imposent une limite de réflexion fixe de manière externe, ignorant l'évolution de l'état de la trace de raisonnement.
• Compression du côté du Modèle : Les techniques telles que le réglage fin (fine-tuning) pénalisant la longueur ou le saut de jetons (token skipping) nécessitent un réentraînement coûteux du modèle de base et peuvent dégrader les performances sur des tâches de non-raisonnement.
• Arrêt Précoce Heuristique : Les méthodes dynamiques actuelles reposent souvent sur la classification ou le seuillage (par exemple, arrêter lorsque la confiance dépasse une valeur fixe), ce qui ne tient pas compte de l'arbitrage entre le coût des jetons supplémentaires et l'amélioration attendue de la qualité de la réponse.

Les auteurs soutiennent que l'élagage de la CoT est fondamentalement un problème de décision séquentielle qui s'apparente le mieux à un problème d'arrêt optimal, où le système doit décider à chaque étape si le gain attendu en précision justifie la génération d'un jeton supplémentaire.

Méthodologie : OS-Pruner

OS-Pruner est un cadre léger et modulaire (plug-in) conçu pour formuler l'élagage de la CoT comme un problème d'arrêt optimal. Il fonctionne sans modifier le modèle de raisonnement de base qui reste gelé.

1. Formulation du Problème

Étant donné un problème d'entrée $x$ et un préfixe de CoT partiellement généré $y_{\le i}$, OS-Pruner apprend une politique d'arrêt paramétrique $\pi_\theta$ qui produit la probabilité d'arrêt. L'objectif est de maximiser une fonction d'utilité explicite qui arbitre l'exactitude de la réponse finale contre la longueur générée :
$$r(y_{\le i}|x) = A(y_{\le i}|x) - \lambda L(y_{\le i})$$
Où :

• $A(y_{\le i}|x)$ est l'exactitude de la réponse dérivée du préfixe (1 si correct, 0 sinon).
• $L(y_{\le i})$ est le nombre de jetons générés jusqu'à l'étape $i$.
• $\lambda$ est un paramètre scalaire contrôlant la pénalité de longueur, permettant aux utilisateurs de sélectionner un point de fonctionnement spécifique sur la frontière exactitude-efficacité.

2. Avantage Théorique

Le document distingue OS-Pruner des méthodes de seuil de confiance via une Vue de la Fonction de Valeur.

• Les méthodes basées sur la classification s'arrêtent si $A(y_{\le i}|x) \ge \gamma$ (un seuil fixe).
• L'Arrêt Optimal (OS-Pruner) s'arrête si la valeur immédiate de l'arrêt excède la valeur de continuation : $A(y_{\le i}|x) \ge C_\lambda(x, y_{\le i})$.
  La valeur de continuation $C_\lambda$ représente la récompense future attendue moins le coût des jetons supplémentaires. Les auteurs prouvent (Théorème 1) que les classificateurs à seuil fixe peuvent présenter un écart de performance arbitrairement important par rapport à la politique d'arrêt optimale, car ils ne tiennent pas compte de la valeur dépendante de l'état des futurs chemins de raisonnement.

3. Implémentation et Architecture

• Prétraitement des Données : Le modèle de base génère des traces de raisonnement complètes. Pour l'entraînement, le système intervient après chaque étape de raisonnement, force le modèle à produire une réponse finale (empêchant la « réflexion » dans la section de la réponse) et l'évalue par rapport à la vérité de terrain via Math-Verify. Cela crée une supervision fine pour chaque préfixe.
• Architecture de la Politique : La politique est une tête linéaire légère ($\sigma(W^\top h_l + b)$) appliquée au dernier état caché $h_l$ du préfixe. Pour améliorer les performances, les $n$ dernières couches (typiquement 2) d'auto-attention du modèle de base sont soumises à un réglage fin, tandis que le reste du modèle demeure gelé.
• Efficacité de l'Entraînement : Comme le modèle de base est causal, l'état caché d'un paragraphe n'est pas affecté par les jetons suivants. Cela permet de traiter l'intégralité de la trace de raisonnement en un seul passage avant (forward pass), la couche linéaire étant appliquée en parallèle aux états cachés pertinents.
• Entraînement Progressif : Pour éviter de rester bloqué dans des minima locaux où la politique ne s'arrêterait jamais, l'entraînement commence avec un $\lambda$ élevé (coût élevé pour la longueur) et l'amorce progressivement vers des valeurs plus basses.

4. Inférence

Lors de l'inférence, la politique est invoquée uniquement aux limites des étapes de raisonnement (par exemple, à la fin d'un paragraphe). Le système vérifie la probabilité d'arrêt ; si la décision est de s'arrêter, il déclenche immédiatement la réponse finale. Cette conception garantit une haute efficacité sur les systèmes de service de LLM modernes (ex: vLLM) en exploitant le traitement de pré-remplissage (pre-fill) parallèle et en minimisant la recomputation du cache KV.

Contributions Clés

1. Formulation Rigoureuse : Le document fournit une formalisation de l'élagage de la CoT comme un problème d'arrêt optimal avec une fonction de récompense explicite arbitrant l'exactitude et la longueur.
2. Plug-in Léger : OS-Pruner est une politique d'arrêt basée sur l'entraînement qui opère sur les préfixes de raisonnement sans modifier le modèle de raisonnement de base gelé, ne nécessitant que le réglage fin d'une petite tête de politique et de quelques couches de transformateur.
3. Performance Empirique : Les expériences démontrent qu'OS-Pruner réduit substantiellement la longueur du raisonnement tout en préservant l'exactitude sur divers modèles et benchmarks.
4. Complémentarité : Il est démontré que la méthode apporte des gains supplémentaires même lorsqu'elle est appliquée à des modèles déjà optimisés pour la concision via des méthodes côté modèle.

Résultats Expérimentaux

Les auteurs ont évalué OS-Pruner sur trois modèles (DeepSeek-R1-Distill-Qwen-7B, GPT-OSS-20B, et DRPO-7B) sur des benchmarks de difficultés variées (GSM8K, MATH-500, AIME).

• Réduction de la Longueur : OS-Pruner a obtenu une réduction de 20 à 60 % de la longueur de génération avec un sacrifice minimal de l'exactitude.
  • Sur DeepSeek-R1-Distill-Qwen-7B, il a réduit la longueur de 59,3 % sur GSM8K et de 52,8 % sur MATH-500, avec des changements d'exactitude de seulement -0,7 % et +2,7 % respectivement.
  • Sur GPT-OSS-20B (un modèle déjà optimisé pour la concision), il a tout de même obtenu des réductions significatives (ex: 21 % sur MATH-500) tout en maintenant une exactitude compétitive.
• Comparaison aux Références (Baselines) : OS-Pruner a systématiquement surpassé les bases d'entraînement (Classification Probe, Answer Convergence, FlashThink) dans le score d'Exactitude-Efficacité (AES). Il se situe sur ou près de la frontière de Pareto, alors que les bases à seuil fixe dégradaient plus brutalement à mesure que les budgets de longueur étaient réduits.
• Robustesse : La méthode a bien performé à travers différents niveaux de difficulté, étant plus agressive sur les tâches faciles (où la sur-réflexion est prédominante) et plus conservatrice sur les tâches difficiles (AIME), où le calcul supplémentaire est précieux.

Signification et Revendications

Le document affirme qu'OS-Pruner offre une approche plus naturelle et optimale de l'élagage de la CoT que les méthodes heuristiques ou statiques existantes. En formulant le problème comme un processus de décision séquentiel, il évalue dynamiquement le point de terminaison suffisant en fonction de l'état de raisonnement spécifique plutôt que par un seuil global.

Les auteurs soulignent qu'OS-Pruner :

• Réduit le Coût Computationnel : En coupant les jetons inutiles, il diminue la latence d'inférence et l'empreinte énergétique.
• Fournit un Contrôle Fin : Le paramètre $\lambda$ permet aux utilisateurs de régler explicitement l'arbitrage entre l'effort de raisonnement et l'exactitude.
• Fonctionne avec les Modèles Existants : Il ne nécessite pas de réentraînement du générateur de base, ce qui en fait un « plug-in » pratique pour les modèles de raisonnement existants.
• Complète les Méthodes Côté Modèle : Même les modèles entraînés pour la concision (comme DRPO) bénéficient de l'élagage supplémentaire, suggérant que la compression côté modèle et l'arrêt optimal traitent des aspects différents de l'efficacité.

Le travail conclut que, bien qu'il se concentre sur le raisonnement mathématique, le cadre est applicable à des contextes plus larges, bien que des travaux futurs soient nécessaires pour explorer les tâches de codage et de raisonnement scientifique ainsi que le passage à l'échelle vers les modèles de pointe (frontier models).