Mixture-of-Depths: Dynamische toewijzing van rekenkracht in transformer-gebaseerde taalmodelle

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

Samenvatting

Transformer-gebaseerde taalmodellen verdelen FLOPs uniform over invoerreeksen. In dit werk tonen we aan dat transformers in plaats daarvan kunnen leren om FLOPs (of rekenkracht) dynamisch toe te wijzen aan specifieke posities in een reeks, waarbij de toewijzing langs de reeks wordt geoptimaliseerd voor verschillende lagen over de diepte van het model. Onze methode handhaaft een totaal rekenbudget door het aantal tokens (k) dat kan deelnemen aan de self-attention en MLP-berekeningen in een bepaalde laag te beperken. De te verwerken tokens worden door het netwerk bepaald met behulp van een top-k routeringsmechanisme. Omdat k vooraf is gedefinieerd, gebruikt deze eenvoudige procedure een statische rekenstructuur met bekende tensorformaten, in tegenstelling tot andere conditionele rekentechnieken. Desalniettemin, omdat de identiteiten van de k tokens fluïde zijn, kan deze methode FLOPs niet-uniform verdelen over de tijd- en modeldieptedimensies. Zo is de rekenuitgaven volledig voorspelbaar in totaal, maar dynamisch en contextgevoelig op tokenniveau. Modellen die op deze manier zijn getraind, leren niet alleen om rekenkracht dynamisch toe te wijzen, maar doen dit ook efficiënt. Deze modellen evenaren de basisprestaties voor equivalente FLOPs en wall-clock tijden om te trainen, maar vereisen een fractie van de FLOPs per forward pass, en kunnen tijdens het post-trainingsampling tot wel 50% sneller zijn per stap.

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.