TRIM: Hybrid Inference via Targeted Stepwise Routing in Multi-Step Reasoning Tasks

1. Problématique

Les tâches de raisonnement multi-étapes (comme la résolution de problèmes mathématiques complexes) sont vulnérables aux échecs en cascade : une seule erreur à une étape critique peut entraîner l'échec complet de la solution.
Les méthodes de routage actuelles pour les grands modèles de langage (LLM) opèrent au niveau de la requête entière : elles attribuent une question à un seul modèle (soit un modèle petit et peu coûteux, soit un modèle grand et performant). Cette approche présente deux défauts majeurs :

• Inefficacité des coûts : Elle suppose que le modèle puissant est nécessaire à chaque token, alors que la plupart des étapes sont des continuations routinières.
• Rigidité : Elle ne permet pas d'intervenir spécifiquement là où le raisonnement dérape, ce qui est crucial dans les tâches séquentielles.

L'objectif est de trouver un équilibre entre la qualité de la réponse et le coût d'inférence (nombre de tokens générés par le modèle coûteux) en intervenant de manière ciblée.

2. Méthodologie : TRIM

Les auteurs proposent TRIM (Targeted routing in multi-step reasoning tasks), un cadre d'inférence hybride qui opère au niveau de chaque étape de raisonnement plutôt qu'au niveau de la requête.

Fonctionnement Global

Le système utilise deux modèles :

• $M_w$ (Weak/Cheap) : Un modèle petit et peu coûteux qui génère la majorité des étapes.
• $M_s$ (Strong/Expensive) : Un modèle grand et performant utilisé uniquement pour régénérer les étapes critiques.

À chaque étape $t$ du processus de génération :

1. Le modèle $M_w$ propose une continuation.
2. Un Routeur évalue cette étape en utilisant un Modèle de Récompense de Processus (PRM) pour estimer la probabilité de correction de l'étape.
3. Le routeur décide de l'action :
   • Continuer : Accepter la sortie de $M_w$.
   • Régénérer : Remplacer l'étape par une génération de $M_s$ (intervention ciblée).

Stratégies de Routage Proposées

Le papier explore plusieurs politiques de décision au sein de TRIM :

1. TRIM-Thr (Seuil Myope) :

   • Une politique simple qui compare le score du PRM de l'étape courante à un seuil fixe $k$.
   • Si le score < $k$, on régénère avec $M_s$. Sinon, on continue avec $M_w$.
   • C'est une approche locale qui ne considère pas l'historique complet.

2. Politiques Entraînées par RL (TRIM-Seq et TRIM-Agg) :

   • TRIM-Seq : Utilise une séquence complète de caractéristiques (scores de correction et nombre de tokens de chaque étape) pour apprendre une politique via Renforcement Learning (PPO). Elle prend en compte le contexte séquentiel.
   • TRIM-Agg : Utilise une représentation agrégée (ex: score minimum des étapes précédentes, score actuel, coût) pour réduire la complexité. Elle apprend à arbitrer entre la précision à long terme et le coût, montrant des performances proches de l'oracle avec moins de tokens coûteux.

3. TRIM-POMDP (Approche basée sur les MDP Partiellement Observables) :

   • Reconnaît que les scores du PRM sont bruyants et imparfaits.
   • Formule le problème comme un POMDP où l'état latent (vraie correction de la trajectoire) n'est pas directement observable, mais estimé via les observations bruyantes du PRM.
   • L'état latent est catégorisé en : $S_0$ (trajectoire correcte), $S_1$ (erreur irrécupérable), $S_2$ (erreur récente mais récupérable).
   • Utilise un solveur (SARSOP) pour calculer une politique optimale qui maximise la récompense finale tout en minimisant le coût, même avec des estimations incertaines.

3. Contributions Clés

1. Preuve de concept de l'intervention ciblée : Démonstration qu'un petit nombre d'interventions stratégiques aux étapes critiques suffit à corriger les trajectoires de raisonnement, évitant ainsi le besoin d'utiliser un modèle puissant pour l'intégralité de la génération.
2. Cadre flexible et généralisable : Développement de plusieurs politiques (du seuil simple aux approches RL et POMDP) qui s'adaptent à différents budgets de coûts.
3. Généralisation transversale : Les politiques entraînées sur un jeu de données (ex: AIME) se généralisent remarquablement bien à d'autres benchmarks (OlympiadBench, Minerva Math), suggérant que les motifs de difficulté au niveau des étapes sont des caractéristiques fondamentales du raisonnement et non spécifiques à un dataset.
4. Efficacité computationnelle : L'architecture permet d'amortir les coûts de pré-remplissage (prefill) via des techniques de décodage spéculatif, rendant l'overhead de routage négligeable.

4. Résultats Expérimentaux

Les évaluations ont été menées sur des benchmarks mathématiques difficiles : MATH-500, AIME, OlympiadBench et Minerva Math.

• Efficacité Coût-Performance :

  • Sur MATH-500, la stratégie simple TRIM-Thr surpasse les méthodes de routage par requête existantes avec une efficacité de coût 5 fois supérieure.
  • Les politiques avancées (TRIM-Agg, TRIM-POMDP) atteignent les performances du modèle coûteux ($M_s$) en utilisant 80 % de moins de tokens coûteux.
  • Sur AIME (benchmark plus difficile), TRIM atteint une efficacité de coût allant jusqu'à 6 fois supérieure aux méthodes de base.

• Généralisation :

  • Une politique entraînée uniquement sur AIME obtient une efficacité de coût 11,68 fois supérieure sur OlympiadBench, surpassant largement les méthodes de routage par requête (comme RouteLLM) qui souffrent d'un sur-ajustement aux spécificités des datasets.

• Robustesse au bruit :

  • TRIM-POMDP excelle particulièrement dans les régimes à faible budget et sur des datasets bruyants, car il modélise explicitement l'incertitude des scores du PRM, contrairement aux méthodes basées sur des seuils fixes.

5. Signification et Impact

Ce travail marque un changement de paradigme dans l'optimisation de l'inférence des LLM :

• Passage du niveau "Requête" au niveau "Étape" : Il démontre que l'hétérogénéité de la difficulté au sein d'une même réponse est la clé pour optimiser les coûts.
• Optimisation des ressources : TRIM permet d'utiliser des modèles très performants (et coûteux) de manière économiquement viable pour des applications réelles, en ne les sollicitant que lorsque c'est strictement nécessaire pour éviter un échec catastrophique.
• Fondement théorique : L'utilisation de POMDP pour gérer l'incertitude des modèles de récompense ouvre la voie à des systèmes de raisonnement plus robustes et adaptatifs, capables de fonctionner même lorsque les signaux de vérification sont imparfaits.

En résumé, TRIM transforme l'inférence hybride d'une simple sélection de modèle en un processus de décision séquentiel dynamique, offrant des gains d'efficacité drastiques sans sacrifier la qualité des réponses sur des tâches complexes.