SwitchHead: Ускорение трансформеров с помощью внимания на основе смеси экспертов

SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention

December 13, 2023

Авторы: Róbert Csordás, Piotr Piękos, Kazuki Irie, Jürgen Schmidhuber

cs.AI

Аннотация

Дорогостоящие слои самовнимания в современных трансформерах требуют памяти и вычислений, квадратично зависящих от длины последовательности. Существующие методы аппроксимации обычно демонстрируют низкую производительность и не обеспечивают значительного ускорения на практике. В данной работе мы представляем SwitchHead — новый метод, который снижает требования к вычислениям и памяти, обеспечивает ускорение в реальном времени, при этом сохраняя качество языкового моделирования на уровне базовых трансформеров с тем же бюджетом параметров. SwitchHead использует слои "Смеси экспертов" (MoE) для проекций значений и выходных данных и требует в 4–8 раз меньше матриц внимания по сравнению со стандартными трансформерами. Наш новый подход к вниманию также может быть объединен с MoE-слоями MLP, что приводит к созданию эффективной полностью MoE-модели трансформера "SwitchAll". Наш код доступен публично.

English

The costly self-attention layers in modern Transformers require memory and compute quadratic in sequence length. Existing approximation methods usually underperform and fail to obtain significant speedups in practice. Here we present SwitchHead - a novel method that reduces both compute and memory requirements and achieves wall-clock speedup, while matching the language modeling performance of baseline Transformers with the same parameter budget. SwitchHead uses Mixture-of-Experts (MoE) layers for the value and output projections and requires 4 to 8 times fewer attention matrices than standard Transformers. Our novel attention can also be combined with MoE MLP layers, resulting in an efficient fully-MoE "SwitchAll" Transformer model. Our code is public.