SLAB: Эффективные трансформеры с упрощенным линейным вниманием и прогрессивной переопределенной пакетной нормализацией

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

May 19, 2024

Авторы: Jialong Guo, Xinghao Chen, Yehui Tang, Yunhe Wang

cs.AI

Аннотация

Трансформеры стали основополагающими архитектурами как для задач обработки естественного языка, так и для задач компьютерного зрения. Однако высокая вычислительная стоимость делает их довольно сложными для развертывания на устройствах с ограниченными ресурсами. В данной статье исследуются вычислительные узкие места эффективного трансформера, а именно слои нормализации и модули внимания. LayerNorm обычно используется в архитектурах трансформеров, но не является вычислительно эффективным из-за статистических вычислений во время вывода. Однако замена LayerNorm более эффективным BatchNorm в трансформере часто приводит к ухудшению производительности и краху обучения. Для решения этой проблемы мы предлагаем новый метод под названием PRepBN для постепенной замены LayerNorm на перепараметризованный BatchNorm во время обучения. Более того, мы предлагаем упрощенный линейный модуль внимания (SLA), который прост в реализации, но эффективен для достижения высокой производительности. Обширные эксперименты по классификации изображений, а также по обнаружению объектов, демонстрируют эффективность нашего предложенного метода. Например, наш SLAB-Swin достигает точности на уровне 83.6% top-1 на ImageNet-1K с задержкой 16.2 мс, что на 2.4 мс меньше, чем у Flatten-Swin, при точности на 0.1% выше. Мы также оценили наш метод для задачи языкового моделирования и получили сопоставимую производительность и более низкую задержку. Коды доступны по ссылкам https://github.com/xinghaochen/SLAB и 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.