O Transformer com Formato: Modelos de Atenção no Limite Infinito de Profundidade e Largura

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

Resumo

Na teoria de aprendizado profundo, a matriz de covariância das representações serve como um proxy para examinar a treinabilidade da rede. Motivados pelo sucesso dos Transformers, estudamos a matriz de covariância de um modelo de atenção baseado em Softmax modificado com conexões de salto (skip connections) no limite proporcional de profundidade e largura infinitas. Mostramos que, na inicialização, a distribuição limite pode ser descrita por uma equação diferencial estocástica (SDE) indexada pela razão profundidade-largura. Para alcançar um limite estocástico bem definido, o mecanismo de atenção do Transformer é modificado centralizando a saída do Softmax na identidade e escalonando os logits do Softmax por um parâmetro de temperatura dependente da largura. Examinamos a estabilidade da rede por meio da SDE correspondente, mostrando como a escala tanto do drift quanto da difusão pode ser elegantemente controlada com o auxílio de conexões residuais. A existência de uma SDE estável implica que a estrutura de covariância é bem comportada, mesmo para profundidade e largura muito grandes, prevenindo assim os notórios problemas de degenerescência de posto em modelos de atenção profundos. Por fim, mostramos, por meio de simulações, que a SDE fornece uma descrição surpreendentemente boa do modelo de tamanho finito correspondente. Batizamos essas modificações arquiteturais com o nome de Transformer moldado (shaped Transformer).

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.