[Traduction] MoD (Mixture-of-Depths) : une approche pour optimiser le calcul des modèles de langage basés sur Transformer, et MoDE (MoD+MoE)

(discuss.pytorch.kr)

7 points par ninebow 2024-04-13 | Aucun commentaire pour le moment. | Partager sur WhatsApp

PyTorchKR

Des méthodes MoE (Mixture-of-Experts), qui traitent les tokens d’entrée en sélectionnant un ou plusieurs blocs « experts », comme Jamba, Qwen1.5-MoE ou encore DBRX, sont de plus en plus utilisées ces derniers temps. Ces approches MoE offrent de meilleures performances en allouant efficacement les ressources de calcul grâce à un contrôle dynamique d’une partie de très grands LLM.
Si le MoE ajustait dynamiquement la « largeur (width) » du modèle en choisissant certains experts parmi plusieurs, le MoD (Mixture-of-Depths) présenté ici adopte une approche qui ajuste dynamiquement la « profondeur (depth) » du modèle. Autrement dit, pour certains tokens, le modèle n’exécute pas tous les calculs sur l’ensemble des couches, mais effectue le calcul de façon dynamique en sautant certaines couches. Voyons cela ensemble.

MoD (Mixture-of-Depths) : une approche pour optimiser le calcul des modèles de langage basés sur Transformer (Mixture-of-Depths: Dynamically allocating compute in transformer-based language models)

Introduction

Traditionnellement, les modèles Transformer allouent les mêmes ressources de calcul à toutes les parties de la séquence d’entrée. Mais l’article sur le Transformer MoD (Mixture-of-Depths) présenté ici tente d’allouer dynamiquement des ressources de calcul à certaines positions de séquence dans le Transformer, et de faire varier cette allocation au fil des couches. Pour mettre en œuvre cette stratégie Mixture-of-Depths (MoD), les auteurs utilisent un mécanisme de routage top-k pour décider quels tokens traiter et ainsi maîtriser la quantité totale de calcul. Cette technique offre des performances comparables, voire supérieures, à celles des modèles existants tout en demandant bien moins de FLOPs, et elle permet jusqu’à 50 % d’accélération lors de l’échantillonnage après l’entraînement.

Architecture Transformer présentée dans l’article Attention is All You Need

Les modèles Transformer récents et les LLM (Large Language Models) continuent de croître en taille et en complexité. Ces modèles utilisent des milliards de paramètres pour se rapprocher de capacités de compréhension et de génération du langage au niveau humain, avec à la clé des performances remarquables. Cependant, ces progrès s’accompagnent de coûts de calcul et de consommation énergétique importants, et l’augmentation de la taille des modèles accroît fortement les temps d’entraînement et d’inférence.

Cela vient du fait que le Transformer effectue la même quantité de calcul pour tous les tokens de la séquence d’entrée. Or, l’idée de départ est que tous les tokens ne fournissent pas forcément la même quantité d’information : certains jouent un rôle central dans la compréhension du contexte, tandis que d’autres sont relativement moins importants.

Le Transformer MoD alloue dynamiquement le calcul à chaque token, ce qui permet au modèle de concentrer ses ressources là où elles sont nécessaires. Cela réduit le coût global de calcul du modèle tout en améliorant l’efficacité du traitement de l’information, en accordant davantage d’attention aux tokens importants. En particulier, pour les grands modèles comme les LLM, l’approche MoD peut devenir un moyen important d’économiser du temps et des ressources lors de l’entraînement comme de l’inférence.

Présentation de la technique MoD (Mixture-of-Depths)

La technique MoD intègre, dans chaque couche du Transformer, un mécanisme de routage top-k qui sélectionne dynamiquement quels tokens seront traités. Ainsi, seule une partie des tokens sélectionnés suit le chemin de calcul complet, tandis que les autres sautent une ou plusieurs couches.

