Deep learning : les promesses de la technique « Mixture of Experts »

Switch Transformer : Google cherche à faire plus avec moins

Ainsi, le modèle Switch Transformer a été conçu pour augmenter le nombre de paramètres sans augmenter la quantité de calcul nécessaire. Dans le détail, les chercheurs de Google ont adapté la technique « Mixture of experts » (MoE) à leurs besoins.

Un expert est une couche ou une sous-couche d’un réseau de neurones spécialisé dans un domaine.

Cette technique ensembliste, théorisée dans les années 1990, consiste au départ à faire appel à des « experts » associés à un lot de paramètres utilisés pour diviser l’espace du problème à traiter, en des régions homogènes de taille réduite. En clair, il s’agit de spécialiser certains éléments ou couches d’un réseau de neurones dans le traitement de tâches subalternes. Un expert est donc une couche ou une sous-couche d’un réseau de neurones spécialisé dans un domaine. Cet expert doit répondre à un problème et un seul, par exemple classer un type d’images, des pronoms ou des phrases suivant les paramètres associés.

Les experts sont souvent rassemblés dans des sous-couches d’un réseau de neurones plus vaste. In fine, la combinaison de l’analyse de chacun de ces spécialistes doit permettre d’obtenir une prédiction fine. Outre le fait de diviser le mégamodèle en un puzzle consistant, il est possible de configurer certaines des pièces pour améliorer les résultats du modèle global dans certains domaines.

Ici, Google cherche à appliquer cette technique au traitement du langage naturel, qui se prête parfaitement à cette subdivision des problèmes. En effet, un modèle NLP doit « comprendre » des termes et le contexte dans lequel ils sont utilisés. Par exemple, un expert peut traiter les pronoms, tandis qu’un deuxième comprendra les verbes, un autre les noms et ainsi de suite.

 Mais avec l’algorithme MoE spécifique au NLP que Google avait développé précédemment, il faut théoriquement router un jeton (ou token, issu de la tokenisation, à savoir une portion plus ou moins grande d’un texte) vers au moins deux experts à chaque transformation, c’est-à-dire entre chaque passage de couche.

Switch Transformer, lui, est un modèle épars qui n’active qu’un sous-ensemble des poids d’un réseau. Pour cela, les chercheurs ont remplacé la couche de réseau de neurones à propagation avant (Feedforward neural network ou FFN) par une couche FFN éparse ainsi que l’architecture de routage du MoE. Un token est uniquement routé vers un expert de la couche FFN le plus à même de traduire ou de comprendre un mot, une phrase ou un paragraphe.

Chaque expert est indépendant et peut être placé sur des machines différentes via un mécanisme de sharding. Résultat, Switch Transformer accélère jusqu’à sept fois le temps d’entraînement par rapport au modèle T5 de Google, tout en utilisant le même nombre d’opérations en virgule flottante par seconde (FLOPS). Du même coup, lors du test, seulement 0,10 % des paramètres étaient activés pour chaque token, contre 100 % de ceux d’OpenAI. En revanche, l’hypothèse selon laquelle l’augmentation du nombre de paramètres améliore la qualité du modèle ne se vérifie pas forcément. Si dans la plupart des cas, Switch Transformer obtient les mêmes résultats ou surperforme face à T5 lors de benchmark NLP, dans certains cas, le ratio paramètres-performance n’est pas satisfaisant, selon l’article de recherche.

Cette technique est-elle efficace pour des modèles NLP beaucoup moins volumineux ? Les gains ne seraient pas aussi importants qu’avec un mégamodèle conventionnel, mais le routage des tokens vers deux experts suffiraient déjà à améliorer les performances, selon les chercheurs de Google.

Microsoft combine MoE et optimisation du parallélisme

Google n’est pas le seul à embrasser la technique Mixture of Experts. Microsoft et OpenAI sont en train de l’adapter à GPT-3. Un papier de recherche consacrée à ce sujet a été publié le 14 janvier 2022. Si les scientifiques posent la même question et cherchent à obtenir une réduction drastique de la consommation de VRAM pour des modèles incluant des milliers de milliards de paramètres, ils tiennent à rappeler que GPT-3 et Switch Transformer n’ont pas grand-chose à voir. Selon eux, les applications de la MoE aux tâches d’encodage - décodage, somme toute classiques dans le domaine du NLP, sont beaucoup plus documentées dans le domaine académique que les modèles de génération de langage, plus récents.

Aussi, les chercheurs de Microsoft veulent obtenir un meilleur ratio entre les performances et les paramètres tout en appliquant la technique à l’entraînement et à l’inférence.

