Der Shaped Transformer: Aufmerksamkeitsmodelle im Grenzbereich unendlicher Tiefe und Breite

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

June 30, 2023

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

cs.AI

Zusammenfassung

In der Theorie des Deep Learnings dient die Kovarianzmatrix der Repräsentationen als Stellvertreter, um die Trainierbarkeit des Netzwerks zu untersuchen. Motiviert durch den Erfolg von Transformern, untersuchen wir die Kovarianzmatrix eines modifizierten Softmax-basierten Aufmerksamkeitsmodells mit Skip-Connections im proportionalen Grenzfall unendlicher Tiefe und Breite. Wir zeigen, dass die Grenzverteilung bei der Initialisierung durch eine stochastische Differentialgleichung (SDE) beschrieben werden kann, die durch das Verhältnis von Tiefe zu Breite indiziert wird. Um einen wohldefinierten stochastischen Grenzwert zu erreichen, wird der Aufmerksamkeitsmechanismus des Transformers modifiziert, indem die Softmax-Ausgabe auf die Identität zentriert und die Softmax-Logits durch einen breitenabhängigen Temperaturparameter skaliert werden. Wir untersuchen die Stabilität des Netzwerks durch die entsprechende SDE und zeigen, wie die Skala sowohl der Drift als auch der Diffusion elegant mit Hilfe von Residualverbindungen gesteuert werden kann. Die Existenz einer stabilen SDE impliziert, dass die Kovarianzstruktur auch bei sehr großer Tiefe und Breite gutartig ist, wodurch die bekannten Probleme der Rangentartung in tiefen Aufmerksamkeitsmodellen verhindert werden. Schließlich zeigen wir durch Simulationen, dass die SDE eine überraschend gute Beschreibung des entsprechenden Modells mit endlicher Größe liefert. Wir prägen den Namen Shaped Transformer für diese architektonischen Modifikationen.

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.