DRP: Distilled Reasoning Pruning with Skill-aware Step Decomposition for Efficient Large Reasoning Models

1. Énoncé du problème

Les Modèles de Raisonnement à Grande Échelle (LRM), tels que o1 d'OpenAI et DeepSeek-R1, excellent dans des tâches complexes en générant de longues trajectoires de raisonnement de type Chaîne de Pensée (CoT). Cependant, ils souffrent de « surréflexion », produisant des sorties excessivement verbeuses avec des étapes redondantes, des retours en arrière inutiles et des boucles répétitives. Cela entraîne :

• Des coûts d'inférence élevés : Une utilisation importante de tokens et une latence accrue.
• Une dégradation de la précision : Un raisonnement redondant peut égarer le modèle vers des conclusions incorrectes.
• Un décalage de distillation : Lors de la distillation de connaissances d'un enseignant « Short-CoT » (concis, poli) vers un élève « Long-CoT » (verbeux, réflexif), un écart de learnabilité apparaît. L'élève peine à imiter directement le style concis de l'enseignant, ce qui se traduit souvent par de mauvaises performances ou une baisse de précision.

Les solutions existantes, comme l'élagage au moment de l'inférence (arrêt prématuré de la génération), risquent de stopper le raisonnement trop tôt, tandis que la distillation standard échoue souvent à combler le fossé structurel entre les élèves verbeux et les enseignants concis.

2. Méthodologie : Distilled Reasoning Pruning (DRP)

DRP est un cadre hybride qui unifie l'élagage au moment de l'inférence avec la distillation par ajustement fin supervisé (SFT). Au lieu d'entraîner simplement l'élève sur la réponse finale de l'enseignant ou sur un CoT court, DRP utilise l'enseignant pour élaguer et affiner les propres traces de raisonnement verbeuses de l'élève.

Le processus se compose de trois étapes (illustrées dans la Figure 2 du papier) :

A. Décomposition des étapes basée sur les compétences

• Le modèle élève génère une trace CoT initiale et verbeuse ($T$).
• Un LLM Enseignant (par exemple, GPT-4o) décompose cette trace en segments non chevauchants.
• Innovation clé : Contrairement à une simple segmentation de phrases, l'enseignant attribue une étiquette de compétence fonctionnelle à chaque étape (par exemple, « Lecture de la quantité donnée », « Représentation algébrique », « Inférence logique»). Cela garantit que les étapes sont cohérentes sémantiquement et alignées fonctionnellement.

B. Élagage et réécriture au niveau des étapes

Le modèle enseignant évalue chaque étape étiquetée par compétence et applique l'une des quatre actions suivantes :

1. Conserver : Garder l'étape inchangée.
2. Supprimer : Retirer les étapes redondantes, non informatives ou de retouche.
3. Réécrire : Condenser l'étape tout en préservant l'intention logique.
4. Fusionner : Combiner des étapes adjacentes en une seule unité atomique.

• L'enseignant reconstruit ensuite une trace de raisonnement fluide et concise ($\hat{T}$) qui maintient le ton original de l'élève mais élimine la verbosité. Il met également à jour le résumé de la réponse finale ($\hat{A}$) si nécessaire.

C. Ajustement fin supervisé (SFT)

• Le modèle élève est affiné sur des paires de $(Entrée, \hat{R})$, où $\hat{R}$ est la réponse élaguée et affinée par l'enseignant.
• L'objectif est de minimiser la vraisemblance négative de la trace affinée, enseignant à l'élève à intérioriser des stratégies de raisonnement alignées sur les compétences et économes en tokens, sans perdre la profondeur structurelle requise pour les problèmes complexes.

3. Contributions clés

1. Cadre DRP : Une méthode novatrice qui comble l'« écart de learnabilité » en élaguant le chemin de raisonnement de l'élève à l'aide d'un enseignant, plutôt que de forcer l'élève à imiter un style de raisonnement complètement différent de celui de l'enseignant.
2. Décomposition consciente des compétences : Introduit un mécanisme pour segmenter le raisonnement en unités cohérentes sémantiquement et étiquetées fonctionnellement. Cela fournit une base stable pour un élagage précis, surpassant la segmentation naïve.
3. Insight sur la cohérence structurelle : Démontre que les CoT d'entraînement efficaces doivent être à la fois informatifs et structurellement cohérents avec le processus de raisonnement de l'élève. Préserver la structure étape par étape (même si elle est raccourcie) est plus efficace qu'une compression naïve.

4. Résultats expérimentaux

Les auteurs ont évalué DRP sur les modèles DeepSeek-R1-Distill-Qwen (1,5B et 7B) à travers des benchmarks in-domain (GSM8K) et out-of-distribution (OOD) (MATH500, AIME24, AMC23).

Principales constatations :

• Efficacité des tokens : DRP a considérablement réduit l'utilisation de tokens sans sacrifier la précision.
  • GSM8K (7B) : Tokens réduits de 917 à 328 (-64 %), tandis que la précision est passée de 91,7 % à 94,1 %.
  • AIME24 (7B) : Tokens réduits de 43 % sans baisse de précision (15/30 corrects).
  • AMC (1,5B) : Tokens réduits de 4574 à 3657, avec une précision passant de 22/40 à 27/40.
• Comparaison avec les références :
  • vs. Basé sur les prompts (TALE) : DRP a surpassé TALE sur des tâches OOD complexes où les performances de TALE chutaient significativement.
  • vs. Basé sur SFT (CoT-Valve) : DRP a obtenu une meilleure précision et compression. CoT-Valve a souvent subi des compromis de précision lors de la compression.
  • vs. Basé sur RL (ThinkPrune) : Bien que ThinkPrune compresse bien, DRP a atteint une précision plus élevée, probablement grâce à un guidage enseignant de haute qualité.
• Études d'ablation :
  • Décomposition des compétences : Le retrait des étiquettes de compétences (segmentation par défaut) a réduit les performances, confirmant que l'alignement sémantique est crucial.
  • Distillation directe : L'entraînement direct sur des CoT courts d'enseignant (sans élaguer la trace de l'élève) a entraîné une chute brutale de la précision sur les tâches OOD, prouvant que la préservation structurelle est vitale.
  • Robustesse de l'enseignant : La méthode a fonctionné efficacement avec divers enseignants (GPT-4o, Gemini 2.0, DeepSeek-V3), bien que GPT-4o ait donné les meilleurs résultats.

5. Importance et conclusion

• Compromis Efficacité vs Précision : DRP remet en question l'idée que l'efficacité doit se faire au détriment de la précision. En affinant le processus de raisonnement plutôt que de simplement raccourcir la sortie, il atteint les deux.
• Évolutivité : La méthode est particulièrement efficace pour les petits modèles (1,5B), compensant une capacité limitée grâce à une supervision structurée de haute qualité.
• Robustesse : Elle atténue les « boucles dégénérées » (où les modèles restent bloqués à générer un texte infini), conduisant à une inférence plus stable.
• Impact futur : Ce travail fournit un modèle pour entraîner des modèles de raisonnement efficaces capables de gérer des tâches complexes à long horizon sans la surcharge computationnelle de la « surréflexion », rendant le raisonnement avancé plus accessible et rentable.

En résumé, DRP résout l'inefficacité des Modèles de Raisonnement à Grande Échelle en utilisant un enseignant pour éliminer chirurgicalement la redondance du propre chemin de raisonnement d'un élève, guidé par des compétences fonctionnelles, aboutissant à des modèles plus rapides, moins chers et souvent plus précis.