SLAB: Transformadores eficientes con atención lineal simplificada y normalización por lotes re-parametrizada progresiva

SLAB: Efficient Transformers with Simplified Linear Attention and Progressive Re-parameterized Batch Normalization

May 19, 2024

Autores: Jialong Guo, Xinghao Chen, Yehui Tang, Yunhe Wang

cs.AI

Resumen

Los Transformers se han convertido en arquitecturas fundamentales tanto para tareas de procesamiento de lenguaje natural como de visión por computadora. Sin embargo, el alto costo computacional hace que sea bastante difícil implementarlos en dispositivos con recursos limitados. Este artículo investiga los módulos que representan cuellos de botella computacional en los transformers eficientes, es decir, las capas de normalización y los módulos de atención. LayerNorm se utiliza comúnmente en las arquitecturas de transformers, pero no es computacionalmente eficiente debido al cálculo de estadísticas durante la inferencia. Sin embargo, reemplazar LayerNorm con BatchNorm, que es más eficiente, en los transformers a menudo conduce a un rendimiento inferior y a colapsos durante el entrenamiento. Para abordar este problema, proponemos un método novedoso llamado PRepBN para reemplazar progresivamente LayerNorm con BatchNorm re-parametrizado durante el entrenamiento. Además, proponemos un módulo de atención lineal simplificado (SLA) que es simple pero efectivo para lograr un rendimiento sólido. Experimentos extensos en clasificación de imágenes y detección de objetos demuestran la efectividad de nuestro método propuesto. Por ejemplo, nuestro SLAB-Swin obtiene un 83.6% de precisión top-1 en ImageNet-1K con una latencia de 16.2ms, que es 2.4ms menos que Flatten-Swin con un 0.1% más de precisión. También evaluamos nuestro método para la tarea de modelado de lenguaje y obtenemos un rendimiento comparable con una latencia más baja. Los códigos están disponibles públicamente en https://github.com/xinghaochen/SLAB y https://github.com/mindspore-lab/models/tree/master/research/huawei-noah/SLAB.

English

Transformers have become foundational architectures for both natural language and computer vision tasks. However, the high computational cost makes it quite challenging to deploy on resource-constraint devices. This paper investigates the computational bottleneck modules of efficient transformer, i.e., normalization layers and attention modules. LayerNorm is commonly used in transformer architectures but is not computational friendly due to statistic calculation during inference. However, replacing LayerNorm with more efficient BatchNorm in transformer often leads to inferior performance and collapse in training. To address this problem, we propose a novel method named PRepBN to progressively replace LayerNorm with re-parameterized BatchNorm in training. Moreover, we propose a simplified linear attention (SLA) module that is simple yet effective to achieve strong performance. Extensive experiments on image classification as well as object detection demonstrate the effectiveness of our proposed method. For example, our SLAB-Swin obtains 83.6% top-1 accuracy on ImageNet-1K with 16.2ms latency, which is 2.4ms less than that of Flatten-Swin with 0.1% higher accuracy. We also evaluated our method for language modeling task and obtain comparable performance and lower latency.Codes are publicly available at https://github.com/xinghaochen/SLAB and https://github.com/mindspore-lab/models/tree/master/research/huawei-noah/SLAB.