Pour ce faire, ils développent un modèle intitulé Pyramid-Residual MoE (PR-MoE). À la différence d’un MoE standard employé dans le NLP, les deux dernières couches de PR-MoE ont le double d’experts des couches précédentes. À chaque passage de couche, le token est routé vers un module perceptron multicouche fixe et un expert spécifique. Cette approche permet de réduire le nombre de paramètres nécessaires tout en rendant leur utilisation plus efficiente. Il s’agit de répartir équitablement les paramètres par token, mais aussi de « corriger » la sortie de la couche perceptron à l’aide de l’expert, avec le même volume de communication en provenance du routeur.

 Résultat, « PR-MoE peut réduire jusqu’à 3 fois la taille des paramètres des modèles MoE sans que la qualité du modèle soit modifiée et avec un minimum d’exigences de calcul », écrivent les chercheurs de Microsoft.

Dans certains cas, le coût de l’entraînement serait cinq fois moins élevé : PR-MoE a été utilisé pour entraîner un modèle NLG comprenant 6,7 milliards de paramètres à l’aide d’un lot de 128 GPU Nvidia A100, alors que l’approche précédente ne permettait pas de dépasser 1,3 milliard de paramètres avec la même configuration.

« Bien que des études récentes traitent de la réduction des FLOPs, [...] contrairement à l’entraînement, la latence et le coût de l’inférence ne dépendent pas uniquement du calcul. »

Mais les chercheurs de Microsoft obtiennent ces gains significatifs quand ils combinent PR-MoE et DeepSpeed, une bibliothèque open source d’optimisation de réseaux de neurones. DeepSpeed permet de mieux partitionner les éléments d’un modèle sur une architecture HPC, ici les experts, et de limiter la consommation de mémoire vidéo. En clair, il s’agit d’optimiser le parallélisme des données dans une architecture MoE. Microsoft avait déjà mené des recherches en ce sens, mais PR-MoE promet de récolter des résultats similaires avec moins de ressources de calcul.

DeepSpeed est également à l’origine de DeepSpeed-MoE Inference. Ici, il s’agit de partitionner, de grouper et de router efficacement les tokens vers les experts. Pour ce faire, les chercheurs ont optimisé la planification des communications à travers une coordination fine du parallélisme et le code du réseau de neurones.

« Bien que des études récentes traitent de la réduction des FLOPs, il est pertinent de mentionner que, contrairement à l’entraînement, la latence et le coût de l’inférence ne dépendent pas uniquement du calcul », notent les data scientists de Microsoft.

« Une inférence efficace dépend de la taille du modèle, de la bande passante de la mémoire et de la capacité d’un système à lire efficacement les données en mémoire. »

Cela permet de servir le modèle NLG 7,3 fois plus rapidement à un coût inférieur par rapport « à d’autres systèmes MoE ». L’inférence serait 4,5 fois plus rapide et jusqu’à 9 fois moins chère.

Le risque d’un déséquilibre entre les poids d’un réseau de neurones

Les chercheurs de Microsoft souhaitent aller encore plus loin dans ce domaine en réduisant les mégamodèles NLG optimisés avec l’architecture MoE. Comment ? Par l’application d’une technique de distillation de l’apprentissage.

« Nous formons d’abord un modèle de MoE dit “enseignant”. Nous réduisons la profondeur de chaque branche d’experts dans le modèle enseignant pour obtenir un étudiant correspondant », expliquent les data scientists de Microsoft. « De cette façon, le modèle final de l’élève possède la même architecture à faible densité (ou éparse) que le modèle MoE de l’enseignant, sauf que chaque branche d’experts a une profondeur réduite. Pour cette raison, nous appelons le modèle résultant Mixture-of-Students (MoS) ».

En réalité, c’est tout un pan de la recherche en data science qui s’active autour de ce sujet de la MoE. Avant cela, Meta AI (Facebook AI) a publié sa propre librairie open source et poursuit des recherches pour mieux répartir les jetons entre les différents experts.

Car l’emploi de modèles basés sur la technique MoE demeure complexe. Les promesses séduisent les géants technologiques, qui en profitent pour livrer les résultats de leur recherche et donc faire progresser l’adoption de réseaux de neurones plus efficients. Pour autant, la répartition des jetons entre les experts représente un risque pour l’efficacité de la généralisation de l’usage des modèles développés. Or c’est bien l’objectif premier de Google, de Facebook, de Microsoft ou encore AWS que de mettre au point des algorithmes capables de répondre à la majorité des problématiques liées au traitement du langage naturel.