Unified Normalization for Accelerating and Stabilizing Transformers

TL;DR

This paper introduces Unified Normalization (UN), a hardware-efficient normalization method for Transformers that accelerates inference, stabilizes training, and maintains performance comparable to Layer Normalization, with significant speed and memory benefits.

Contribution

The paper proposes a novel Unified Normalization technique that addresses the inefficiencies and performance issues of existing normalization methods in Transformers.

Findings

UN achieves about 31% inference speedup on GPU.

UN reduces memory usage by nearly 18%.

UN maintains comparable performance to Layer Normalization.

Abstract

Solid results from Transformers have made them prevailing architectures in various natural language and vision tasks. As a default component in Transformers, Layer Normalization (LN) normalizes activations within each token to boost the robustness. However, LN requires on-the-fly statistics calculation in inference as well as division and square root operations, leading to inefficiency on hardware. What is more, replacing LN with other hardware-efficient normalization schemes (e.g., Batch Normalization) results in inferior performance, even collapse in training. We find that this dilemma is caused by abnormal behaviors of activation statistics, including large fluctuations over iterations and extreme outliers across layers. To tackle these issues, we propose Unified Normalization (UN), which can speed up the inference by being fused with other linear operations and achieve comparable performance on par with LN. UN strives to boost performance by calibrating the activation and gradient statistics with a tailored fluctuation smoothing strategy. Meanwhile, an adaptive outlier filtration strategy is applied to avoid collapse in training whose effectiveness is theoretically proved and experimentally verified in this paper. We demonstrate that UN can be an efficient drop-in alternative to LN by conducting extensive experiments on language and vision tasks. Besides, we evaluate the efficiency of our method on GPU. Transformers equipped with UN enjoy about 31% inference speedup and nearly 18% memory reduction. Code will be released at https://github.com/hikvision-research/Unified-Normalization.