Résumé technique : Vérification intégrée aux outils pour la mise à l'échelle du calcul au moment de l'inférence dans les petits modèles de langage

Énoncé du problème

Les avancées récentes dans la mise à l'échelle du calcul au moment de l'inférence (test-time compute scaling) ont démontré qu'allouer un calcul supplémentaire lors de l'inférence peut considérablement améliorer les capacités de raisonnement des petits modèles de langage (sLM). Cependant, les approches existantes de mise à l'échelle reposent fortement sur de grands modèles de langage (LLM) agissant comme vérificateurs (par exemple, des modèles de récompense de processus ou des modèles critiques) pour sélectionner la meilleure solution parmi plusieurs candidats. Cette dépendance vis-à-vis de grands vérificateurs compromet l'efficacité et les avantages de coût liés à l'utilisation des sLM.

Un goulot d'étranglement critique identifié dans ce travail est l'incapacité des sLM à vérifier de manière fiable leurs propres sorties générées, particulièrement dans les tâches nécessitant la mémorisation de faits arithmétiques ou de connaissances factuelles. Alors que les modèles plus grands peuvent vérifier les solutions en utilisant uniquement le raisonnement par chaîne de pensée (chain-of-thought), les sLM peinent même pour des calculs simples à mesure que la complexité augmente. Les auteurs émettent l'hypothèse que cet échec provient de la capacité limitée des sLM à mémoriser le vaste nombre de faits de calcul requis pour la vérification, plutôt que d'un manque de capacité de raisonnement logique.

Méthodologie : Vérification intégrée aux outils (T1)

Pour remédier au goulot d'étranglement de la mémorisation, les auteurs proposent la Vérification Intégrée aux Outils (T1), un cadre à deux étapes conçu pour décharger les étapes lourdes en mémorisation vers des outils externes tout en conservant le sLM pour la vérification finale.

Le cadre à deux étapes

1. Étape de filtrage basée sur l'outil (ToolV) :
   • Le sLM génère du code exécutable (par exemple, Python) ou des requêtes de recherche basés sur la solution candidate.
   • Un outil externe (tel qu'un interprète de code ou un récupérateur de connaissances) exécute le code ou récupère les faits pour vérifier l'exactitude numérique et la cohérence factuelle.
   • Cette étape agit comme un filtre binaire ($f(x, y; T, \theta) \in \{0, 1\}$), écartant les solutions présentant des erreurs de calcul ou des faits incorrects.
2. Étape de notation par modèle de récompense :
   • Les solutions filtrées restantes sont notées par un modèle de récompense (soit un modèle de récompense de processus [PRM], soit un modèle de récompense génératif [GenRM]).
   • La sortie finale est sélectionnée comme étant le candidat ayant le score le plus élevé du modèle de récompense.

Distillation de connaissances

Pour améliorer les performances des deux étapes, les auteurs emploient la distillation de connaissances à partir d'un modèle enseignant plus grand (par exemple, GPT-4o-mini).

• Distillation ToolV : Le sLM est affiné (fine-tuned) pour générer des requêtes d'appel d'outils (code) et extraire les résultats de vérification basés sur les sorties de l'enseignant.
• Distillation RM : Le sLM est affiné pour générer des rationales de vérification et des scores, imitant l'évaluation de l'enseignant.
• Multi-LoRA : Des adaptateurs LoRA (Low-Rank Adaptation) distincts sont utilisés pour les tâches basées sur les outils et les tâches de récompense afin de gérer efficacement les objectifs distincts lors de l'apprentissage par affinement.

Analyse théorique

L'article fournit des garanties théoriques soutenant l'approche T1 :

1. Réduction de la mémorisation : En utilisant l'analyse de la théorie de l'information, les auteurs prouvent que sans outils, un vérificateur optimal doit mémoriser un ensemble de données de taille proportionnelle à $O(M^3)$ (où $M$ est la plage de nombres) pour atteindre une erreur proche de zéro. En utilisant un outil externe pour effectuer le calcul, l'information mutuelle requise pour être stockée dans les paramètres du modèle tombe à zéro, car l'outil gère la logique de vérification.
2. Performance de mise à l'échelle : Les auteurs modélisent le vérificateur comme un filtre imparfait avec du bruit. Ils démontrent que l'utilisation d'un filtre basé sur un outil augmente la probabilité d'identifier correctement l'étiquette de vérité terrain dans un processus de sélection "Best-of-N". Théoriquement, cela améliore la monotonicité de la performance de mise à l'échelle au moment de l'inférence, garantissant qu'augmenter le nombre de solutions générées ($N$) produit de meilleurs résultats lorsqu'un filtre fiable est présent.

Résultats expérimentaux

Les expériences ont été menées sur des benchmarks de raisonnement mathématique (MATH500, GSM8K) et des tâches intensives en connaissances (sous-ensemble MMLU-Pro).

• Gains de performance : T1 permet aux sLM de surpasser significativement des modèles beaucoup plus grands. Plus précisément, un modèle Llama-3.2 1B utilisant T1 sous mise à l'échelle au moment de l'inférence surpasse le modèle Llama-3.1 8B sur le benchmark MATH500.
• Intégration avec les vérificateurs : T1 améliore la précision des PRM et des GenRM. Combiné à un PRM distillé, T1 permet à un modèle de 1B d'égaler ou de dépasser la performance d'un modèle de 8B avec moins de solutions générées.
• Efficacité du filtrage : L'analyse montre que ToolV filtre efficacement les solutions incorrectes causées par des erreurs de calcul, que les vérificateurs génératifs seuls omettent souvent.
• Efficacité de calcul : Même en tenant compte de la surcharge de jetons (token overhead) pour la génération de code (une augmentation d'environ 1,14x du budget de calcul), T1 offre une performance supérieure par rapport à une simple augmentation de la taille du vérificateur ou à l'utilisation de modèles plus grands sans outils.
• Efficacité des données : Le vérificateur basé sur l'outil maintient des performances compétitives même lorsqu'il est distillé avec seulement 10 % des données d'entraînement, soulignant son efficacité.

Signification et revendications

L'article affirme que l'intégration d'outils n'est pas seulement bénéfique mais nécessaire pour permettre une vérification réussie dans les sLM sous mise à l'échelle au moment de l'inférence. Les principales contributions sont :

1. Identification du goulot d'étranglement : Identification systématique des étapes lourdes en mémorisation comme le principal point de défaillance pour la vérification des sLM.
2. Proposition de cadre : Introduction de T1, un cadre à deux étapes fluide qui intègre des outils externes avec des sLM pour les vérificateurs génératifs et les PRM.
3. Justification théorique : Preuve que le déchargement vers des outils réduit la charge de mémorisation et garantit théoriquement une amélioration de la performance de mise à l'échelle.
4. Validation empirique : Démonstration que de petits modèles (par exemple, 1 milliard de paramètres) peuvent atteindre des capacités de raisonnement comparables à des modèles beaucoup plus grands (par exemple, 8 milliards+) en exploitant des outils, préservant ainsi l'efficacité de déploiement des sLM.

Les auteurs concluent que T1 offre une voie pratique pour améliorer la performance des sLM dans les tâches de raisonnement complexes sans le coût prohibitif de l'augmentation des paramètres du modèle, à condition que la charge de vérification soit transférée à des outils externes.