Multi-Task GRPO: Reliable LLM Reasoning Across Tasks

Imaginez que vous entraînez un robot assistant super intelligent pour résoudre toutes sortes d'énigmes, des problèmes mathématiques aux devinettes logiques. Vous voulez que ce robot soit un véritable généraliste, capable de gérer n'importe quoi sans s'embrouiller ou échouer face aux difficultés. Dans le monde de l'intelligence artificielle, cela s'appelle le « raisonnement », et les robots sont appelés des « Grands Modèles de Langage » (LLM). Pour rendre ces modèles meilleurs, les scientifiques utilisent une technique appelée « Apprentissage par Renforcement », qui consiste à donner au robot un « tape sur l'épaule » (une récompense) quand il trouve la bonne réponse et un doux « réessaie » quand il se trompe. Une méthode populaire pour faire cela est le GRPO, une façon ingénieuse de déterminer exactement ce que le robot doit apprendre de chaque essai. Cependant, il y a un piège : si vous demandez au robot d'apprendre dix types d'énigmes différents en même temps, il se bloque souvent. Il peut devenir très bon aux énigmes faciles, mais ignorer complètement les plus difficiles, ou bien rester bloqué sur un type spécifique de devinette tout en oubliant comment faire les autres. La grande question est : comment entraîner un seul cerveau pour qu'il soit aussi bon en tout, en veillant à ce qu'aucune compétence ne soit laissée de côté ?

C'est exactement le problème abordé par une nouvelle méthode appelée MT-GRPO (Multi-Task GRPO). Considérez le processus d'entraînement standard comme un professeur essayant d'enseigner à une classe d'élèves ayant des vitesses d'apprentissage très différentes. Si le professeur se concentre uniquement sur l'élève « moyen », les élèves brillants pourraient s'ennuyer tandis que les élèves en difficulté pourraient décrocher davantage. L'approche standard permet souvent aux tâches faciles de « détourner » l'entraînement, faisant paraître le modèle performant en moyenne alors qu'il échoue secrètement aux défis les plus complexes. Les auteurs de cet article ont réalisé qu'en se contentant de faire la moyenne des scores, on ne crée pas un robot fiable ; il faut s'assurer que le pire performeur de la classe progresse également.

Pour corriger cela, les chercheurs ont inventé une stratégie en deux parties. D'abord, ils ont créé un « système de pondération dynamique ». Imaginez un entraîneur qui surveille constamment le tableau des scores. Si le robot réussit très bien le « Countdown » (un jeu de mathématiques) mais est très mauvais aux « énigmes du Zèbre » (des devinettes logiques), l'entraîneur déplace immédiatement l'attention. Il commence à donner plus de problèmes d'entraînement de la catégorie Zèbre et moins de la catégorie mathématiques. Mais voici la partie délicate : parfois, le robot tente une énigme du Zèbre et obtient un score de « zéro » parce qu'il n'a même pas essayé de la résoudre correctement, ce qui signifie qu'il n'apprend rien de cette tentative. Si l'entraîneur se contente de compter le nombre de problèmes, il pourrait penser qu'il enseigne les énigmes du Zèbre, alors qu'en réalité, le robot n'apprend que grâce aux exercices de mathématiques car les tentatives du Zèbre étaient des « gradients nuls » (un apprentissage vide).

Pour résoudre cela, la seconde partie de leur invention est un « Échantillonneur Préservant le Ratio » (Ratio-Preserving Sampler). C'est comme un bibliothécaire très strict qui veille à ce que, même si le robot jette les tentatives d'entraînement inutiles, la pile finale de problèmes utiles contienne exactement le bon mélange de mathématiques et de logique. Si l'entraîneur veut que 30 % de l'entraînement soit composé d'énigmes de logique, le bibliothécaire s'assure qu'après avoir filtré les déchets, 30 % des problèmes restants de qualité sont toujours des énigmes de logique. Cela empêche les tâches faciles de prendre accidentellement le dessus sur la session d'entraînement simplement parce que les tâches difficiles génèrent plus de tentatives « ratées ».

Les résultats de cette nouvelle méthode sont très prometteurs. Lorsque l'équipe a testé sa méthode sur un mélange de trois tâches de raisonnement différentes, leur nouvelle méthode MT-GRPO a amélioré la précision de la tâche la moins performante de 16 à 28 % par rapport à la méthode standard, et de 6 % par rapport à une autre méthode avancée appelée DAPO. Plus impressionnant encore, elle a atteint un seuil de précision de 50 % sur les tâches les plus difficiles en utilisant 50 % de cycles d'entraînement en moins que la concurrence. Dans un test plus large portant sur neuf tâches différentes, la méthode a continué de surpasser les autres, montrant qu'elle peut monter en puissance pour gérer des défis plus complexes et diversifiés. Les auteurs ont découvert qu'en ajustant un « bouton » (un paramètre appelé $\lambda$), ils pouvaient contrôler à quel point le modèle priorise la correction de ses points faibles par rapport au simple fait de s'améliorer en moyenne.

En résumé, cet article suggère qu'en étant plus intelligents sur le choix des problèmes que le robot pratique et en garantissant que ces problèmes parviennent réellement à la session d'entraînement, nous pouvons construire une IA qui n'est pas seulement bonne dans les choses faciles, mais réellement compétente de manière globale. C'est une étape vers la création d'une IA qui ne se spécialise pas seulement dans un domaine étroit, mais qui peut véritablement raisonner à travers une grande variété de problèmes du monde réel sans laisser aucune de ses compétences derrière elle.