SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention

现代Transformer中昂贵的自注意力层需要的内存和计算量与序列长度的平方成正比。现有的近似方法通常表现不佳，实践中无法获得显著的加速。在这里，我们提出了SwitchHead——一种新颖的方法，它减少了计算和内存需求，并实现了墙钟速度提升，同时与具有相同参数预算的基准Transformer的语言建模性能相匹配。SwitchHead使用专家混合层（MoE）用于值和输出投影，并且所需的注意力矩阵比标准Transformer少4到8倍。我们的新型注意力还可以与MoE MLP层相结合，从而产生高效的完全MoE“SwitchHead”Transformer模型。我们的代码是公开的。

论文旨在解决现代Transformer中昂贵的自注意力层对内存和计算的需求随序列长度呈二次增长的问题。

SwitchHead是一种新颖的方法，通过使用少量的Mixture-of-Experts（MoE）层，同时降低计算和内存需求并实现时间上的加速，达到与基线Transformer相同参数预算的语言建模性能。

SwitchHead使用MoE层进行值和输出投影，并且比标准Transformer需要4到8倍更少的注意力矩阵。该方法还可以与MoE MLP层相结合，形成高效的完全MoE“SwitchHead”Transformer模型。研究者公开了他们的代码。

与这篇论文相关的研究包括：1. 'Attention is All You Need' 2. 'Reformer: The Efficient Transformer' 3. 'Longformer: The Long-Document Transformer'等。