SwitchHead: Acelerando Transformers con Atención de Mezcla de Expertos
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
Resumen
Las costosas capas de auto-atención en los Transformers modernos requieren memoria y cómputo cuadráticos en función de la longitud de la secuencia. Los métodos de aproximación existentes generalmente tienen un rendimiento inferior y no logran obtener mejoras significativas de velocidad en la práctica. Aquí presentamos SwitchHead, un método novedoso que reduce tanto los requisitos de cómputo como de memoria y logra una aceleración en tiempo real, mientras iguala el rendimiento en modelado de lenguaje de los Transformers de referencia con el mismo presupuesto de parámetros. SwitchHead utiliza capas de Mezcla de Expertos (MoE) para las proyecciones de valor y salida, y requiere de 4 a 8 veces menos matrices de atención que los Transformers estándar. Nuestra nueva atención también puede combinarse con capas MLP de MoE, dando como resultado un modelo Transformer completamente MoE, denominado "SwitchAll", altamente eficiente. Nuestro código es 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.