SwitchHead: エキスパート混合アテンションによるTransformerの高速化

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

要旨

現代のTransformerにおける高コストなセルフアテンションレイヤーは、シーケンス長に対して二次的なメモリと計算を必要とします。既存の近似手法は通常、性能が劣り、実際の高速化を実現できません。本論文では、SwitchHeadを提案します。これは、計算量とメモリ要件を削減し、実時間での高速化を達成しながら、同じパラメータ予算を持つベースラインTransformerの言語モデリング性能に匹敵する新しい手法です。SwitchHeadは、値と出力の射影にMixture-of-Experts（MoE）レイヤーを使用し、標準的なTransformerよりも4倍から8倍少ないアテンションマトリックスを必要とします。この新しいアテンションは、MoE MLPレイヤーと組み合わせることも可能で、効率的な完全MoE「SwitchAll」Transformerモデルを実現します。私たちのコードは公開されています。

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.