Mischung von Tiefen: Dynamische Zuweisung von Berechnungen in auf Transformer basierenden Sprachmodellen

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

April 2, 2024

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

cs.AI

Zusammenfassung

Transformer-basierte Sprachmodelle verteilen FLOPs gleichmäßig über Eingabesequenzen. In dieser Arbeit zeigen wir, dass Transformer stattdessen lernen können, FLOPs (oder Berechnungen) dynamisch bestimmten Positionen in einer Sequenz zuzuweisen, wobei die Zuweisung entlang der Sequenz für verschiedene Schichten über die Modelltiefe optimiert wird. Unsere Methode setzt ein Gesamtberechnungsbudget durch Begrenzung der Anzahl von Tokens (k) um, die an der Self-Attention und MLP-Berechnungen in einer bestimmten Schicht teilnehmen können. Die zu verarbeitenden Tokens werden vom Netzwerk mithilfe eines Top-k-Routingmechanismus bestimmt. Da k a priori definiert ist, verwendet dieses einfache Verfahren einen statischen Berechnungsgraphen mit bekannten Tensorgrößen, im Gegensatz zu anderen bedingten Berechnungstechniken. Dennoch, da die Identitäten der k Tokens variabel sind, kann diese Methode FLOPs nicht gleichmäßig über die Zeit- und Modelltiefe-Dimensionen verteilen. Somit ist der Berechnungsaufwand insgesamt vollständig vorhersehbar, aber dynamisch und kontextsensitiv auf der Token-Ebene. Modelle, die auf diese Weise trainiert werden, lernen nicht nur, Berechnungen dynamisch zuzuweisen, sondern tun dies auch effizient. Diese Modelle erreichen die Baseline-Performance für äquivalente FLOPS und Trainingszeiten, benötigen jedoch nur einen Bruchteil der FLOPs pro Vorwärtspass und können während der Post-Training-Samplingphase um mehr als 50\% schneller sein.

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.