Effective Theory of Transformers at Initialization

TL;DR

This paper analyzes how signals propagate in wide, deep Transformers at initialization, providing guidance on optimal hyperparameter scaling, and validates these insights through practical training of Vision and Language Transformers.

Contribution

It offers an effective-theory framework for understanding Transformer initialization and proposes specific hyperparameter scalings validated by experiments.

Findings

Guides on width scalings for initialization and training

Improved training stability and performance in Vision and Language Transformers

Theoretical insights align with empirical results

Abstract

We perform an effective-theory analysis of forward-backward signal propagation in wide and deep Transformers, i.e., residual neural networks with multi-head self-attention blocks and multilayer perceptron blocks. This analysis suggests particular width scalings of initialization and training hyperparameters for these models. We then take up such suggestions, training Vision and Language Transformers in practical setups.