Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs

TL;DR

This paper rigorously derives the non-Gaussian infinite-width limit distribution of a single attention layer in neural networks using Tensor Programs, providing insights beyond traditional Gaussian approximations and aiding the understanding of Transformer models.

Contribution

It introduces a novel theoretical framework to analyze the infinite-width limit of attention layers without simplifying assumptions like infinite heads or special scalings.

Findings

The limit distribution is non-Gaussian and hierarchical.

Numerical experiments confirm the theory's accuracy at finite widths.

The results enable a unified understanding of deep Transformer architectures.

Abstract

In modern theoretical analyses of neural networks, the infinite-width limit is often invoked to justify Gaussian approximations of neuron preactivations (e.g., via neural network Gaussian processes or Tensor Programs). However, these Gaussian-based asymptotic theories have so far been unable to capture the behavior of attention layers, except under special regimes such as infinitely many heads or tailored scaling schemes. In this paper, leveraging the Tensor Programs framework, we rigorously identify the infinite-width limit distribution of variables within a single attention layer under realistic architectural dimensionality and standard $1/\sqrt{n}$-scaling with $n$ dimensionality. We derive the exact form of this limit law without resorting to infinite-head approximations or tailored scalings, demonstrating that it departs fundamentally from Gaussianity. This limiting distribution exhibits non-Gaussianity from a hierarchical structure, being Gaussian conditional on the random similarity scores. Numerical experiments validate our theoretical predictions, confirming the effectiveness of our theory at finite width and accurate description of finite-head attentions. Beyond characterizing a standalone attention layer, our findings lay the groundwork for developing a unified theory of deep Transformer architectures in the infinite-width regime.