Mixture-of-Depths : Allocation dynamique des ressources de calcul dans les modèles de langage basés sur les transformersMixture-of-Depths: Dynamically allocating compute in transformer-based
language models
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.