SLAB: Effiziente Transformer mit vereinfachter linearer Aufmerksamkeit und progressiver reparametrisierter Batch-Normalisierung.

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

May 19, 2024

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

cs.AI

Zusammenfassung

Transformer-Architekturen sind zu grundlegenden Strukturen für sowohl natürliche Sprache als auch Computer Vision Aufgaben geworden. Allerdings macht die hohe Rechenkosten es recht herausfordernd, sie auf ressourcenbeschränkten Geräten einzusetzen. Diese Arbeit untersucht die rechenintensiven Module des effizienten Transformers, d.h., Normalisierungsschichten und Aufmerksamkeitsmodule. LayerNorm wird häufig in Transformer-Architekturen verwendet, ist jedoch aufgrund der statistischen Berechnung während der Inferenz nicht rechenfreundlich. Der Ersatz von LayerNorm durch die effizientere BatchNorm im Transformer führt jedoch oft zu schlechterer Leistung und einem Zusammenbruch im Training. Um dieses Problem zu lösen, schlagen wir eine neuartige Methode namens PRepBN vor, um LayerNorm schrittweise durch reparametrisierte BatchNorm im Training zu ersetzen. Darüber hinaus schlagen wir ein vereinfachtes lineares Aufmerksamkeitsmodul (SLA) vor, das einfach, aber effektiv ist, um starke Leistung zu erzielen. Umfangreiche Experimente zur Bildklassifizierung sowie Objekterkennung zeigen die Wirksamkeit unserer vorgeschlagenen Methode. Zum Beispiel erreicht unser SLAB-Swin eine Top-1-Genauigkeit von 83,6% auf ImageNet-1K mit einer Latenz von 16,2 ms, was 2,4 ms weniger ist als bei Flatten-Swin mit 0,1% höherer Genauigkeit. Wir haben unsere Methode auch für die Sprachmodellierungsaufgabe evaluiert und vergleichbare Leistung und niedrigere Latenz erzielt. Die Codes sind öffentlich verfügbar unter https://github.com/xinghaochen/SLAB und 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.