SLAB: Transformers Eficientes com Atenção Linear Simplificada e Normalização de Lotes com Reparametrização Progressiva

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

Resumo

Os Transformers tornaram-se arquiteturas fundamentais tanto para tarefas de processamento de linguagem natural quanto para visão computacional. No entanto, o alto custo computacional torna bastante desafiador sua implantação em dispositivos com recursos limitados. Este artigo investiga os módulos que representam gargalos computacionais em transformers eficientes, ou seja, camadas de normalização e módulos de atenção. O LayerNorm é comumente usado em arquiteturas de transformers, mas não é computacionalmente amigável devido ao cálculo de estatísticas durante a inferência. No entanto, substituir o LayerNorm pelo BatchNorm mais eficiente em transformers frequentemente resulta em desempenho inferior e colapso durante o treinamento. Para resolver esse problema, propomos um método inovador chamado PRepBN para substituir progressivamente o LayerNorm pelo BatchNorm reparametrizado durante o treinamento. Além disso, propomos um módulo de atenção linear simplificado (SLA) que é simples, mas eficaz para alcançar um forte desempenho. Experimentos extensivos em classificação de imagens e detecção de objetos demonstram a eficácia do nosso método proposto. Por exemplo, nosso SLAB-Swin obtém 83,6% de precisão top-1 no ImageNet-1K com uma latência de 16,2ms, que é 2,4ms menor que a do Flatten-Swin com 0,1% a mais de precisão. Também avaliamos nosso método para a tarefa de modelagem de linguagem e obtivemos desempenho comparável e menor latência. Os códigos estão publicamente disponíveis em https://github.com/xinghaochen/SLAB e 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.