Learning in Compact Spaces with Approximately Normalized Transformer

TL;DR

This paper introduces an approximate normalization technique for transformers that constrains parameters on a hypersphere, leading to faster convergence and better scaling without extra hyperparameters.

Contribution

It proposes a holistic approximate normalization method that removes the need for regularization and hyperparameter tuning in transformer training.

Findings

Up to 40% faster convergence compared to GPT with QK normalization.

Enables training with larger batch sizes while maintaining scaling laws.

Minimal 3% additional runtime cost for the normalization method.

Abstract

The successful training of deep neural networks requires addressing challenges such as overfitting, numerical instabilities leading to divergence, and increasing variance in the residual stream. A common solution is to apply regularization and normalization techniques that usually require tuning additional hyperparameters. An alternative is to force all parameters and representations to lie on a hypersphere. This removes the need for regularization and increases convergence speed, but comes with additional costs. In this work, we propose a more holistic, approximate normalization via simple scalar multiplications motivated by the tight concentration of the norms of high-dimensional random vectors. Additionally, instead of applying strict normalization for the parameters, we constrain their norms. These modifications remove the need for weight decay and learning rate warm-up as well, but do not increase the total number of normalization layers. Our experiments with transformer architectures show up to 40% faster convergence compared to GPT models with QK normalization, with only 3% additional runtime cost. When deriving scaling laws, we found that our method enables training with larger batch sizes while preserving the favorable scaling characteristics of classic GPT architectures.