Figure 1 / À gauche : vue d’ensemble du Transformer MoD (Mixture-of-Depths) : à l’image de l’architecture MoE (Mixture-of-Experts), un routeur sélectionne le chemin de calcul. Mais, contrairement au MoE, il décide ici si le bloc standard (Self-Attention et MLP) doit être exécuté ou non. Si un token d’entrée ($X_i$) n’active pas le calcul complet selon la décision du routeur, il ne passe pas par le calcul de cette couche.
Figure 1 / À droite : résultat des décisions de routage pour une séquence courte réduite à 64 tokens. (Axe X : séquence, Axe Y : couche) Le violet indique que le calcul de la couche a été entièrement exécuté, et l’orange qu’il a été sauté.

Comme l’explique la figure ci-dessus, le MoD implémente un routeur prédictif dans l’architecture Transformer afin de décider, pour chaque token, si le calcul de la couche doit être effectué. Cela permet de concentrer les ressources de calcul sur les tokens relativement « importants ».

Entraînement du Transformer MoD

Le processus d’entraînement du Transformer MoD est, dans son principe, similaire à celui d’un Transformer standard. La séquence d’entrée est fournie au modèle, qui la transmet à travers les couches en effectuant les calculs nécessaires. Cependant, le Transformer MoD inclut une étape supplémentaire dans laquelle le routeur décide, pour chaque token, du chemin de calcul à suivre.

Calcul des poids du routeur : à chaque couche du modèle, le routeur calcule un poids pour chaque token. Ce poids représente la quantité de calcul que le token doit recevoir et repose sur son importance.
Sélection des tokens prioritaires : sur la base des poids calculés, le routeur sélectionne, pour chaque couche, les tokens les mieux classés. Ces tokens sont dirigés vers le chemin où l’intégralité du calcul est effectuée.
Attribution du chemin de calcul : les tokens sélectionnés sont affectés au chemin exécutant le calcul standard (Self-Attention et MLP). Les autres tokens sont transmis via la connexion résiduelle et passent à la couche suivante sans calcul supplémentaire.
Calcul de la perte et rétropropagation : la sortie du modèle est évaluée par la fonction de perte finale, puis les gradients sont rétropropagés dans le modèle afin de mettre à jour les poids. Le routeur est lui aussi entraîné dans ce processus, ce qui lui permet de prendre des décisions de routage plus efficaces au fil du temps.

Méthodes de routage du Transformer MoD (Routing Schemes)

L’implémentation du routage constitue le cœur du Transformer MoD. Les auteurs ont comparé trois méthodes de routage et ont constaté que le routage des tokens en mode Expert-Choice MoD offrait les meilleures performances. Comparons-le avec les autres méthodes présentées :

Routage par choix du token (Token-choice)

Dans le routage par choix du token, chaque token choisit lui-même le chemin qu’il va emprunter. Le modèle génère, pour chaque token, une distribution de probabilité sur différents chemins de calcul possibles (par exemple, différentes couches ou différents blocs de traitement), puis envoie le token vers le chemin correspondant selon cette distribution.

Le routage par choix du token permet à chaque token de sélectionner le chemin de calcul le plus adapté, ce qui offre un routage très flexible. En outre, le modèle peut prendre en compte de manière plus fine l’importance ou le contexte de chaque token, ce qui permet une optimisation à l’échelle individuelle.

Routage par choix de l’expert (Expert-choice)

Dans le routage par choix de l’expert, le modèle sélectionne directement les tokens que chaque chemin (ou « expert ») va traiter. Cela se fait à l’aide d’un mécanisme top-k qui choisit les tokens les mieux classés à partir des poids attribués par le routeur.

Avec ce routage par choix de l’expert, les tokens peuvent être répartis plus uniformément entre les différents chemins, ce qui équilibre la charge de calcul et permet d’utiliser plus efficacement les ressources du modèle. De plus, comme les tokens retenus sont clairement déterminés, la prévisibilité et la stabilité du modèle s’en trouvent améliorées.

Expert-choice MoD

L’Expert-choice MoD applique au Transformer MoD le concept de routage par choix de l’expert. Dans cette approche, le modèle sélectionne, pour chaque couche ou étape de traitement, les tokens prioritaires, et seuls ces tokens suivent le chemin où le calcul est exécuté. Les autres empruntent le chemin moins coûteux via la connexion résiduelle.

