SwitchHead: Acelerando Transformers com Atenção de Mistura de Especialistas

SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention

December 13, 2023

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

cs.AI

Resumo

As camadas de auto-atenção dispendiosas nos Transformers modernos exigem memória e computação quadráticas em relação ao comprimento da sequência. Os métodos de aproximação existentes geralmente têm desempenho inferior e não conseguem obter ganhos significativos de velocidade na prática. Aqui apresentamos o SwitchHead - um método inovador que reduz tanto os requisitos de computação quanto de memória e alcança aceleração em tempo real, enquanto iguala o desempenho de modelagem de linguagem dos Transformers de base com o mesmo orçamento de parâmetros. O SwitchHead utiliza camadas de Mistura de Especialistas (MoE) para as projeções de valor e saída e requer de 4 a 8 vezes menos matrizes de atenção do que os Transformers padrão. Nossa nova abordagem de atenção também pode ser combinada com camadas MLP MoE, resultando em um modelo Transformer totalmente MoE eficiente, chamado "SwitchAll". Nosso código é público.

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.