Le Transformer Structuré : Modèles d'Attention dans la Limite de Profondeur et de Largeur Infinies

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

Résumé

Dans la théorie de l'apprentissage profond, la matrice de covariance des représentations sert de proxy pour examiner la capacité d'entraînement du réseau. Motivés par le succès des Transformers, nous étudions la matrice de covariance d'un modèle d'attention modifié basé sur Softmax avec des connexions résiduelles dans la limite proportionnelle de profondeur et de largeur infinies. Nous montrons qu'à l'initialisation, la distribution limite peut être décrite par une équation différentielle stochastique (EDS) indexée par le rapport profondeur-sur-largeur. Pour atteindre une limite stochastique bien définie, le mécanisme d'attention du Transformer est modifié en centrant la sortie Softmax sur l'identité et en mettant à l'échelle les logits Softmax par un paramètre de température dépendant de la largeur. Nous examinons la stabilité du réseau à travers l'EDS correspondante, montrant comment l'échelle à la fois de la dérive et de la diffusion peut être élégamment contrôlée à l'aide des connexions résiduelles. L'existence d'une EDS stable implique que la structure de covariance est bien comportée, même pour des profondeurs et largeurs très importantes, évitant ainsi les problèmes notoires de dégénérescence de rang dans les modèles d'attention profonds. Enfin, nous montrons, à travers des simulations, que l'EDS fournit une description étonnamment précise du modèle de taille finie correspondant. Nous baptisons ces modifications architecturales le Transformer façonné.

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.