Comme seuls les tokens sélectionnés sont effectivement calculés, la quantité totale de FLOPs consommés par le modèle diminue fortement. En outre, cette approche peut maintenir, voire améliorer, les performances du modèle en garantissant le calcul complet pour les tokens importants tout en minimisant les calculs inutiles. Enfin, grâce au routage dynamique assuré par le routeur, le modèle peut allouer ses ressources de calcul de manière souple selon les caractéristiques de la séquence d’entrée.

Mise en œuvre du routage

La mise en œuvre du routage peut se décomposer en trois grandes étapes :

Détermination des embeddings de tokens et des poids du routeur : pour chaque couche du modèle, on dispose d’un ensemble d’embeddings de tokens $X^i$ pour une séquence donnée. Pour chaque embedding de token $x^l_i$, le routeur calcule, via une transformation linéaire, un poids scalaire $r^l_i$ qui représente la probabilité de participation de ce token au calcul.
Détermination des poids prioritaires : à partir de l’ensemble des poids du routeur des tokens $R^l$, on cherche le percentile correspondant au $\beta$-ième poids, noté $P_\beta(R^l)$. Ici, $\beta$ détermine la proportion de tokens que le modèle traitera à chaque couche. Ainsi, les tokens dont le poids du routeur est supérieur à $P_\beta(R^l)$ participent au calcul.
Sélection du chemin de calcul : pour chaque token, si le poids du routeur $r^l_i$ est supérieur ou égal à $P_\beta(R^l)$, le calcul de la couche est exécuté ; sinon, le token est transmis directement à la couche suivante. Cela permet de n’effectuer le calcul que pour les tokens sélectionnés, tandis que les autres suivent un chemin sans coût de calcul supplémentaire.

Lors de l’implémentation de ce routage, les points suivants doivent être particulièrement pris en compte :

Décision dynamique fondée sur les poids : le routage est déterminé dynamiquement à partir du poids de chaque token, ce qui reflète son importance. Les tokens importants reçoivent davantage de calcul afin d’optimiser les performances du modèle, tandis que les tokens moins importants réduisent les coûts en passant par le chemin résiduel.
Efficacité de calcul et optimisation des performances : grâce à ce mécanisme de routage, le Transformer MoD peut utiliser efficacement ses ressources de calcul tout en maintenant, voire en améliorant, ses performances. En particulier, le traitement des seuls tokens importants dans les couches les plus coûteuses réduit la consommation globale de FLOPs.
Simplicité et efficacité de l’implémentation : le routage repose sur une transformation linéaire et sur un calcul de percentile, ce qui le rend relativement simple et efficace. Cela permet une allocation dynamique des ressources de calcul sans ajouter de surcharge importante lors de l’entraînement ou de l’inférence.

Exemple de fonctionnement du Transformer MoD

À partir de ce que nous avons vu jusqu’ici, prenons un exemple pour expliquer le fonctionnement du Transformer MoD.

Supposons par exemple qu’une séquence d’entrée comporte 100 tokens et que le routeur fixe la capacité à 10. Dans ce cas, à chaque couche, le routeur sélectionne les 10 tokens les plus importants, et seuls ceux-ci sont traités via le chemin de calcul standard.

Les 90 tokens restants sautent le calcul et sont directement transmis à la couche suivante via la connexion résiduelle. Grâce à ce processus, le Transformer MoD réduit la consommation totale de FLOPs tout en garantissant le calcul nécessaire pour les tokens importants.

Comparaison des performances du Transformer MoD

Performances à l’entraînement

Optimisation des hyperparamètres : des expériences d’entraînement d’une variante du Transformer MoD alternant des blocs de routage à capacité de 12,5 % avec des blocs d’attention complets ont permis d’identifier le réglage d’hyperparamètres optimal. Cette configuration a offert de meilleures performances que la baseline optimale dans une analyse isoFLOP.

Efficacité de calcul : certaines variantes de MoD ont atteint les mêmes performances que les modèles existants avec moins de FLOPs, tout en affichant une vitesse de traitement supérieure de 66 %.

Analyse du routage

