De Shaped Transformer: Aandachtsmodellen in de Limiet van Oneindige Diepte en Breedte

The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit

June 30, 2023

Auteurs: Lorenzo Noci, Chuning Li, Mufan Bill Li, Bobby He, Thomas Hofmann, Chris Maddison, Daniel M. Roy

cs.AI

Samenvatting

In de diepgaande leerstheorie dient de covariantiematrix van de representaties als een benadering om de trainbaarheid van het netwerk te onderzoeken. Geïnspireerd door het succes van Transformers, bestuderen we de covariantiematrix van een aangepast Softmax-gebaseerd aandachtmodel met skip-verbindingen in de proportionele limiet van oneindige diepte en breedte. We tonen aan dat bij initialisatie de limietverdeling kan worden beschreven door een stochastische differentiaalvergelijking (SDE) geïndexeerd op de diepte-tot-breedte-verhouding. Om een goed gedefinieerde stochastische limiet te bereiken, wordt het aandachtmechanisme van de Transformer aangepast door de Softmax-uitvoer te centreren op de identiteit en de Softmax-logits te schalen met een breedteafhankelijke temperatuurparameter. We onderzoeken de stabiliteit van het netwerk via de bijbehorende SDE, waarbij we aantonen hoe zowel de schaal van de drift als de diffusie elegant kan worden geregeld met behulp van restverbindingen. Het bestaan van een stabiele SDE impliceert dat de covariantiestructuur goed gedrag vertoont, zelfs voor zeer grote diepte en breedte, waardoor de beruchte problemen van rangdegeneratie in diepe aandachtmodellen worden voorkomen. Tot slot tonen we, via simulaties, aan dat de SDE een verrassend goede beschrijving geeft van het corresponderende eindige-grootte-model. We introduceren de naam 'shaped Transformer' voor deze architecturale aanpassingen.

English

In deep learning theory, the covariance matrix of the representations serves as a proxy to examine the network's trainability. Motivated by the success of Transformers, we study the covariance matrix of a modified Softmax-based attention model with skip connections in the proportional limit of infinite-depth-and-width. We show that at initialization the limiting distribution can be described by a stochastic differential equation (SDE) indexed by the depth-to-width ratio. To achieve a well-defined stochastic limit, the Transformer's attention mechanism is modified by centering the Softmax output at identity, and scaling the Softmax logits by a width-dependent temperature parameter. We examine the stability of the network through the corresponding SDE, showing how the scale of both the drift and diffusion can be elegantly controlled with the aid of residual connections. The existence of a stable SDE implies that the covariance structure is well-behaved, even for very large depth and width, thus preventing the notorious issues of rank degeneracy in deep attention models. Finally, we show, through simulations, that the SDE provides a surprisingly good description of the corresponding finite-size model. We coin the name shaped Transformer for these architectural modifications.