Advancing General-Purpose Reasoning Models with Modular Gradient Surgery

Résumé Technique : Modular Gradient Surgery (MGS)

1. Problématique (Le Problème)

L'objectif actuel du développement des grands modèles de langage (LLM) est de créer des modèles de raisonnement généralistes (LRM) capables de maîtriser des domaines très divers (mathématiques, chat ouvert, suivi d'instructions) via l'apprentissage par renforcement (RL).

Cependant, les chercheurs font face à un défi majeur : l'hétérogénéité des domaines. Lorsqu'on tente d'entraîner un modèle sur plusieurs domaines simultanément, on observe une interférence croisée qui limite les performances. L'étude identifie deux stratégies classiques inefficaces :

• SEQUENTIAL RL (RL Séquentiel) : Entraîner le modèle sur un domaine après l'autre. Cela provoque deux phénomènes : le "Forgetting" (oubli des compétences du premier domaine au profit du second) et la "Rigidity" (incapacité du modèle à apprendre de nouvelles compétences car les paramètres ont été trop spécialisés).
• MIXED RL (RL Mixte/Multi-tâches) : Mélanger les données de tous les domaines dans chaque lot (batch). Cela provoque des conflits de gradients : la direction de mise à jour optimale pour les mathématiques peut être diamétralement opposée à celle du chat, annulant ainsi les progrès des deux côtés.

2. Méthodologie (La Solution : MGS)

Pour résoudre ces conflits, les auteurs introduisent la Modular Gradient Surgery (MGS).

L'idée repose sur l'observation que les transformeurs sont des architectures hautement modulaires : les couches d'attention gèrent principalement le routage de l'information et le raisonnement, tandis que les couches MLP (Multi-Layer Perceptron) sont davantage liées à la mémoire paramétrique et au stockage des connaissances. Les conflits de gradients ne sont donc pas uniformes sur tout le modèle, mais localisés dans certains modules.

Fonctionnement technique :
Au lieu d'appliquer une correction de gradient globale (comme la méthode PCGrad qui traite tout le modèle comme un seul vecteur), MGS décompose le gradient par module (Attention, MLP, LayerNorm) :

1. Calcul des gradients pour chaque tâche au sein d'un lot.
2. Détection des conflits : Pour chaque module $\phi$, on calcule le produit scalaire des gradients des tâches $i$ et $j$. Si le produit est négatif ($\nabla_{\phi} L_i \cdot \nabla_{\phi} L_j < 0$), il y a conflit.
3. Projection orthogonale : Si un conflit est détecté, le gradient de la tâche $i$ est projeté sur l'espace orthogonal au gradient de la tâche $j$. Cela permet de conserver uniquement la composante du gradient qui ne nuit pas à l'autre tâche.
4. Agrégation modulaire : Les gradients "corrigés" sont ensuite réassemblés pour mettre à jour le modèle.

3. Contributions Clés

• Analyse systématique : Une étude approfondie des échecs du RL multi-domaine (interférences de mode et conflits de gradients).
• Innovation algorithmique : L'introduction de la "chirurgie de gradient modulaire", qui adapte les techniques de multi-tâches classiques à la structure spécifique des transformeurs.
• Preuve de passage à l'échelle : Démontrer que la méthode fonctionne non seulement sur deux tâches, mais aussi sur trois (Math, Chat, Instruction Following) et sur des entraînements prolongés.

4. Résultats Expérimentaux

Les tests ont été menés sur les familles de modèles Llama-3.1-8B et Qwen-2.5-7B.

• Amélioration globale : MGS surpasse nettement le mélange naïf et le RL séquentiel. Sur Llama, il apporte une amélioration moyenne de 4,5 points ; sur Qwen, de 4,3 points.
• Équilibre des compétences : Contrairement au RL séquentiel qui sacrifie souvent le chat pour les maths (ou vice versa), MGS permet d'atteindre des scores élevés et équilibrés sur tous les benchmarks (MATH, WildBench, IFEval).
• Robustesse : La méthode reste efficace même lorsque le nombre d'étapes d'entraînement est doublé, là où les méthodes classiques commencent à stagner ou à diverger.
• Efficacité computationnelle : Le surcoût en mémoire et en temps de calcul est jugé négligeable grâce à l'utilisation de frameworks de calcul parallèle comme FSDP.

5. Signification et Impact

Ce travail est crucial car il fournit une "recette" stable et efficace pour transformer des modèles de base en modèles de raisonnement généralistes. En résolvant le problème de l'imbalance et des conflits de gradients au niveau des modules, MGS ouvre la voie à des modèles capables de raisonner de manière complexe (via des chaînes de pensée/CoT) tout en conservant une grande polyvalence conversationnelle et une capacité à suivre des instructions précises.