SwitchHead : Accélération des Transformers grâce à une attention basée sur un mélange d'experts

SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention

December 13, 2023

Auteurs: Róbert Csordás, Piotr Piękos, Kazuki Irie, Jürgen Schmidhuber

cs.AI

Résumé

Les couches d'auto-attention coûteuses dans les Transformers modernes nécessitent une mémoire et des calculs quadratiques par rapport à la longueur de la séquence. Les méthodes d'approximation existantes sous-performent généralement et ne parviennent pas à obtenir des accélérations significatives en pratique. Nous présentons ici SwitchHead - une méthode novatrice qui réduit à la fois les besoins en calcul et en mémoire, tout en obtenant une accélération en temps réel, tout en égalant les performances de modélisation linguistique des Transformers de référence avec le même budget de paramètres. SwitchHead utilise des couches de Mixture-of-Experts (MoE) pour les projections de valeur et de sortie, et nécessite 4 à 8 fois moins de matrices d'attention que les Transformers standard. Notre nouvelle attention peut également être combinée avec des couches MLP MoE, aboutissant à un modèle Transformer entièrement MoE efficace appelé "SwitchAll". Notre code est public.

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.