Résumé technique : TrajFusion pour le passage à l'échelle du raisonnement mathématique

Problématique

Les grands modèles de langage (LLM) ont réalisé des progrès significatifs dans le raisonnement mathématique, principalement grâce au réglage fin supervisé (SFT) sur des trajectoires de solutions de haute qualité. Le paradigme dominant pour la construction de ces données est le Réajustement par Échantillonnage de Rejet (RFT - Rejection Sampling Fine-Tuning). Dans le RFT, un modèle enseignant génère plusieurs chemins de raisonnement candidats pour un problème donné, et seuls ceux ayant les réponses finales correctes sont conservés pour l'entraînement ; toutes les trajectoires incorrectes sont rejetées.

Bien qu'efficace, cette approche traite la supervision comme un filtre binaire. Elle exclut systématiquement les erreurs générées par l'enseignant, créant un fossé dans la manière dont les échecs de raisonnement sont modélisés pendant l'entraînement. Les auteurs soutiennent que les trajectoires incorrectes contiennent des informations diagnostiques précieuses — telles que des erreurs de calcul courantes, des hypothèses manquantes ou des chemins d'argumentation logiquement plausibles mais invalides — qui pourraient aider les modèles à reconnaître et à éviter des modes d'échec spécifiques, particulièrement dans les tâches de raisonnement de longue forme ou ambiguës. Les méthodes actuelles traitent ces erreurs comme du bruit à éliminer plutôt que comme une supervision structurée à exploiter.

Méthodologie : TrajFusion

Le papier propose TrajFusion, une stratégie de réglage fin qui reformule le RFT d'un mécanisme de filtrage binaire en un processus de construction de supervision structurée. TrajFusion intègre les trajectoires de raisonnement incorrectes sélectionnées avec les trajectives correctes dans un seul échantillon d'entraînement fusionné, modélisant explicitement un processus de raisonnement par essais et erreurs.

Composants clés

