Mezcla-de-Profundidades: Asignación dinámica de cómputo en modelos de lenguaje basados en transformadores

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

Resumen

Los modelos de lenguaje basados en Transformers distribuyen los FLOPs de manera uniforme a lo largo de las secuencias de entrada. En este trabajo demostramos que, en cambio, los Transformers pueden aprender a asignar dinámicamente FLOPs (o cómputo) a posiciones específicas de una secuencia, optimizando la asignación a lo largo de la secuencia para diferentes capas a través de la profundidad del modelo. Nuestro método impone un presupuesto total de cómputo limitando el número de tokens (k) que pueden participar en los cálculos de autoatención y MLP en una capa dada. Los tokens que se procesan son determinados por la red mediante un mecanismo de enrutamiento top-k. Dado que k se define a priori, este procedimiento simple utiliza un gráfico de cálculo estático con tamaños de tensor conocidos, a diferencia de otras técnicas de cómputo condicional. Sin embargo, dado que las identidades de los k tokens son fluidas, este método puede gastar FLOPs de manera no uniforme a lo largo de las dimensiones de tiempo y profundidad del modelo. Por lo tanto, el gasto de cómputo es completamente predecible en su totalidad, pero dinámico y sensible al contexto a nivel de token. No solo los modelos entrenados de esta manera aprenden a asignar cómputo dinámicamente, sino que lo hacen de manera eficiente. Estos modelos igualan el rendimiento de referencia para FLOPs y tiempos de entrenamiento equivalentes, pero requieren una fracción de los FLOPs por paso hacia adelante y pueden ser hasta un 50\% más rápidos durante el muestreo posterior al entrenamiento.

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.