Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer; Azalia Mirhoseini; Krzysztof Maziarz; Andy Davis; Quoc Le; Geoffrey Hinton; Jeff Dean

Articles · arXiv 1701.06538

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean

Publié le 2017-01-23cs.LGcs.CLcs.NEstat.ML arXiv PDF

Lire avec traduction et IA

Résumé en trois lignes

🎯 Les réseaux de neurones denses subissent une explosion quadratique des coûts lors de l'expansion de la capacité du modèle pour absorber des jeux de données massifs.
🔧 Introduction d'une couche Sparsely-Gated Mixture-of-Experts (MoE) utilisant un routage Noisy Top-K et des pertes auxiliaires d'équilibrage de charge réparties sur des fragments de GPU.
📈 Mise à l'échelle de modèles de langue et de traduction jusqu'à 137 milliards de paramètres, améliorant la perplexité des meilleurs LM publiés de 24 % à 39 % pour un coût de calcul réduit.

Résumé

La capacité d'un réseau de neurones à absorber l'information est limitée par son nombre de paramètres. Le calcul conditionnel, où certaines parties du réseau sont activées au cas par cas pour chaque exemple, a été proposé en théorie comme un moyen d'augmenter considérablement la capacité du modèle sans accroissement proportionnel du calcul. En pratique, cependant, il pose d'importants défis algorithmiques et de performance. Dans ce travail, nous relevons ces défis et concrétisons enfin la promesse du calcul conditionnel, en multipliant par plus de 1000 la capacité du modèle avec seulement des pertes mineures d'efficacité de calcul sur des clusters de GPU modernes. Nous introduisons une couche de mélange d'experts à porte creuse (Sparsely-Gated Mixture-of-Experts, MoE), composée de jusqu'à plusieurs milliers de sous-réseaux feed-forward. Un réseau de routage (gating) entraînable détermine une combinaison creuse de ces experts à utiliser pour chaque exemple. Nous appliquons le MoE aux tâches de modélisation du langage et de traduction automatique, où la capacité du modèle est essentielle pour absorber les vastes quantités de connaissances présentes dans les corpus d'entraînement. Nous présentons des architectures de modèles dans lesquelles un MoE contenant jusqu'à 137 milliards de paramètres est appliqué de manière convolutionnelle entre des couches LSTM empilées. Sur de grands benchmarks de modélisation du langage et de traduction automatique, ces modèles obtiennent des résultats nettement supérieurs à l'état de l'art pour un coût de calcul réduit.

Résumé original (anglais)

The capacity of a neural network to absorb information is limited by its number of parameters. Conditional computation, where parts of the network are active on a per-example basis, has been proposed in theory as a way of dramatically increasing model capacity without a proportional increase in computation. In practice, however, there are significant algorithmic and performance challenges. In this work, we address these challenges and finally realize the promise of conditional computation, achieving greater than 1000x improvements in model capacity with only minor losses in computational efficiency on modern GPU clusters. We introduce a Sparsely-Gated Mixture-of-Experts layer (MoE), consisting of up to thousands of feed-forward sub-networks. A trainable gating network determines a sparse combination of these experts to use for each example. We apply the MoE to the tasks of language modeling and machine translation, where model capacity is critical for absorbing the vast quantities of knowledge available in the training corpora. We present model architectures in which a MoE with up to 137 billion parameters is applied convolutionally between stacked LSTM layers. On large language modeling and machine translation benchmarks, these models achieve significantly better results than state-of-the-art at lower computational cost.

Guide de lecture

Problème et motivation

Les modèles d'apprentissage profond bénéficient grandement de la capacité en paramètres lorsqu'ils sont entraînés sur de vastes jeux de données. Cependant, les architectures standard activent l'intégralité du réseau pour chaque token ou exemple d'entrée. Cela conduit à une augmentation intenable du coût de calcul à mesure que la taille du réseau et le volume des données augmentent.

