DLLG: Dynamic Logit-Level Gating of LLM Experts

Résumé Technique : DLLG : Gating Dynamique au Niveau des Logits pour les Experts de LLM

Énoncé du Problème
La prolifération de modèles de langage spécialisés (LLM) optimisés pour des domaines étroits (ex. : raisonnement mathématique, génération de code) présente un défi : comment intégrer des experts « boîte noire » entraînés indépendamment dans un système unifié qui tire parti de leurs forces complémentaires sans coûts de réentraînement prohibitifs. Les approches existantes souffrent d'une tension persistante entre adaptativité, robustesse et praticité :

• Les méthodes basées sur le routage reposent sur une sélection d'experts précoce et rigide. Cela conduit à un engagement prématuré ; si l'expert choisi échoue sur une sous-tâche spécifique au sein d'une seule réponse, l'erreur est irrécupérable. De plus, de nombreux routeurs nécessitent un calibrage intensif spécifique aux benchmarks.
• L'Ensemblage Heuristique (ex. : pondération par la perplexité, vote de jetons) utilise des signaux proxys au moment de l'inférence. Ces métriques sont souvent empiriques et mal alignées avec la véritable correction de la tâche, menant à des sorties sous-optimales.
• La Fusion de Paramètres (ex. : model souping, task arithmetic) fusionne statiquement les poids. Cela sacrifie la modularité et peut induire des interférences destructrices lors de la combinaison d'experts ayant des directions de gradient conflictuelles, entraînant une dégradation des performances.
• Les architectures Mixture-of-Experts (MoE) nécessitent généralement un entraînement conjoint de bout en bout des portes (gates) et des experts, ce qui limite leur capacité à utiliser des spécialistes pré-entraînés prêts à l'emploi.

Méthodologie : DLLG
Les auteurs proposent DLLG (Dynamic Logit-Level Gating), un cadre qui réalise une fusion de logits dynamique au niveau du jeton (token) en utilisant un modèle de gating léger entraîné sur une supervision parcellaire au niveau de la réponse.

1. Architecture :

   • Experts Gelés : $K$ LLM spécialisés restent gelés. À chaque étape de décodage $t$, étant donné un prompt $x$ et un préfixe $y_{<t}$, chaque expert $E_i$ produit un état caché $h^{(i)}_t$ et les logits du jeton suivant $\ell^{(i)}_t$.
   • Modèle de Gating : Un modèle léger ($g_\phi$) traite les états cachés concaténés de tous les experts. Il projette ces représentations dans un espace d'embedding partagé (via une projection de faible rang) et les traite à travers une tige autoregressif (ex. : Qwen2.5-0.5B-Instruct avec des adaptateurs LoRA).
   • Prédiction de Poids : Une tête de prédiction produit des poids de fusion par jeton $w_{t,i}$ via une transformation sigmoïde. Ces poids sont appliqués aux logits des experts : $\ell_t = \sum_{i=1}^K w_{t,i} \ell^{(i)}_t$. Les poids ne sont pas contraints de sommer à un, permettant des contributions souples et flexibles.

2. Entraînement (Supervision au Niveau de la Réponse) :

   • Le cadre traite l'absence d'annotations au niveau du jeton en utilisant des signaux de correction au niveau de la réponse ($s_i \in \{0, 1\}$) dérivés de vérificateurs automatiques (ex. : tests unitaires pour le code, vérification de réponses pour les mathématiques).
   • Teacher Forcing : Pendant l'entraînement, les experts génèrent des états cachés conditionnés par le préfixe de vérité terrain $y_{<t}$. Le modèle de gating prédit les poids $w_{t,i}$ pour chaque position de jeton.
   • Objectif : Le modèle est entraîné pour minimiser l'erreur quadratique moyenne (MSE) entre les poids prédits et le signal de correction de la réponse diffusé : $\mathcal{L} = \frac{1}{T} \sum_{t=1}^T \sum_{i=1}^K (w_{t,i} - s_i)^2$. Cela enseigne au gate à aligner la fusion par jeton avec le succès global de la trajectoire.

3. Inférence :

   • DLLG fonctionne de manière autoregressive. À chaque étape, le modèle de gating prédit les poids de fusion basés sur les états cachés actuels des experts, les logits sont combinés, et le jeton suivant est sélectionné. Cela permet une utilisation récupérable des experts ; si un expert commet une erreur tôt, le modèle de gating peut dynamiquement déplacer le poids vers un expert plus compétent lors des étapes suivantes.

Principales Contributions

1. Proposition de Cadre : Introduction de DLLG, un cadre de fusion de logits dynamique qui fait le pont entre la génération au niveau du jeton et la supervision parcellaire au niveau de la réponse, permettant une utilisation fine des experts sans annotation au niveau du jeton ni apprentissage par renforcement en ligne.
2. Analyse de Mécanisme : Démonstration que le mécanisme de gating appris adapte dynamiquement les contributions des experts pendant le déploiement, ajustant les poids en fonction de l'évolution du contexte (ex. : passage du raisonnement mathématique à la génération de code au sein d'une même réponse).
3. Validation Empirique : Évaluation sur divers benchmarks de raisonnement et de code (GSM8K, MATH, HumanEval, etc.), montant des gains constants par rapport aux baselines de routage, d'ensemblage heuristique et de fusion de paramètres.

Résultats

• Performance : Sur les échelles de modèles 0.5B et 1.5B, DLLG atteint systématiquement la performance moyenne la plus élevée par rapport à des baselines solides, incluant RouterDC, EmbedLLM, diverses méthodes d'ensemblage de logits (GaC, UniTe) et des techniques de fusion de paramètres (Linear, SLERP, Task Arithmetic).
• Robustesse : Dans les configurations à faible capacité (0.5B) où les experts individuels sont faibles, DLLG surpasse les méthodes de routage et les méthodes heuristiques fragiles, démontrant sa capacité à atténuer les erreurs précoces grâce à la fusion souple.
• Généralisation : La méthode montre de solides performances sur des benchmarks hors domaine (ex. : HumanEval, MBPP, BBH), indiquant que les règles de fusion apprises se généralisent au-delà de la distribution d'entraînement spécifique.
• Ablation : Les expériences confirment qu'une architecture de gating suffisamment expressive (incorporant un LLM auxiliaire et une projection de faible rang) est nécessaire pour une fusion efficace au niveau du jeton, surpassant les simples têtes MLP.

Signification et Revendications
L'article affirme que DLLG offre un paradigme robuste et évolutif pour l'intégration d'experts spécialisés. En apprenant la fusion au niveau des logits à partir d'une supervision faible, il atteint une adaptativité de type MoE sans le coût d'un réentraînement des experts ou le risque d'interférence de paramètres. Les auteurs soulignent que DLLG préserve la modularité des experts et permet une utilisation fine et récupérable des experts, surmontant les limites de l'engagement prématuré du routage et la nature statique de la fusion de paramètres. Ce travail suggère que l'apprentissage d'une utilisation fine des experts à partir d'une supervision faible est une voie pratique pour exploiter les modèles de langage de plus en plus spécialisés.