Il Transformer Modellato: Modelli di Attenzione nel Limite di Profondità e Ampiezza Infinita

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

June 30, 2023

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

cs.AI

Abstract

Nella teoria del deep learning, la matrice di covarianza delle rappresentazioni funge da proxy per esaminare l'addestrabilità della rete. Ispirati dal successo dei Transformer, studiamo la matrice di covarianza di un modello di attenzione basato su Softmax modificato con connessioni skip nel limite proporzionale di profondità e larghezza infinite. Dimostriamo che all'inizializzazione la distribuzione limite può essere descritta da un'equazione differenziale stocastica (SDE) indicizzata dal rapporto profondità-larghezza. Per ottenere un limite stocastico ben definito, il meccanismo di attenzione del Transformer viene modificato centrando l'output del Softmax sull'identità e scalando i logit del Softmax con un parametro di temperatura dipendente dalla larghezza. Esaminiamo la stabilità della rete attraverso la corrispondente SDE, mostrando come la scala sia della deriva che della diffusione possa essere elegantemente controllata con l'aiuto delle connessioni residue. L'esistenza di una SDE stabile implica che la struttura di covarianza sia ben comportata, anche per profondità e larghezza molto grandi, prevenendo così i noti problemi di degenerazione del rango nei modelli di attenzione profondi. Infine, mostriamo, attraverso simulazioni, che la SDE fornisce una descrizione sorprendentemente accurata del corrispondente modello a dimensione finita. Coniamo il nome shaped Transformer per queste modifiche architetturali.

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.