SwitchHead: Accelerazione dei Transformer con Attenzione a Miscela di Esperti

SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention

December 13, 2023

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

cs.AI

Abstract

I costosi strati di self-attention nei moderni Transformer richiedono memoria e calcoli quadratici rispetto alla lunghezza della sequenza. I metodi di approssimazione esistenti generalmente ottengono prestazioni inferiori e non riescono a ottenere significativi miglioramenti di velocità nella pratica. Qui presentiamo SwitchHead, un metodo innovativo che riduce sia i requisiti di calcolo che di memoria e ottiene un'accelerazione in termini di tempo reale, mantenendo le prestazioni di modellazione del linguaggio dei Transformer di riferimento con lo stesso budget di parametri. SwitchHead utilizza strati Mixture-of-Experts (MoE) per le proiezioni dei valori e degli output e richiede da 4 a 8 volte meno matrici di attenzione rispetto ai Transformer standard. La nostra nuova attenzione può anche essere combinata con strati MLP MoE, risultando in un efficiente modello Transformer completamente MoE chiamato "SwitchAll". Il nostro codice è pubblico.

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.