Mixture-of-Depths: Alocação dinâmica de computação em modelos de linguagem baseados em transformers

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

April 2, 2024

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

cs.AI

Resumo

Modelos de linguagem baseados em Transformers distribuem FLOPs uniformemente ao longo das sequências de entrada. Neste trabalho, demonstramos que os Transformers podem, em vez disso, aprender a alocar dinamicamente FLOPs (ou computação) para posições específicas em uma sequência, otimizando a alocação ao longo da sequência para diferentes camadas ao longo da profundidade do modelo. Nosso método impõe um orçamento total de computação limitando o número de tokens (k) que podem participar dos cálculos de self-attention e MLP em uma determinada camada. Os tokens a serem processados são determinados pela rede usando um mecanismo de roteamento top-k. Como k é definido a priori, este procedimento simples utiliza um grafo de computação estático com tamanhos de tensor conhecidos, ao contrário de outras técnicas de computação condicional. No entanto, como as identidades dos k tokens são fluidas, este método pode gastar FLOPs de forma não uniforme ao longo das dimensões de tempo e profundidade do modelo. Assim, o gasto de computação é totalmente previsível no total, mas dinâmico e sensível ao contexto no nível do token. Não apenas os modelos treinados dessa forma aprendem a alocar computação dinamicamente, mas o fazem de forma eficiente. Esses modelos correspondem ao desempenho de linha de base para FLOPs e tempos de treinamento equivalentes, mas exigem uma fração dos FLOPs por passagem direta e podem ser até 50% mais rápidos para avançar durante a amostragem pós-treinamento.

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.