Bien que le calcul conditionnel (activer uniquement des sous-ensembles d'un modèle par exemple) ait été proposé en théorie, les implémentations antérieures se sont heurtées à des limites sur les accélérateurs modernes en raison de :

La préférence des GPU pour une arithmétique uniforme plutôt que pour des branchements conditionnels.
Le « shrinking batch problem » (problème de rétrécissement de lot), où le routage de sous-ensembles d'exemples laisse les sous-réseaux individuels avec des tailles de lots minuscules et inefficaces.
Des goulets d'étranglement de bande passante réseau entre les périphériques.
Un routage déséquilibré, où une poignée de sous-réseaux dominent tandis que les autres restent sous-utilisés.

Idées clés

Couche MoE à porte creuse (Sparsely-Gated MoE Layer) : Des sous-réseaux feed-forward (« experts ») sont placés entre des couches récurrentes (comme des LSTM empilés). Seule une sélection creuse du top- $k$ des experts est évaluée par position.
Noisy Top-K Gating : Un bruit gaussien ajustable est ajouté aux logits de routage avant la sélection du top- $k$ et l'application du softmax afin de fournir des gradients continus rétropropageables et d'encourager une exploration équilibrée entre les experts.
Parallélisme combiné des données et du modèle : Les couches standard et le mécanisme de porte s'exécutent en parallélisme de données sur l'ensemble des périphériques, tandis que les experts sont partitionnés en parallélisme de modèle sur les GPU. Les lots de tous les périphériques sont combinés pour que les experts reçoivent des mini-lots agrégés suffisamment grands.
Contraintes d'équilibrage : Des pertes auxiliaires explicites sur l'importance et la charge des experts empêchent le réseau de routage de s'effondrer vers des solutions dégénérées qui surutilisent une poignée d'experts.

Méthode (avec les équations clés)

Une couche MoE se compose de $n$ réseaux d'experts feed-forward $E_{1}, \dots, E_{n}$ et d'un réseau de porte $G (x)$ . La sortie de la couche est : $y = i = 1 \sum n G (x)_{i} E_{i} (x)$ Le calcul est évité dès lors que $G (x)_{i} = 0$ .

Noisy Top-K Gating

Pour combiner sparsité et routage différentiable, du bruit est injecté dans les projections linéaires : $H (x)_{i} = (x \cdot W_{g})_{i} + StandardNormal () \cdot Softplus ((x \cdot W_{noise})_{i})$ $KeepTopK (v, k)_{i} = {v_{i} - \infty si v_{i} est parmi les k plus grands \overset{e}{ˊ} l \overset{e}{ˊ} ments de v sinon$ $G (x) = Softmax (KeepTopK (H (x), k))$

Pertes d'équilibrage de charge

Pour imposer des poids de porte totaux égaux, une perte d'importance est définie sur un lot $X$ en utilisant le coefficient de variation ( $C V$ ) : $Importance (X) = x \in X \sum G (x)$ $L_{importance} (X) = w_{importance} \cdot C V (Importance (X))^{2}$ Pour assurer un nombre égal d'exemples par expert, l'estimateur de charge lissé utilise la fonction de répartition normale $Φ$ : $P (x, i) = Φ (\frac{( x \cdot W _{g} ) _{i} - kth_excluding ( H ( x ) , k , i )}{Softplus (( x \cdot W _{noise} ) _{i} )})$ $Load (X)_{i} = x \in X \sum P (x, i)$ $L_{load} (X) = w_{load} \cdot C V (Load (X))^{2}$

Résultats (avec les chiffres les plus importants)

1 Billion Word Benchmark : Un modèle MoE-4096-h à coût de calcul équivalent a atteint une perplexité de test de 34,1 (une réduction de 24 % par rapport aux baselines LSTM comparables en calcul). Un modèle MoE à budget élevé a atteint une perplexité de test de 28,0 en 47 heures sur 32 GPU K40, surpassant le précédent état de l'art de 34,7 (qui nécessitait 59 heures sur 32 GPU K40).
Corpus 100 Billion Word : Le passage à l'échelle jusqu'à 65 536 experts (68 milliards de paramètres) a réduit la perplexité de test de 39 % par rapport à la baseline dense (28,9 contre 47,0), tout en maintenant 0,72 TFLOPS/GPU.
Traduction automatique : Le modèle GNMT augmenté par MoE a atteint 40,56 BLEU sur WMT'14 En $\to$ Fr (+1,34 BLEU par rapport à GNMT sans RL) et 26,03 BLEU sur WMT'14 En $\to$ De (+1,12 BLEU), tout en nécessitant moins de la moitié des opérations par pas de temps.
Traduction multilingue : Un seul modèle MoE multilingue a surpassé une baseline GNMT multilingue sur 11 des 12 paires de langues (jusqu'à +5,84 BLEU sur Coréen $\to$ Anglais) et a réduit la perplexité sur l'ensemble de validation de 19 %.

Limites et questions ouvertes

Une sparsité extrême peut entraîner des rendements décroissants ou des régressions de performance ; la perplexité s'est dégradée lors du passage de 65 536 à 131 072 experts sur le jeu de données de 100 milliards de mots.
Les surcoûts de communication distribuée restent dépendants du ratio dimensions cachées/entrée-sortie, ce qui rend les petites capacités d'experts moins efficaces en termes de bande passante.
L'application de MoE de manière récurrente à l'intérieur des étapes de récurrence d'un RNN brise l'optimisation de déroulement de lot (batch unrolling) à travers les positions de la séquence, nécessitant des stratégies de recalcul des activations.

Termes clés

Sparsely-Gated Mixture-of-Experts Layer (MoE): Une couche de réseau de neurones contenant plusieurs sous-réseaux experts parallèles où un réseau de routage entraînable sélectionne uniquement un sous-ensemble creux d'experts pour chaque exemple d'entrée.
Conditional Computation: Un paradigme architectural où des parties du réseau de neurones sont activées ou désactivées dynamiquement pour chaque exemple afin d'augmenter la capacité en paramètres sans augmentation proportionnelle du calcul.
Noisy Top-K Gating: Un mécanisme de routage qui ajoute un bruit gaussien ajustable aux logits avant de ne conserver que les k meilleures valeurs et d'appliquer un softmax, permettant la rétropropagation et empêchant l'effondrement du routage.
Shrinking Batch Problem: La réduction de la taille de lot effective reçue par les sous-réseaux individuels lorsque les entrées sont réparties de manière creuse entre de nombreux experts, diminuant l'efficacité arithmétique du matériel.
Importance Loss: Une perte auxiliaire pénalisant le coefficient de variation au carré de la somme des valeurs de porte par expert sur un lot afin d'encourager une allocation égale des experts.
Load Loss: Un objectif auxiliaire minimisant la variance dans le nombre estimé d'assignations pondérées par probabilité des tokens entre les experts pour équilibrer la charge matérielle.
Hierarchical Mixture of Experts: Une structure MoE à deux niveaux où un réseau de routage primaire achemine les entrées vers des clusters MoE secondaires afin de réduire le facteur de branchement et le surcoût de communication distribuée.

Questions

Comment le réseau de routage permet-il la rétropropagation à travers une sélection discrète top-k ?

Comme la sélection top-k est appliquée avant le softmax, les valeurs de porte non nulles pour les experts choisis conservent des dérivées continues par rapport aux poids d'entrée, ce qui permet aux gradients de se propager en arrière à travers le réseau de routage.

Comment le shrinking batch problem est-il résolu sur des périphériques distribués ?

Le modèle combine le parallélisme de données et de modèle : les couches standard traitent des mini-lots indépendants sur chaque périphérique, mais les tokens de tous les périphériques sont agrégés avant d'être routés vers des experts stationnaires hébergés sur différents GPU.

Pourquoi le bruit est-il nécessaire dans le réseau de routage ?

Le bruit gaussien injecté favorise l'exploration entre les experts en début d'entraînement et fournit un estimateur lissé pour la perte d'équilibrage de charge, évitant ainsi un effondrement auto-renforçant vers quelques experts.

Que se passe-t-il lorsque le modèle passe à 131 072 experts ?

La perplexité s'est légèrement dégradée, passant de 28,9 (avec 65 536 experts) à 29,2 (avec 131 072 experts), et l'efficacité de calcul a chuté à 0,30 TFLOPS/GPU, probablement en raison de l'extrême sparsité et de tailles de lots fixes non adaptées à l'échelle.

La présentation de cette page a été rédigée par une IA à partir de l'article et peut contenir des erreurs. Consultez l'article pour les détails.

Lire avec traduction et IA