In plain words: Instead of running every attention head, it picks a few specialized attention blocks per token, cutting the attention matrices computed by up to 8 times. The model matched the usual Transformer's language quality using just 44% of the compute and less memory.
Abstract
Despite many recent works on Mixture of Experts (MoEs) for resource-efficient Transformer language models, existing methods mostly focus on MoEs for feedforward layers. Previous attempts at extending MoE to the self-attention layer fail to match the performance of the parameter-matched baseline. Our novel SwitchHead is an effective MoE method for the attention layer that successfully reduces both the compute and memory requirements, achieving wall-clock speedup, while matching the language modeling performance of the baseline Transformer. Our novel MoE mechanism allows SwitchHead to compute up to 8 times fewer attention matrices than the standard Transformer. SwitchHead can also be combined with MoE feedforward layers, resulting in fully-MoE "SwitchAll" Transformers. For our 262M parameter model trained on C4, SwitchHead matches the perplexity of standard models with only 44% compute and 27% memory usage. Zero-shot experiments on downstream tasks confirm the performance of SwitchHead, e.g., achieving more than 3.5% absolute improvements on BliMP compared to the baseline with an equal compute resource.
Róbert Csordás, Piotr Piękos, Kazuki Irie, Jürgen Schmidhuber
arXiv:2312.07987 · cs.LG, cs.CL, cs.NE · submitted Dec 13, 2023 · updated Sep 30, 2024
abstract · pdf · html · Accepted to NeurIPS 2024