SLAB: Efficiënte Transformers met Vereenvoudigde Lineaire Aandacht en Progressieve Herparameterisatie van Batch Normalisatie

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

Samenvatting

Transformers zijn fundamentele architecturen geworden voor zowel natuurlijke taalverwerking als computervisietaken. De hoge rekenkosten maken het echter behoorlijk uitdagend om ze te implementeren op apparaten met beperkte middelen. Dit artikel onderzoekt de rekenkundige knelpunten van efficiënte transformers, namelijk normalisatielagen en aandachtmodules. LayerNorm wordt vaak gebruikt in transformer-architecturen, maar is niet rekenkundig vriendelijk vanwege de statistische berekeningen tijdens de inferentie. Het vervangen van LayerNorm door de efficiëntere BatchNorm in transformers leidt echter vaak tot inferieure prestaties en instorting tijdens de training. Om dit probleem aan te pakken, stellen we een nieuwe methode voor genaamd PRepBN om LayerNorm geleidelijk te vervangen door her-geparameteriseerde BatchNorm tijdens de training. Bovendien stellen we een vereenvoudigde lineaire aandachtmodule (SLA) voor die eenvoudig maar effectief is om sterke prestaties te behalen. Uitgebreide experimenten op het gebied van beeldclassificatie en objectdetectie tonen de effectiviteit van onze voorgestelde methode aan. Onze SLAB-Swin behaalt bijvoorbeeld een top-1 nauwkeurigheid van 83,6% op ImageNet-1K met een latentie van 16,2ms, wat 2,4ms minder is dan die van Flatten-Swin met een 0,1% hogere nauwkeurigheid. We hebben onze methode ook geëvalueerd voor taalmodelleertaken en vergelijkbare prestaties en lagere latentie behaald. De codes zijn openbaar beschikbaar op https://github.com/xinghaochen/SLAB en 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.