Parcimonie des décisions de routage : avec un routage à capacité de 12,5 %, la plupart des tokens contournaient le bloc et seule une petite partie des tokens importants le traversait. Cela se retrouvait également dans la distribution des poids du routeur, le réseau ayant parfois tendance à privilégier certains tokens pour le routage selon la profondeur.

Performances en inférence

Inférence auto-régressive : lors du passage d’un schéma de routage top-k non causal à l’entraînement à une approche causale fondée sur la prédiction en inférence, la baisse de performance a été presque inexistante. Cela semble s’expliquer par le fait que ce problème de prédiction est facile à apprendre, avec plus de 97 % de précision atteints dès le début de l’entraînement.

Mixture-of-Depths-and-Experts (MoDE) = MoD + MoE

MoDE (Mixture-of-Depths-and-Experts) désigne l’intégration de la technique MoD (Mixture-of-Depths) et de la technique MoE (Mixture of Experts). Cette approche vise à combiner les avantages des deux méthodes afin d’optimiser simultanément l’efficacité de calcul et les performances des modèles Transformer. Là où le MoD choisit, pour chaque token, d’effectuer un calcul plus profond ou plus superficiel, le MoE traite le token via l’« expert » le plus adapté parmi plusieurs. MoDE intègre ces deux approches pour permettre à un token soit de sauter certains blocs, soit d’exécuter sélectivement un calcul plus profond, tout en choisissant parmi plusieurs experts le mode de traitement optimal.

Mise en œuvre de MoDE

MoDE propose deux grands modes d’implémentation : le MoDE par étapes (Staged MoDE) et le MoDE intégré (Integrated MoDE).

Staged MoDE (MoDE par étapes) : le MoDE par étapes applique d’abord la technique MoD pour déterminer le chemin de calcul de chaque token, puis applique la technique MoE pour router, vers les différents blocs experts, les tokens traités sur ce chemin sélectionné. Le modèle peut ainsi d’abord ajuster la profondeur du calcul, puis choisir l’expert optimal à chaque étape.
Integrated MoDE (MoDE intégré) : le MoDE intégré réunit les techniques MoD et MoE dans une seule étape de routage. Autrement dit, lors du routage d’un token, une seule décision permet de sélectionner simultanément la profondeur du chemin de calcul et l’expert chargé du traitement. Cette approche peut être plus simple et plus efficace à implémenter, tout en offrant l’avantage de déterminer en une fois le chemin de calcul et l’expert les plus adaptés à chaque token.

Avantages de MoDE

Efficacité de calcul : MoDE gère efficacement le coût global de calcul du modèle en décidant dynamiquement, pour chaque token, de la quantité de calcul nécessaire et de l’expert qui le traitera.
Optimisation des performances : en attribuant chaque token au chemin de calcul et à l’expert les plus adaptés, MoDE maximise les performances du modèle. Cela peut être particulièrement utile pour résoudre des problèmes complexes ou traiter des données de nature variée.
Flexibilité et extensibilité : MoDE offre la flexibilité nécessaire pour s’appliquer à diverses architectures de modèles et à différents types de tâches, tout en fournissant un framework extensible capable d’améliorer à la fois l’efficacité de calcul et les performances.

MoDE présente une direction prometteuse pour l’évolution future des modèles Transformer et apporte une contribution importante aux recherches qui cherchent à maximiser les performances des modèles tout en minimisant leur coût de calcul.

Pour aller plus loin

https://arxiv.org/abs/2404.02258

Ce texte s’appuie sur un résumé réalisé à l’aide d’un modèle GPT ; il peut donc contenir des formulations qui diffèrent du contenu ou de l’intention du texte original. Si le sujet vous intéresse, n’hésitez pas à consulter également la source. Si vous remarquez des passages maladroits ou incorrects pendant votre lecture, merci de le signaler en commentaire.

⚠️Publicité⚠️ : Ce texte préparé par le groupe d’utilisateurs PyTorch Corée vous a-t-il été utile ? Si vous devenez membre, nous vous enverrons les principaux articles par e-mail ! (Le réglage par défaut est Weekly, mais vous pouvez aussi passer à Daily.)