1. Échantillonnage et Analyse : Pour chaque problème $x$, le modèle enseignant échantillonne $K$ trajectoires. Celles-ci sont partitionnées en ensembles corrects ($Y^+$) et incorrects ($Y^-$) à l'aide d'un vérificateur automatique.
2. Sélection Adaptative : La méthode détermine de manière adaptative le nombre de trajectoires incorrectes à inclure en se basant sur deux statistiques :
   • Taux d'erreur ($r(x)$) : La fréquence des tentatives incorrectes.
   • Diversité des erreurs ($u(x)$) : Le nombre de réponses finales uniques parmi les tentatives incorrectes (mesuré par l'entropie de Shannon ou le nombre de réponses uniques).
     Le nombre de trajectoires incorrectes conservées $k(x)$ est calculé comme suit :
     $$k(x) = \min(k_{max}, \lfloor \alpha \cdot r(x) \cdot u(x) \rfloor)$$
     Cela garantit que les problèmes présentant des erreurs diverses et fréquentes reçoivent une supervision plus riche, tandis que les problèmes aux erreurs rares ou peu informatives dégénèrent naturellement vers un RFT classique.
3. Construction d'Échantillons Fusionnés : Les trajectoires incorrectes sélectionnées sont entrelacées avec la trajectoire correcte en utilisant des prompts de réflexion fixes (ex: « Attendez... cela ne semble pas correct. Réessayons. »). La structure finale est :
   $$T(x) = [y_{e1}, \rho_1, \dots, y_{ek}, \rho_k, y^*]$$
   où $y_{ei}$ sont les chemins incorrects, $\rho_i$ sont les prompts de réflexion, et $y^*$ est la solution correcte.
4. Entraînement : Les séquences fusionnées sont entraînées via la prédiction du prochain jeton (next-token prediction) standard, sans modification de l'architecture du modèle ou de l'objectif d'entraînement.

Principales Contributions

• Formalisation d'un fossé de supervision : Les auteurs identent et formalisent l'exclusion systématique des trajectoires incorrectes générées par l'enseignant dans le RFT, arguant que ces erreurs fournissent une valeur diagnostique pour la modélisation des échecs de raisonnement.
• Cadre TrajFusion : Un cadre de supervision structuré qui fusionne de manière adaptative les trajectoires incorrectes et correctes. Il traite la diversité des erreurs comme un signal de contrôle explicite, permettant à la méthode de passer à l'échelle de manière robuste sur des problèmes de difficulté variable sans nécessiter de modèles auxiliaires, de générateurs de critiques ou d'étiquettes de difficulté externes.
• Validation Empirique : La méthode démontre des améliorations constantes par rapport au RFT classique sur plusieurs benchmarks de raisonnement mathématique (GSM8K, MATH, CollegeMath, DeepMind-Mathematics, OlympiadBench, TheoremQA) et des jeux de données de niveau compétition (AIME24, AIME25, MATH-500).

Résultats Expérimentaux

Les auteurs ont évalué TrajFusion sur six benchmarks de raisonnement mathématique (GSM8K, MATH, CollegeMath, DeepMind-Mathematics, OlympiadBench, TheoremQA) et des jeux de données de niveau compétition (AIME24, AIME25, MATH-500).

• Gains de Performance : TrajFusion a systématiquement surpassé le RFT classique. Sur le jeu de données MATH avec LLaMA3-8B (100K données), TrajFusion a atteint 51,4 % de précision contre 48,5 % pour le RFT. Sur DeepSeekMath-7B, il a atteint 52,3 % contre 50,2 %.
• Impact sur la Complexité : Les améliorations ont été les plus marquées sur les benchmarks de raisonnement de longue forme et complexes (MATH, OlympiadBench, TheoremQA), où les chemins de raisonnement sont plus longs et les modes d'échec plus diversifiés. Les gains sur des jeux de données plus simples comme GSM8K étaient plus faibles mais restaient positifs.
• Efficacité des Données : TrajFusion a obtenu des résultats compétitifs avec seulement 15K exemples d'entraînement, surpassant des bases de référence utilisant des jeux de données nettement plus larges (par exemple, MMIQC avec 2,3 millions d'exemples).
• Passage à l'échelle du Contexte Long : Appliqué à une version à contexte de 32K de Qwen2.5-Math-7B utilisant des données de longue forme distillées de DeepSeek-R1, TrajFusion a amélioré le Pass@1 sur AIME24 de 40,0 % (RFT) à 43,3 %.
• Études d'Ablation :
  • L'augmentation du nombre d'erreurs fusionnées ($k_{max}$) a amélioré les performances jusqu'à un certain point (rendements décroissants observés à $k_{max}=6$).
  • La sélection des erreurs basée sur la diversité (HighDiv) a donné de meilleurs résultats que la sélection basée uniquement sur la fréquence (LowDiv), suggérant que les modes d'échec structurellement distincts sont plus informatifs.
  • Les modèles entraînés avec TrajFusion ont initialement généré des séquences plus longues, puis sont devenus plus concis tout en augmentant la fréquence des jetons de réflexion, indiquant une utilisation plus efficace des motifs structurels.

Signification et Revendications

Le papier affirme que TrajFusion offre un moyen simple et efficace de passer à l'échelle le raisonnement mathématique en exploitant des signaux de supervision généralement rejetés. En reformulant l'échantillonnage par rejet comme un processus de construction de trajectoires structurées d'essais et d'erreurs, la méthode fournit une supervision plus riche sans modifications architecturales ou pipelines complexes d'apprentissage par renforcement.

Les auteurs soulignent que TrajFusion est particulièrement précieux pour les problèmes de raisonnement complexes et de longue forme, où distinguer un raisonnement valide d'un raisonnement erroné est essentiel. Ils affirment que la nature adaptative de TrajFusion garantit sa robustesse : elle enrichit la supervision lorsque les erreurs sont informatives et revient en toute sécurité au comportement standard du RFT lorsqu'elles ne le sont pas. Cette approche comble le fossé entre le filtrage binaire de correction et la réalité nuancée des échecs de raisonnement, permettant aux modèles d'apprendre de leurs propres erreurs dans un cadre supervisé.