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

DeepMind недавно выпустил статью про улучшение лосса у LLM при том же бюджете на компьют.

Основано всё на том, что трансформеры используют компьют неоптимально: не на все токены нужно прогонять все слои, не на все токены нужно аттендиться.

Роутер в Mixture-of-Experts определяет, в какой MLP подать токен. В этой же работе роутер Mixture-of-Depths определяет, прогонять ли токен через Self-Attention & MLP или оставить его как есть. Получается, что через слой MoD проходит только часть токенов (часть входной последовательности).

То есть на обучении (когда делаем forward pass сразу на всей последовательности)
(1) через каждый слой (= глубину LLM) проходит заранее заданное число токенов K, которое может быть сильно меньше общей длины (авторы 12.5% используют, правда слой MoD чередуется с обычным транcформером, где всё используется)
(2) Self-Attention вычисляется на этом же подмножестве из выбранных top-K токенов

Роутер обучаемый, для каждого токена выдает одно число: его важность для данного слоя.

Ещё одно преимущество подхода: константа K задается заранее, поэтому вычислительный граф при обучении не меняется (статичный). В Self-Attention & MLP всегда пойдет K токенов. Это дает доп выигрыш в скорости.

Таким образом, варьируя K можно делать трейдофф между качеством и скоростью, достигать лучшего лосса при заданном бюджете на компьют.

В общем, архитектура трансформера достаточно неоптимальная и заметное ускорение при том же качестве можно выжимать и на software части, а не только hardware.

@building_singularity