SLAB : Transformers efficaces avec attention linéaire simplifiée et normalisation par lots re-paramétrée progressive

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

May 19, 2024

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

cs.AI

Résumé

Les Transformers sont devenus des architectures fondamentales pour les tâches de traitement du langage naturel et de vision par ordinateur. Cependant, leur coût computationnel élevé rend leur déploiement sur des appareils à ressources limitées particulièrement difficile. Cet article étudie les modules qui constituent des goulots d'étranglement computationnels dans les Transformers efficaces, à savoir les couches de normalisation et les modules d'attention. Le LayerNorm est couramment utilisé dans les architectures de Transformers, mais il n'est pas favorable sur le plan computationnel en raison du calcul des statistiques pendant l'inférence. Cependant, remplacer le LayerNorm par le BatchNorm, plus efficace, dans les Transformers conduit souvent à une performance inférieure et à un effondrement pendant l'entraînement. Pour résoudre ce problème, nous proposons une nouvelle méthode appelée PRepBN pour remplacer progressivement le LayerNorm par un BatchNorm reparamétré pendant l'entraînement. De plus, nous proposons un module d'attention linéaire simplifié (SLA) qui est simple mais efficace pour atteindre de solides performances. Des expériences approfondies sur la classification d'images ainsi que sur la détection d'objets démontrent l'efficacité de notre méthode proposée. Par exemple, notre SLAB-Swin obtient une précision top-1 de 83,6 % sur ImageNet-1K avec une latence de 16,2 ms, ce qui est 2,4 ms de moins que celle de Flatten-Swin avec une précision supérieure de 0,1 %. Nous avons également évalué notre méthode pour la tâche de modélisation du langage et obtenu des performances comparables avec une latence réduite. Les codes sont disponibles publiquement à l'adresse https://github.com/xinghaochen/SLAB et 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.