SwitchHead: Beschleunigung von Transformern durch Mixture-of-Experts-Attention

SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention

December 13, 2023

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

cs.AI

Zusammenfassung

Die kostspieligen Self-Attention-Schichten in modernen Transformern erfordern Speicher und Rechenleistung, die quadratisch mit der Sequenzlänge wachsen. Bestehende Approximationsmethoden schneiden in der Regel schlechter ab und erzielen in der Praxis keine signifikanten Beschleunigungen. Hier präsentieren wir SwitchHead – eine neuartige Methode, die sowohl den Rechen- als auch den Speicherbedarf reduziert und eine Beschleunigung der Echtzeitleistung erreicht, während sie die Sprachmodellierungsleistung von Baseline-Transformern mit demselben Parameterbudget erreicht. SwitchHead verwendet Mixture-of-Experts (MoE)-Schichten für die Wert- und Ausgangsprojektionen und benötigt 4 bis 8 Mal weniger Attention-Matrizen als Standard-Transformer. Unsere neuartige Attention kann auch mit MoE-MLP-Schichten kombiniert werden, was zu einem effizienten vollständig-MoE „SwitchAll“-Transformer-Modell führt. Unser Code ist öffentlich zugänglich.

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.