ASTRO Améliore la Raisonnerie de Llama 3 de 16% à 20% Sans Modification Architecturale

Les chercheurs de Meta AI et de l'Université de Washington ont développé ASTRO (Autoregressive Search-Taught Reasoner), une nouvelle méthode de post-formation visant à améliorer les capacités de raisonnement de Llama-3.1-70B-Instruct, sans modifier son architecture. ASTRO enseigne au modèle à effectuer des recherches contextualisées, à s'interroger et à revenir sur ses erreurs, des mécanismes souvent associés au raisonnement humain et auxalgorithmes de recherche symbolique traditionnels. Cette approche permet d'augmenter les performances mathématiques de Llama-3 sur plusieurs benchmarks concurrentiels, avec des gains allant de 16% à 20%. Génération de Chaînes de Pensée Guidée par la Recherche Le processus d'ASTRO débute par une recherche Monte Carlo Tree Search (MCTS) sur des trajectoires de résolution de problèmes mathématiques. Cette recherche explore à la fois les chemins corrects et incorrects. La principale innovation réside dans la procédure de clonage : les arbres de recherche entiers sont transformés en chaînes de pensée (CoT) linéaires, encodant naturellement les échecs et les réparations grâce à la réflexion et au retour en arrière. Ces traces linéarisées sont reformulées en langage naturel et servent alors de base pour un entrainement supervisé fin (SFT). Résultat : le modèle ne se contente pas de résoudre les problèmes étape par étape, mais il réévalue également sa démarche. Il est capable de rebrousser chemin après une auto-évaluation pour corriger des erreurs intermédiaires. Par exemple, lorsque sa confiance interne diminue, le modèle peut intercaler des phrases comme « Retournons à l’étape où nous avons établi l'équation ». Entrainement Supervisé Fin : Injection de Priors de Recherche L'ASTRO entraîne Llama-3.1-70B-Instruct sur 36 100 CoT de solutions réarrangées, issues de différents jeux de données comme MATH, AMC/AIME et AoPS. Ce modèle entraîné avec ASTRO-SFT obtient des scores compétitifs ou supérieurs à ceux des modèles de base et des variantes SPOC/Step-KTO, qui n’ont pas bénéficié de priorités de recherche explicites. Importamment, même l’entrainement supervisé fin (SFT) sans apprentissage par renforcement permet d’améliorer les performances, en exposant le modèle à des données de raisonnement structurées par la recherche. Apprentissage par Renforcement avec Initialisation Sensible à la Recherche L'étape suivante consiste en un apprentissage par renforcement (RL), initialisé avec le point de contrôle SFT, en utilisant une version modifiée de Group Relative Policy Optimization (GRPO). Contrairement à l'apprentissage par renforcement basé sur les préférences, ASTRO utilise des signaux de récompense vérifiables (+1 pour une réponse correcte, -1 pour une réponse incorrecte) sur 8 700 prompts de difficulté modérée. Au cours de la formation, les générations de CoT deviennent plus longues, passant d'environ 1 800 à 6 000 tokens, ce qui témoigne d'une exploration plus profonde. Le modèle ASTRO-RL resultant atteint des performances qui rivalisent voire surpassent celles des modèles avec des comptages de paramètres plus élevés, validant ainsi l’importance de l'initialisation sensiblement recherche-dire d'ASTRO. Fréquence de Retour en Arrière et Succès en Raisonnement Une observation empirique remarquable est la corrélation positive entre la fréquence de retour en arrière et les performances. Au fur et à mesure de la formation, ASTRO-RL montre des actions de correction plus fréquentes et une exploration plus approfondie. Les coefficients de corrélation de Pearson sur les diverses normes dépassent 0,8, indiquant que la réflexion et le retour en arrière ne sont pas de simples comportements superficiels, mais étroitement liés à une meilleure précision. Expériences Contrôlées et Impact Plus Large Des expériences contrôlées comparent ASTRO avec des modèles entraînés uniquement sur des solutions directes de CoT (sans priorités de recherche). Même si ces modèles sont formés sur les mêmes ensembles de problèmes et les mêmes arbres de recherche, ASTRO se distingue constamment. Par exemple, ASTRO-RL surpasse Direct-RL avec : MATH : +13% AMC : +15% AoPS : +16% De plus, les sorties d'ASTRO peuvent être visualisées sous forme de graphes dirigés, où chaque nœud représente une étape de raisonnement et chaque arête capture les transitions, les réflexions et les corrections, ce qui facilite une meilleure interprétatibilité des résultats. Conclusion ASTRO démontre qu’il est possible d'améliorer de manière significative le raisonnement des grands modèles linguistiques (LLMs) tels que Llama-3, non pas en augmentant leur taille ou en prolongeant leur préformation, mais en utilisant des techniques de post-formation bien fondées. En imitant les algorithmes de recherche en langage naturel, ASTRO permet aux modèles de réfléchir avant de répondre, de mettre en question leurs étapes et de se corriger pendant le processus de raisonnement. Ce cadre de travail établit un nouveau standard pour l'ajustement fin des LLMs ouverts, les rapprochant ainsi davantage d’un raisonnement humain grâce à des comportements inspirés par la recherche.