Résumé technique : Apprendre à sélectionner, et non à réapprendre : MoR-LoRA à routage dur pour les mélanges de LoRA de raisonnement

Énoncé du problème

L'article traite du défi consistant à composer des adaptateurs Low-Rank Adaptation (LoRA) entraînés indépendamment en un seul grand modèle de langage (LLM) pour une adaptation multi-domaine, particulièrement lorsque les données d'entraînement originales ne peuvent pas être partagées pour des raisons de confidentialité ou de réglementation.

Un problème critique survient lors de l'application du routage "soft" (souple) classique de Mixture-of-Experts (MoE) à des adaptateurs LoRA gelés et entraînés indépendamment. Les modules LoRA sont entraînés sous l'hypothèse d'une mise à jour additive à échelle unitaire ($W_{updated} = W_{pretrained} + AB$). Cependant, le routage MoE "soft" combine les experts par des sommes pondérées ($W_{updated} = W_{pretrained} + \sum \omega_i A_i B_i$), où les poids de routage $\omega_i < 1$. Cela crée un décalage d'échelle : la mise à jour effective appliquée lors de l'intégration diffère de la mise à jour à échelle unitaire sous laquelle l'expert a été initialement entraîné.

Les solutions existantes, telles que LoRA-Mixer, tentent de résoudre cela en réentraînant les adaptateurs avec des pertes de préservation auxiliaires. Cette approche augmente le nombre de paramètres entraînables, nécessite plus de données et risque d'altérer les connaissances encodées dans les experts originaux. De plus, ce défi est particulièrement prononcé pour les modèles ajustés par instruction (instruction-tuned) dotés de Chaînes de Pensée (CoT), où le fine-tuning conventionnel peut dégrader les comportements de raisonnement internes, et où les méthodes de composition existantes échouent souvent à préserver ces traces de raisonnement spécialisées.

Méthodologie

Les auteurs proposent le MoR-LoRA à routage dur, un cadre en deux étapes conçu pour composer des experts LoRA de raisonnement gelés par une sélection dure à échelle unitaire, évitant ainsi la nécessité de réentraîner les experts.

Étape I : Entraînement des experts de raisonnement

Des adaptateurs LoRA spécifiques à un domaine sont entraînés indépendamment en utilisant l'Apprentissage par Renforcement avec Feedback Vérifiable (RLVF). Contrairement au fine-tuning supervisé (SFT) qui encourage l'imitation de sorties de référence, le RLVF optimise la correction en utilisant des récompenses automatiquement vérifiables (par exemple, vérifier la réponse finale). Ce paradigme encourage l'émergence de CoT structurées sans écraser les schémas de raisonnement alignés du modèle de base. Les experts résultants sont ensuite gelés.

Étape II : Intégration par routage dur

Les experts gelés sont intégrés dans un modèle partagé via un Mélange à Routage Dur (Hard-Routed Mixture).

• Routage dur Top-1 : Pour chaque jeton (token), le routeur sélectionne exactement un expert. Cela garantit que l'expert sélectionné est appliqué avec une échelle unitaire, préservant l'hypothèse originale de LoRA ($W_{updated} = W_{pretrained} + 1 \cdot A_j B_j$).
• Estimateur de Passage Direct (Straight-Through Estimator - STE) : Puisque la sélection discrète (argmax) n'est pas dérivable, les auteurs utilisent le STE. Lors de la passe avant (forward pass), la décision dure est utilisée ; lors de la passe arrière (backward pass), les gradients sont propagés à travers les probabilités "soft". Cela permet un entraînement de bout en bout du routeur tout en maintenant le comportement discret et à échelle unitaire pendant l'inférence.
• Efficacité des paramètres : Seuls un routeur linéaire partagé léger et un petit module d'attention LoRA sont entraînés. Les experts préentraînés restent gelés.
• Distillation : Le routeur est entraîné via un fine-tuning supervisé (SFT) sur un petit ensemble de données de traces de raisonnement distillées à partir des experts gelés. L'objectif est d'apprendre quel expert sélectionner, et non de réapprendre la capacité de raisonnement elle-même.

