Mixture-of-Depths: Allocazione dinamica delle risorse computazionali nei modelli linguistici basati su transformer

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

April 2, 2024

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

cs.AI

Abstract

I modelli linguistici basati su Transformer distribuiscono i FLOP in modo uniforme lungo le sequenze di input. In questo lavoro dimostriamo che i transformer possono invece imparare ad allocare dinamicamente i FLOP (o il calcolo) a posizioni specifiche di una sequenza, ottimizzando l'allocazione lungo la sequenza per diversi strati in profondità nel modello. Il nostro metodo impone un budget totale di calcolo limitando il numero di token (k) che possono partecipare ai calcoli di self-attention e MLP in un determinato strato. I token da elaborare sono determinati dalla rete mediante un meccanismo di routing top-k. Poiché k è definito a priori, questa semplice procedura utilizza un grafo di calcolo statico con dimensioni tensoriali note, a differenza di altre tecniche di calcolo condizionale. Tuttavia, poiché l'identità dei k token è fluida, questo metodo può impiegare i FLOP in modo non uniforme lungo le dimensioni temporali e di profondità del modello. Pertanto, la spesa computazionale è interamente prevedibile in totale, ma dinamica e sensibile al contesto a livello di token. Non solo i modelli addestrati in questo modo imparano ad allocare dinamicamente il calcolo, ma lo fanno in modo efficiente. Questi modelli eguagliano le prestazioni di riferimento per FLOP e tempi di addestramento equivalenti, ma richiedono una frazione dei FLOP per ogni passaggio in avanti e possono essere fino al 50% più veloci durante il campionamento post-addestramento.

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.