El Transformer con Forma: Modelos de Atención en el Límite de Profundidad y Ancho Infinitos

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

June 30, 2023

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

cs.AI

Resumen

En la teoría del aprendizaje profundo, la matriz de covarianza de las representaciones sirve como un indicador para examinar la capacidad de entrenamiento de la red. Motivados por el éxito de los Transformers, estudiamos la matriz de covarianza de un modelo de atención basado en Softmax modificado con conexiones residuales en el límite proporcional de profundidad y anchura infinitas. Demostramos que, en la inicialización, la distribución límite puede describirse mediante una ecuación diferencial estocástica (SDE) indexada por la relación profundidad-anchura. Para lograr un límite estocástico bien definido, el mecanismo de atención del Transformer se modifica centrando la salida del Softmax en la identidad y escalando los logits del Softmax con un parámetro de temperatura dependiente de la anchura. Examinamos la estabilidad de la red a través de la SDE correspondiente, mostrando cómo la escala tanto de la deriva como de la difusión puede controlarse elegantemente con la ayuda de las conexiones residuales. La existencia de una SDE estable implica que la estructura de covarianza está bien comportada, incluso para profundidades y anchuras muy grandes, evitando así los notorios problemas de degeneración de rango en los modelos de atención profundos. Finalmente, mostramos, mediante simulaciones, que la SDE proporciona una descripción sorprendentemente buena del modelo de tamaño finito correspondiente. Acuñamos el nombre Transformer conformado para estas modificaciones arquitectónicas.

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.