Résumé Technique : Asymétrie Géométrique dans la Spécialisation des MoE

Énoncé du Problème

Bien que les architectures à Mélange d'Experts (MoE) soient devenues un mécanisme standard pour mettre à l'échelle les modèles Transformer en permettant des nombres de paramètres élevés avec des coûts de calcul fixes, la structure géométrique interne de la spécialisation des experts reste mal comprise. Les analyses existantes se concentrent principalement sur les statistiques de routage, l'équilibrage de charge ou les performances au niveau du système. Un vide critique existe dans la compréhension de la manière dont les experts diffèrent dans l'espace fonctionnel (les transformations qu'ils effectuent) par rapport à l'espace de représentation (la manière dont les jetons acheminés sont organisés dans les états cachés).

L'intuition actuelle suggère que les couches MoE partitionnent la distribution d'entrée en clusters disjoints où chaque expert agit comme un approximeur local. Cependant, les preuves empiriques suggèrent que les experts opèrent souvent sur des régimes de jetons se chevauchant, soulevant la question : Quelle structure géométrique le routage induit-il à l'intérieur d'un Transformer MoE préentraîné ? Plus précisément, les experts implémentent-ils des fonctions distinctes, et comment leurs représentations sont-elles liées les unes aux autres ?

Méthodologie

Les auteurs introduisent un cadre unifié Jacobian-PCA-Grassmann pour analyser les couches MoE à la fois dans les espaces fonctionnels et de représentation. Ce cadre est appliqué à des modèles préentraînés (Mistral, Qwen) et à des modèles Transformer contrôlés à 3 couches entraînés à partir de zéro.

1. Géométrie Fonctionnelle : Jacobiens Locaux aux Experts

Pour caractériser la fonction d'experts individuels indépendamment de la dynamique de routage, les auteurs analysent le jacobien local à l'expert ($J_e(x) = \partial f_e(x) / \partial x$).

• Métrique : Ils calculent le jacobien moyen ($\bar{J}_e$) pour chaque expert sur un ensemble de données et le vectorisent.
• Alignement : L'alignement fonctionnel par paires est mesuré via le produit scalaire de Frobenius normalisé (similarité cosinus) entre ces jacobiens vectorisés.
• Interprétation : Une faible similarité cosinus indique une décorrélation fonctionnelle, ce qui signifie que les experts implémentent des transformations localement distinctes plutôt que des variantes redondantes d'une même fonction.

2. Géométrie de Représentation : PCA Acheminée et Distances Grassmanniennes

Pour caractériser la géométrie des représentations acheminées, les auteurs collectent les états cachés pondérés par les poids de routage ($h_e(x) = g_e(x)h(x)$).

• PCA : Une Analyse en Composantes Principales (PCA) standard est effectuée sur les états cachés acheminés pour identifier les directions principales et les spectres de variance.
• Distance de Sous-espace : La séparation géométrique entre les experts est quantifiée à l'aide de la Distance Géodésique Grassmannienne entre les sous-espaces PCA des $n$ premiers experts différents.
• Interprétation : Cette métrique capture la séparation multidirectionnelle entre des sous-espaces entiers. Une distance proche de zéro implique des rôles fonctionnels redondants, tandis qu'une distance approchant le maximum théorique ($\frac{\pi}{2}\sqrt{n}$) implique une orthogonalité.

3. Expériences de Routage Contrôlé

Pour isoler l'effet des mécanismes de routage des dynamiques d'entraînement, les auteurs comparent deux régimes sur un modèle contrôlé à 3 couches avec des paramètres d'experts fixes :

• Routage Top-k : Sélection sparse et discrète.
• Routage Fully-Soft : Mélange dense et continu.

Résultats Clés

1. Asymétrie Structurelle : Décorrélation Fonctionnelle vs Chevauchement de Représentation

À travers les Transformers MoE préentraînés (Mistral-8x7B et Qwen1.5-MoE-A2.7B), les auteurs observent une asymétrie structurelle cohérente :

• Forte Décorrélation Fonctionnelle : L'alignement des jacobiens entre experts est constamment proche de zéro (moyenne $\approx$ 0,000–0,068). Cela indique que les experts ne sont pas redondants ; ils implémentent des transformations locales distinctes.
• Chevauchement Partiel de Représentation : Malgré la décorrélation fonctionnelle, les représentations acheminées occupent des sous-espaces distincts mais partiellement chevauchants. Les distances grassmanniennes sont substantielles (indiquant une séparation) mais restent significativement inférieures au maximum théorique pour une orthogonalité complète.
• Conclusion : La décorrélation fonctionnelle et le chevauchement de représentation coexistent. Les experts agissent comme des opérateurs localement décorrélés sur des régimes se chevauchant d'une variété partagée, plutôt que de partitionner l'espace d'entrée en clusters disjoints.

2. La Sparsité du Routage comme Molette de Contrôle Géométrique

Les expériences contrôlées révèlent que la netteté du routage façonne directement cette géométrie :

• Routage Top-k : Induit une séparation fonctionnelle plus nette (faible alignement des jacobiens) et de plus grandes distances grassmanniennes (plus grande divergence de sous-espaces).
• Routage Fully-Soft : Produit des structures d'experts entrelacées avec un alignement élevé des jacobiens (indiquant un effondrement vers des mappings redondants) et des distances grassmanniennes effondrées (fort chevauchement de sous-espaces).
• Implication : La sparsité du routage est un facteur clé dans l'imposition de la spécialisation.

3. Comparaison avec les Couches Denses

Les experts MoE présentent un déclin plus raide des valeurs propres dans leurs spectres de variance par rapport aux couches feedforward denses. Les premières composantes principales capturent un pourcentage significativement plus élevé de variance chez les experts MoE, suggérant que le routage concentre les jetons dans des sous-espaces de plus faible dimension et plus focalisés.

4. Cohérence Inter-Modèles

Les motifs géométriques observés (faible alignement des jacobiens + séparation grassmannienne élevée mais sous-maximale) sont cohérents à travers différents modèles (Mistral, Qwen), couches (tôt, milieu, tard) et ensembles de données, suggérant qu'il s'agit d'une propriété robuste des architectures MoE plutôt que d'un artefact de runs d'entraînement spécifiques.

Signification et Revendications

L'article revendique fournir un cadre diagnostique général pour étudier le calcul conditionnel dans les architectures Transformer modernes. Ses contributions principales sont :

1. Outil Diagnostique Unifié : Il offre une méthode systématique pour analyser la spécialisation des experts au-delà des simples statistiques de routage ou des motifs d'activation, caractérisant conjointement l'espace fonctionnel et l'espace de représentation.
2. Interprétation Géométrique des MoE : Les résultats soutiennent une vision des couches MoE implémentant un partitionnement doux d'une variété de représentation partagée. Les experts sont fonctionnellement distincts mais opèrent sur des sous-variétés se chevauchant.
3. Insight Mécanistique : L'étude identifie la netteté du routage (sparsité) comme un mécanisme de contrôle direct du degré de spécialisation, reliant les choix de conception architecturale à des structures géométriques mesurables.
4. Implications pour la Conception de Modèles : Les auteurs suggèrent que les stratégies de fusion ou d'élagage d'experts basées uniquement sur la similarité de représentation (par exemple, le chevauchement d'activation) peuvent être erronées. Parce que les experts peuvent avoir des représentations se chevauchant tout en effectuant des fonctions distinctes, des métriques fonctionnelles (comme l'alignement des jacobiens) peuvent être nécessaires pour identifier correctement la véritable redondance.

L'article maintient une portée modeste, reconnaissant que son analyse est locale (jacobiens d'ordre un) et linéaire (PCA), et ne capture pas le comportement non linéaire global ni l'évolution de la géométrie pendant l'entraînement. Cependant, il établit une perspective géométrique fondamentale sur la manière dont les modèles MoE organisent le calcul.