Amélioration du routage dans un mélange clairsemé d'experts avec un graphe de jetons

Le mélange clairsemé d'experts (Sparse Mixture of Experts, SMoE) s'est imposé comme un élément clé pour atteindre une évolutivité sans précédent dans l'apprentissage profond. En activant seulement un petit sous-ensemble de paramètres par échantillon, le SMoE réalise une augmentation exponentielle du nombre de paramètres tout en maintenant une charge computationnelle constante. Cependant, les modèles SMoE sont sensibles aux fluctuations de routage — changements dans l’acheminement d’une entrée donnée vers son expert cible — lors des phases tardives de l’entraînement du modèle, ce qui conduit à une non-robustesse du modèle. Dans ce travail, nous dévoilons la limitation du SMoE sous l’angle du modèle graphique probabiliste (PGM). À travers ce cadre PGM, nous mettons en lumière l’indépendance dans la sélection des experts pour les jetons, ce qui expose le modèle aux fluctuations de routage et à la non-robustesse. Pour atténuer cette indépendance, nous proposons le nouveau (Similarity-Aware) SMoE, qui prend en compte les interactions entre les jetons durant la sélection des experts. Nous dérivons ensuite un nouveau PGM sous-jacent à un bloc (S)MoE-Attention qui va au-delà d’une simple couche (S)MoE. En tirant parti des similarités entre jetons capturées par la matrice d'attention, nous proposons le (Similarity-Attention-Aware) SMoE innovant, qui utilise la matrice d’attention pour guider l’acheminement des jetons vers les experts appropriés dans le (S)MoE. Nous démontrons théoriquement que le routage sensible à la similarité ou à l'attention aide à réduire l'entropie de la sélection des experts, entraînant des mécanismes de routage des jetons plus stables. Nous validons empiriquement nos modèles sur diverses tâches et domaines, montrant des améliorations significatives dans la réduction des fluctuations de routage, l’augmentation de la précision et la robustesse du modèle par rapport au MoE-Transformer de référence avec un routage des jetons via un gating softmax.