Mixture-of-Depths : Allocation dynamique des ressources de calcul dans les modèles de langage basés sur les transformers

Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

April 2, 2024

Auteurs: David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, Adam Santoro

cs.AI

Résumé

Les modèles de langage basés sur Transformer répartissent uniformément les FLOPs sur les séquences d'entrée. Dans ce travail, nous démontrons que les transformers peuvent plutôt apprendre à allouer dynamiquement les FLOPs (ou calcul) à des positions spécifiques dans une séquence, optimisant ainsi l'allocation le long de la séquence pour différentes couches à travers la profondeur du modèle. Notre méthode impose un budget de calcul total en limitant le nombre de tokens (k) pouvant participer aux calculs d'auto-attention et de MLP à une couche donnée. Les tokens à traiter sont déterminés par le réseau à l'aide d'un mécanisme de routage top-k. Puisque k est défini a priori, cette procédure simple utilise un graphe de calcul statique avec des tailles de tenseurs connues, contrairement à d'autres techniques de calcul conditionnel. Néanmoins, puisque les identités des k tokens sont fluides, cette méthode peut dépenser les FLOPs de manière non uniforme à travers les dimensions temporelles et de profondeur du modèle. Ainsi, la dépense de calcul est entièrement prévisible en somme totale, mais dynamique et sensible au contexte au niveau des tokens. Non seulement les modèles entraînés de cette manière apprennent à allouer dynamiquement le calcul, mais ils le font efficacement. Ces modèles égalent les performances de référence pour des FLOPs et temps d'entraînement équivalents, mais nécessitent une fraction des FLOPs par passage avant, et peuvent être jusqu'à 50\% plus rapides à exécuter lors de l'échantillonnage post-entraînement.

English

Transformer-based language models spread FLOPs uniformly across input sequences. In this work we demonstrate that transformers can instead learn to dynamically allocate FLOPs (or compute) to specific positions in a sequence, optimising the allocation along the sequence for different layers across the model depth. Our method enforces a total compute budget by capping the number of tokens (k) that can participate in the self-attention and MLP computations at a given layer. The tokens to be processed are determined by the network using a top-k routing mechanism. Since k is defined a priori, this simple procedure uses a static computation graph with known tensor sizes, unlike other conditional computation techniques. Nevertheless, since the identities of the k tokens are fluid, this method can expend FLOPs non-uniformly across the time and model depth dimensions. Thus, compute expenditure is entirely predictable in sum total, but dynamic and context-sensitive at the token-level. Not only do models trained in this way learn to dynamically allocate compute, they do so efficiently. These models match baseline performance for equivalent FLOPS and wall-clock times to train, but require a fraction of the FLOPs per forward pass, and can be upwards of 50\% faster to step during post-training sampling.