Contributions clés

1. Identification du décalage d'échelle : L'article identifie que l'application du routage MoE "soft" à des adaptateurs LoRA gelés et entraînés indépendamment introduit un décalage d'échelle qui altère la fonction effective des adaptateurs.
2. Routage dur à échelle unitaire : Les auteurs proposent d'utiliser un routage dur Top-1 avec un estimateur de passage direct (STE) pour préserver l'application à échelle unitaire des experts, permettant une composition modulaire sans réentraînement.
3. Cadre en deux étapes : Un pipeline novateur qui sépare l'acquisition du raisonnement (via RLVF) de la sélection de l'expert (via le routage dur), en gardant les experts gelés lors de l'intégration.
4. Validation empirique : L'étude démontre que cette approche préserve le comportement de l'expert tout en nécessitant nettement moins de paramètres entraînables que les bases de comparaison de routage "soft".

Résultats expérimentaux

Les expériences ont été menées sur cinq benchmarks (GSM8K, ARC-Challenge, MedQA, BoolQ, CoLA) et sur plusieurs échelles de modèles (1B, 3B, 8B) en utilisant les familles LLaMA-3, Gemma-3 et Qwen-2.5.

• Supériorité du RLVF : Les experts entraînés par RLVF ont nettement surpassé les experts entraînés par SFT sur les tâches de raisonnement intensif pour les modèles capables (3B et 8B), alors que le SFT a parfois dégradé les performances de raisonnement dans les grands modèles ajustés par instruction.
• Performance vs Paramètres : Le MoR-LoRA à routage dur a obtenu des performances comparables ou supérieures aux bases de comparaison de routage "soft" (LoRAMixer, MoLE) tout en entraînant substantiellement moins de paramètres. Par exemple, sur le modèle 3B avec des experts RLVF, la méthode proposée a entraîné environ 73M de paramètres pour atteindre une précision moyenne de 72,07, tandis que LoRAMixer nécessitait environ 606M de paramètres pour atteindre 65,67.
• Analyse du comportement de routage : L'analyse du routage "soft" Top-2 a révélé qu'il converge souvent vers une sélection proche du Top-1 (poids dominant moyen ~0,71), suggérant que le routage dur capture le comportement voulu de manière plus directe et efficace.
• Efficacité des données : L'étape d'intégration (Étape II) n'a nécessité qu'environ 1 000 échantillons par jeu de données pour atteindre une performance proche de l'entraînement sur l'ensemble du jeu de données, indiquant que le routeur apprend les schémas de sélection plutôt que la connaissance du domaine.
• Capacité de domaine mixte : Dans les évaluations de domaines mixtes (par exemple, un prompt contenant à la fois des mathématiques et de la compréhension de lecture), le routage dur au niveau du jeton (token-level) a surpassé le routage au niveau du prompt, démontrant la flexibilité de la sélection de différents experts au sein d'une même entrée.

Signification et affirmations

L'article affirme que la l'adaptation modulaire peut être réalisée en apprenant quand utiliser chaque expert plutôt qu'en modifiant ce que chaque expert sait.

Les auteurs soutiennent que pour les experts LoRA préentraînés et gelés, le routage dur à échelle unitaire fournit une abstraction simple et efficace. En traitant l'intégration principalement comme un problème de routage, la méthode évite le coût computationnel et la dégradation potentielle des connaissances associés au réentraînement des experts. Les résultats suggèrent que le principal avantage de la composition provient de la sélection de l'expert approprié pour le contexte, plutôt que d'une modification extensive des paramètres des experts.

Le travail souligne que si le routage au niveau du prompt est suffisant pour les entrées propres à un seul domaine, le routage dur au niveau du jeton offre la flexibilité nécessaire pour les entrées complexes à domaines mixtes. Les auteurs notent modestement que la méthode n'est pas entièrement exempte d'entraînement, car l'Étape II nécessite quelques données étiquetées pour apprendre la politique de routage, et qu'elle ne supporte actuellement pas l'intégration "zero-shot" réelle de nouveaux domaines sans données de routage supplémentaires.