Смесь глубин: Динамическое выделение вычислительных ресурсов в языковых моделях на основе трансформеров

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

April 2, 2024

Авторы: David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, Adam Santoro

cs.AI

Аннотация

Языковые модели на основе трансформеров равномерно распределяют операции с плавающей запятой (FLOPs) по входным последовательностям. В данной работе мы демонстрируем, что трансформеры могут вместо этого научиться динамически распределять FLOPs (или вычисления) на конкретные позиции в последовательности, оптимизируя распределение вдоль последовательности для различных слоев модели. Наш метод обеспечивает общий бюджет вычислений путем ограничения количества токенов (k), которые могут участвовать в самовнимании и вычислениях MLP на данном слое. Токены для обработки определяются сетью с использованием механизма маршрутизации top-k. Поскольку k определяется заранее, этот простой метод использует статический вычислительный граф с известными размерами тензоров, в отличие от других методов условных вычислений. Тем не менее, поскольку идентификаторы k токенов изменчивы, этот метод может неравномерно расходовать FLOPs по временным и глубинным измерениям модели. Таким образом, расход вычислений полностью предсказуем в общей сумме, но динамичен и контекстно-чувствителен на уровне токенов. Модели, обученные таким образом, не только учатся динамически распределять вычисления, но делают это эффективно. Эти модели соответствуют базовой производительности для эквивалентных FLOPS и времени обучения на стенде, но требуют долю FLOPs на каждый прямой проход и могут быть на 50\% быстрее на шаге во время пост-тренировочной выборки.

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.