SwitchHead: Versnelling van Transformers met Mixture-of-Experts Aandacht

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

Samenvatting

De kostbare self-attention lagen in moderne Transformers vereisen geheugen en rekencapaciteit die kwadratisch toenemen met de sequentielengte. Bestaande benaderingsmethoden presteren meestal ondermaats en slagen er niet in om significante snelheidswinst te behalen in de praktijk. Hier presenteren we SwitchHead - een nieuwe methode die zowel de reken- als geheugenvereisten vermindert en een snelheidswinst in werkelijke tijd bereikt, terwijl de taalmodelprestaties van baseline Transformers met hetzelfde parameterbudget worden geëvenaard. SwitchHead gebruikt Mixture-of-Experts (MoE) lagen voor de value- en outputprojecties en vereist 4 tot 8 keer minder aandachtmatrices dan standaard Transformers. Onze nieuwe aandacht kan ook worden gecombineerd met MoE MLP lagen, wat resulteert in een efficiënt volledig-MoE "SwitchAll" Transformer model. Onze code is openbaar.

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.