- Mixture-of-Depths (MoD) de Google DeepMind est une méthode qui permet aux modèles de langage transformer de ne pas utiliser les mêmes FLOPs pour tous les tokens : à chaque couche, seuls les tokens nécessaires participent aux calculs d’attention et de MLP
- Un routeur propre à chaque couche produit un poids scalaire pour chaque token ; seuls les tokens compris dans une capacité top-k définie à l’avance traversent le bloc, tandis que les autres le contournent via une connexion résiduelle
- En fixant k à l’avance, le graphe de calcul et la taille des tenseurs restent statiques, tout en permettant une allocation dynamique du calcul par token selon le contexte
- Dans les expériences, à isoFLOP, MoD a obtenu une loss plus faible que le transformer de base pour le même temps d’entraînement, ou a réduit les FLOPs par forward pass à performance égale, accélérant ainsi les étapes d’entraînement et d’échantillonnage
- La meilleure configuration consistait à appliquer des blocs à 12,5 % de capacité un bloc sur deux, et, pour l’échantillonnage auto-régressif, un routage fondé sur un prédicteur a réduit le problème de non-causalité du top-k
Le gaspillage de calcul visé par MoD
- Un modèle de langage transformer classique utilise la même quantité de calcul pour tous les tokens lors du forward pass
- MoD part du principe que tous les tokens et toutes les séquences n’exigent pas le même temps ni le même effort pour la prédiction, et réduit le budget de calcul global tout en allouant le calcul aux tokens qui en ont besoin
- Le calcul conditionnel existant peut créer des graphes de calcul dynamiques, qui ne correspondent pas toujours aux graphes de calcul statiques privilégiés par le matériel actuel
- Si l’utilisateur fixe le budget de calcul total avant l’entraînement, le modèle apprend, à l’intérieur de ce budget, où utiliser le calcul selon les tokens et les couches
- La quantité de calcul, l’usage mémoire et les FLOPs par forward pass peuvent être prédits à l’avance
- Les tokens concernés par le routage varient selon le contexte d’entrée
Structure de routage et méthode de réduction du calcul
- MoD utilise un routeur comme un transformer MoE, mais au lieu de choisir l’un de plusieurs experts, il envoie les tokens vers l’un de deux chemins
- le calcul de bloc standard, incluant l’attention et le MLP
- la connexion résiduelle, qui transmet la valeur telle quelle
- La connexion résiduelle a un faible coût de calcul, et la sortie du bloc est déterminée par l’entrée
- Si la capacité d’un bloc est fixée à une valeur inférieure au nombre total de tokens T, les FLOPs par forward pass diminuent par rapport à un transformer de base
- Par exemple, si la capacité est réduite à T/2, le produit matriciel query-key de l’attention passe de T² à (T/2)², soit environ 25 % des FLOPs
- Une capacité à 0 rend le modèle très rapide, mais dégrade les performances car la plupart des paramètres du transformer ne sont pas utilisés ; une capacité à T revient au transformer de base
- L’objectif est de trouver, entre ces deux extrêmes, un point qui offre des performances égales ou supérieures à celles du transformer de base tout en accélérant les étapes
Routage expert-choice top-k
- À chaque couche l, le routeur génère un poids scalaire pour chaque embedding de token
- Selon la capacité C définie par l’utilisateur, seuls les tokens figurant dans le top-k des tokens de la séquence participent aux calculs d’attention et de MLP
- Les tokens non sélectionnés sautent le calcul du bloc et sont transmis à la couche suivante via la connexion résiduelle
- Cette méthode repose sur l’expert-choice routing
- Avec le token-choice routing, ce sont les tokens qui choisissent leur chemin, ce qui peut créer des problèmes de load balancing
- Avec l’expert-choice routing, chaque chemin choisit ses tokens top-k, ce qui remplit exactement la capacité
- Dans MoD, les chemins étant en pratique au nombre de deux — « calcul du bloc » et « contournement » —, une seule opération top-k permet de séparer les tokens en deux ensembles
- Les poids du routeur sont multipliés par la sortie du bloc, se trouvent sur le gradient path et sont influencés par la descente de gradient durant l’entraînement du modèle de langage
Le problème de non-causalité dans l’échantillonnage
- L’opération top-k de l’expert-choice routing est non causale
- Pour décider si un token entre dans le top-k, il faut aussi connaître les poids du routeur des tokens suivants
- Lors de l’échantillonnage auto-régressif, les informations des tokens futurs ne peuvent pas être utilisées
- Deux méthodes ont été expérimentées pour éviter ce problème
- Une loss auxiliaire de binary cross-entropy pour que la sigmoid de la sortie du routeur prédise si le token sera sélectionné dans le top-k
- Un petit MLP predictor auxiliaire recevant la même entrée que le routeur, mais avec stop-gradient appliqué
- La première méthode affecte l’objective principal de modélisation du langage d’environ 0,2 à 0,3 %
- La deuxième méthode, fondée sur un predictor, n’a pas affecté l’objective de modélisation du langage et n’a pas eu d’effet significatif sur la step speed
- La tâche auxiliaire de prédiction atteint expérimentalement rapidement 99 % de précision
Résultats expérimentaux et observations
- Les premières expériences ont consisté à rechercher les hyperparamètres de MoD avec un budget total de 6e18 FLOPs
- La configuration MoD optimale appliquait le routage un bloc sur deux et fixait top-k à 256
- Pour une longueur de séquence de 2048, cela représente 256 tokens, soit seulement 12,5 %, qui traversent l’attention et le MLP
- Les 1792 autres tokens, soit 87,5 %, contournent le bloc
- La variante MoD à 220M paramètres a obtenu des performances légèrement meilleures que le baseline optimal à isoFLOP, tout en étant plus de 60 % plus rapide par étape pendant l’entraînement ; sur le même matériel, le wall-clock time total d’entraînement était presque identique
- Dans les analyses isoFLOP à 6e18, 2e19 et 1e20 FLOPs, les modèles FLOP-optimal de MoD avaient également tendance à comporter davantage de paramètres que le baseline
- Il existe des variantes de MoD qui atteignent les mêmes performances, ou une loss plus faible, avec moins de FLOPs par forward pass
- L’amélioration de la vitesse par étape vient de deux facteurs
- Comme certains tokens contournent le bloc, les FLOPs per parameter diminuent à taille de modèle identique
- Il est possible de choisir une variante MoD plus petite qui obtient les mêmes performances que le baseline
- MoD avec stochastic routing a montré des performances nettement inférieures au transformer de base et à MoD standard, ce qui indique l’importance d’un routage appris
- L’évaluation auto-régressive a utilisé 256 000 séquences et 500M tokens de held-out data ; le passage à un routage fondé sur un predictor n’a entraîné qu’une faible baisse de performance
- MoD peut aussi être combiné avec MoE pour former Mixture-of-Depths-and-Experts (MoDE)
- Le staged MoDE décide, avant l’attention, s’il faut envoyer les tokens vers le bloc ou les contourner
- L’integrated MoDE intègre le routage MoD en ajoutant un expert « no-op » parmi les experts MLP existants
- L’analyse du routage a montré que certains tokens traversent fréquemment plusieurs blocs, tandis que d’autres les contournent chaque fois que possible
- Des résultats préliminaires d’analyse indiquent que les tokens qui traversent souvent les blocs sont corrélés à une entropie plus élevée de la prédiction de sortie, ce qui suggère qu’il pourrait s’agir de tokens plus difficiles à prédire
1 commentaires
Commentaires Hacker News
Un routage plus complexe va très probablement devenir plus courant
En particulier, je pense qu’on finira un jour par aller vers du routage récursif qui repasse par un mélange d’experts. À l’avenir, le « chain-of-thought » se produira sans doute de manière récursive à l’intérieur du modèle
Côté idée d’« expert », le papier Sparse Universal Transformers de l’an dernier est encore plus proche : il combine Universal Transformer et mélange clairsemé d’experts, avec un mécanisme de gating qui décide quels blocs transformer utiliser et dans quel ordre
Ce n’est pas mon domaine de spécialité, mais si j’ai bien compris c’est difficile à entraîner correctement, et il faut plus de calcul total à l’inférence pour obtenir des résultats comparables à un transformer classique. Cela dit, c’est une direction intéressante, et le fait qu’il y ait une borne supérieure au nombre d’étapes de calcul par token est à mon avis l’un des gros défauts de l’architecture transformer classique
Si on choisit un nombre aléatoire de fois ou qu’on teste plusieurs profondeurs de récursion, la sortie devient « floue ». Autrement dit, on ne sait plus si la sortie d’une couche doit fournir une information cruciale au résultat final, ou produire la meilleure entrée possible pour le prochain tour de récursion
Le raisonnement plus profond viendra sans doute davantage de la dynamique au niveau du token qu’au niveau de la couche. Il y a par exemple le récent papier Quiet-STaR, dans lequel le modèle produit des tokens de justification qu’il écarte ensuite : https://arxiv.org/abs/2403.09629
C’est peut-être le papier le plus important de 2024
L’idée de vouloir un modèle qui n’utilise pas la même quantité de calcul pour tous les tokens existe depuis longtemps, mais c’est la première fois que je vois un mécanisme vraiment convaincant pour y parvenir
Version d’introduction simplifiée :
Imaginez un assistant intelligent capable de comprendre et de traiter la parole. En temps normal, cet assistant accorde la même attention à chaque mot, quelle que soit son importance pour le sens global
Maintenant, supposons qu’on ait trouvé un moyen d’apprendre à cet assistant à utiliser ses « ressources mentales » plus intelligemment. Au lieu de donner la même attention à tous les mots, il se concentre davantage sur ceux qui sont les plus importants pour comprendre le sens, et ajuste ce focus à la volée selon le contexte
On impose aussi une limite au total des « ressources mentales » qu’il peut utiliser à un instant donné, pour éviter la surcharge. C’est comme lui donner un budget et lui dire : « tu ne peux consacrer des ressources qu’à un certain nombre de mots à la fois »
L’assistant doit alors décider quels mots sont les plus importants
Même avec cette limite, il ajuste l’utilisation de ses ressources de façon souple. Il en consacre plus à certains mots, moins à d’autres, et hiérarchise selon la situation
En apprenant de cette manière, l’assistant devient capable de prêter attention de façon intelligente tout en restant efficace. Il comprend aussi bien qu’un assistant qui traite tous les mots de la même façon, mais en utilisant moins de ressources globales, ce qui accélère fortement les réponses et le traitement de nouvelles informations
J’ai fait un petit récapitulatif ici à partir de ce que j’ai compris :
https://lifeinthesingularity.com/p/googles-breakthroughs-in-...
C’est très proche du mélange d’experts. Sauf qu’au lieu de router les tokens vers plusieurs experts, on les envoie vers « un expert unique qu’on peut sauter dynamiquement »
« This is more computationally efficient than performing a full content-based lookup across an entire memory buffer for each step in the future, and could be one step towards drastically increasing the context-length available for making a prediction. »
Est-ce que c’est ce qui permettrait une fenêtre de contexte de 10 millions de tokens ? Ou est-ce qu’ils parlent simplement de fenêtres de contexte plus longues à l’avenir ?
Après avoir essayé de comprendre et d’implémenter quelques algorithmes RASP, j’en suis arrivé à la conclusion que certaines fonctions nécessitent un certain nombre de couches transformer pour fonctionner
En suivant cette logique, on voit que les fonctions apprises par les transformers peuvent être réparties sur plusieurs têtes. Répéter ce type de fonction pourrait être très utile pour comprendre et résoudre un problème, mais les méthodes d’inférence actuelles ne permettent pas d’exécuter de façon répétée des groupes successifs de têtes. Cet article semble vraiment aller dans une direction prometteuse
[1] https://arxiv.org/pdf/2106.06981.pdf
[2] https://www.youtube.com/watch?v=t5LjgczaS80
Le seul inconvénient au fait que les LLM progressent si vite, c’est peut-être que les papiers sortent plus rapidement que quiconque — en tout cas en dehors de Google — ne peut apprendre et tester les améliorations
Quand j’ai commencé le deep learning, c’était l’époque de ReLU et du dropout, et on pouvait modifier une ou deux lignes de code sur une GeForce 1080 grand public pour tester en quelques heures si ça améliorait les choses. Aujourd’hui, on a l’impression qu’il faut attendre des semaines que des équipes comme Mistral aient le temps de faire l’essai
Moi, je me concentre sur les approches de quantification et je teste avec des GPU anciens, de générations précédentes
hu-po fait des reviews approfondies d’articles d’IA en livestream
Je recommande vivement ; il couvre ici le papier mixture-of-depths dont il est question : https://www.youtube.com/watch?v=Teru_qIdB8Y