Mixture-of-Experts as Soft Clustering: A Dual Jacobian-PCA Spectral Geometry Perspective

1. Problématique et Contexte

Les architectures Mixture-of-Experts (MoE) sont largement adoptées dans les grands modèles de langage (LLM) pour leur efficacité computationnelle, permettant d'activer un sous-ensemble de paramètres par entrée. Cependant, l'impact géométrique du routage des experts sur la structure des fonctions apprises et des représentations internes reste mal compris.

La question centrale est de savoir comment le mécanisme de routage (qu'il soit Top-k ou fully-soft) modifie :

• La sensibilité locale de la fonction (courbure).
• La distribution de la variance dans l'espace des représentations.
• L'alignement fonctionnel entre les différents experts.

Contrairement aux travaux précédents axés sur les lois d'échelle ou l'efficacité, cet article vise à isoler les effets géométriques du routage indépendamment de la performance sur une tâche spécifique.

2. Méthodologie

L'auteur propose une approche analytique basée sur deux sondes géométriques complémentaires, appliquées dans un environnement contrôlé (MLP-MoE avec entrées synthétiques gaussiennes) pour permettre un calcul exact des Jacobiens et une estimation stable de la covariance.

A. Sonde Spectrale du Jacobien (Géométrie de la fonction)

• Calcul : Analyse des valeurs singulières du Jacobien local de chaque expert ($J_e(x) = \partial f_e(x) / \partial x$).
• Métriques :
  • La plus grande valeur singulière ($\sigma_1$) comme indicateur de la sensibilité locale maximale (norme spectrale).
  • La décroissance du spectre pour évaluer l'anisotropie.
  • L'alignement cosinus entre les Jacobiens moyens des experts pour mesurer le chevauchement fonctionnel.

B. Analyse en Composantes Principales Pondérée (Géométrie de la représentation)

• Calcul : PCA appliquée aux états cachés ($h(x)$) entrant dans la couche MoE, où les échantillons sont pondérés par les poids de routage ($g_e(x)$).
• Objectif : Caractériser comment la variance est distribuée dans l'espace des représentations spécifiques à chaque expert.
• Métriques : Ratio de variance expliquée, courbes de variance cumulative et rang effectif (nombre de composantes nécessaires pour atteindre un seuil de variance, ex: 90%).

C. Comparaisons Architecturales

L'étude compare trois configurations sous capacité égale :

1. Modèles denses (baseline).
2. MoE avec routage Top-k (routage dur/éparse).
3. MoE avec routage fully-soft (routage probabiliste complet).

3. Résultats Clés

A. Réduction de la Sensibilité Locale (Jacobian)

• Sensibilité : Les experts MoE présentent des valeurs singulières principales ($\sigma_1$) significativement plus faibles (1,5 à 1,7 fois moins) que les modèles denses.
• Stabilité : La norme du Jacobien se stabilise rapidement lors de l'entraînement pour les MoE, tandis qu'elle continue de croître pour les modèles denses.
• Orthogonalité : L'alignement cosinus entre les Jacobiens moyens des experts est proche de zéro. Cela indique que les experts apprennent des transformations fonctionnelles distinctes avec un faible chevauchement, plutôt que des variantes d'une carte partagée.

B. Redistribution de la Variance (PCA)

• Rang Effectif : Dans le cadre synthétique (données isotropes), les représentations locales des experts ont un rang effectif plus élevé que le modèle dense. La variance est répartie sur plus de directions principales, contrairement au modèle dense qui concentre la variance sur quelques composantes.
• Effet du Routage :
  • Routage Top-k : Produit des structures de variance plus concentrées (rang effectif plus faible) car les experts opèrent sur des régions d'entrée plus étroites.
  • Routage Fully-soft : Produit des structures de variance plus larges et un rang effectif plus élevé, car les experts sont exposés à des entrées plus hétérogènes.

C. Validation sur Données Réelles (WikiText-2)

Une validation sur un transformateur de 3 couches entraîné sur du texte naturel révèle une inversion géométrique :

• Courbure : La réduction de la courbure (Jacobian plus plat) persiste pour les experts MoE par rapport au dense.
• PCA : Contrairement au cas synthétique, les experts MoE sur données réelles atteignent 90% de la variance cumulative avec moins de composantes (rang effectif plus faible) que le modèle dense.
• Interprétation : Les données de langage résident sur une variété de basse dimension. Le routage MoE agit comme un clustering doux, partitionnant les tokens en sous-variétés plus simples où chaque expert peut opérer localement avec une dimensionnalité réduite, tandis que le modèle dense doit gérer toute la complexité via une seule transformation.

4. Contributions Principales

1. Cadre Géométrique : Interprétation des MoE comme une partition douce de l'espace des fonctions en "cartes" locales d'experts qui se chevauchent.
2. Méthodologie de Sondage : Introduction d'une approche combinée (Spectre Jacobien + PCA pondérée) pour quantifier la sensibilité locale et la structure de la variance.
3. Découverte de l'Inversion Géométrique : Mise en évidence du fait que l'effet du routage sur le rang effectif dépend de la nature des données (isotrope vs variété de langage), tout en maintenant une réduction de la courbure fonctionnelle.
4. Hypothèses Testables : Proposition de critères de conception pour les grands modèles (ex: 70B+), notamment l'existence d'un "coude géométrique" pour le nombre d'experts et la réduction potentielle des hallucinations grâce à des mappings plus plats.

5. Signification et Implications

Ce travail fournit une base mécaniste pour comprendre pourquoi et comment les MoE fonctionnent au-delà de la simple économie de calcul.

• Interprétation du Clustering : Le routage agit comme un mécanisme de clustering doux qui décompose la complexité globale en sous-problèmes géométriquement plus simples (courbure réduite).
• Robustesse et Hallucinations : La réduction de la norme du Jacobien suggère que les MoE pourraient être intrinsèquement plus robustes aux petites perturbations, potentiellement réduisant les erreurs d'accumulation (hallucinations) dans la génération de contexte long.
• Diversité de l'Ensemble : Le faible alignement entre les Jacobiens des experts soutient l'idée que le routage Top-k favorise une diversité fonctionnelle supérieure par rapport au routage fully-soft, ce qui est crucial pour la fiabilité des ensembles.
• Conception Architecturale : Les résultats suggèrent que le choix de la "dureté" du routage (Top-k vs Soft) ne contrôle pas seulement l'activation, mais aussi la dimensionnalité interne et la structure de la variance des représentations, offrant de nouveaux leviers pour optimiser les transformateurs à contexte long.

En résumé, l'article démontre que les MoE ne sont pas seulement des outils d'efficacité, mais qu'ils réorganisent fondamentalement la géométrie de l'apprentissage en aplatissant la courbure locale et en adaptant la dimensionnalité des représentations à la structure sous-jacente des données.