Papers
Attention Is All You Need
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin · NeurIPS 2017 (31st Conference on Neural Information Processing Systems, então chamada NIPS), 2017
O problema: redes recorrentes são sequenciais por construção
Na época do paper, os modelos dominantes para transdução de sequências (como tradução automática) eram redes recorrentes ou convolucionais com estrutura de codificador e decodificador, e os melhores ainda ligavam codificador e decodificador por um mecanismo de atenção. As redes recorrentes, incluindo LSTM e GRU, processam a sequência ao longo das posições dos símbolos: o estado oculto da posição t é calculado a partir do estado da posição anterior e da entrada em t. Segundo os autores, essa natureza inerentemente sequencial impede a paralelização dentro de um mesmo exemplo de treino, o que pesa mais em sequências longas, porque a memória limita quantos exemplos podem ser agrupados em um lote. Truques de fatoração e computação condicional melhoraram a eficiência, mas a restrição fundamental da computação sequencial permanecia. Modelos convolucionais (Extended Neural GPU, ByteNet, ConvS2S) reduzem a sequencialidade ao calcular representações de todas as posições em paralelo, mas o número de operações para relacionar duas posições distantes cresce com a distância: linearmente no ConvS2S e logaritmicamente no ByteNet, o que dificulta aprender dependências de longo alcance. A atenção, por sua vez, já permitia modelar dependências sem considerar a distância entre posições, mas quase sempre era usada junto de uma rede recorrente. A proposta do paper é o Transformer, uma arquitetura que dispensa recorrência e convoluções e se apoia só em atenção para relacionar entrada e saída.
Self-attention: consulta, chaves e valores
O Transformer mantém a estrutura de codificador e decodificador, mas cada uma das duas pilhas é feita de 6 camadas idênticas. Cada camada do codificador tem duas subcamadas: atenção própria (self-attention) com várias cabeças e uma rede feed-forward aplicada a cada posição de forma separada e idêntica. Em volta de cada subcamada há uma conexão residual seguida de normalização de camada, e todas as saídas têm dimensão 512. O decodificador acrescenta uma terceira subcamada, que atende à saída do codificador, e mascara a atenção própria para que a previsão da posição i dependa só das saídas conhecidas em posições anteriores a i, preservando a geração auto-regressiva. A atenção é descrita como um mapeamento de uma consulta (query) e de um conjunto de pares chave-valor para uma saída: uma soma ponderada dos valores, cujo peso vem de uma função de compatibilidade entre a consulta e cada chave. O paper usa a atenção por produto escalar escalado: calcula-se o produto da consulta com todas as chaves, divide-se pela raiz da dimensão das chaves e aplica-se softmax para obter os pesos. Os autores suspeitam que, para dimensões grandes, os produtos escalares crescem em magnitude e empurram o softmax para regiões de gradientes muito pequenos, e a divisão serve para compensar isso. Em self-attention, consultas, chaves e valores vêm do mesmo lugar, a saída da camada anterior, e cada posição pode atender a todas as posições dessa camada.
Multi-head attention e a informação de posição
Em vez de uma única função de atenção sobre vetores de dimensão cheia, o paper projeta consultas, chaves e valores h vezes com projeções lineares aprendidas e diferentes, aplica a atenção em paralelo a cada versão projetada, concatena os resultados e projeta de novo. Os autores afirmam que isso permite ao modelo atender, ao mesmo tempo, a informações de diferentes subespaços de representação em diferentes posições, algo que uma cabeça única inibiria por fazer uma média. Na configuração base são h = 8 cabeças com dimensão 64 cada, de modo que o custo total fica parecido com o de uma única cabeça de dimensão cheia. Nos experimentos de variação do modelo, uma única cabeça ficou 0,9 BLEU pior que o melhor ajuste, e a qualidade também cai com cabeças demais. Como não há recorrência nem convolução, o modelo não enxerga a ordem por si só, então o paper soma aos embeddings de entrada codificações posicionais de mesma dimensão. Elas usam senos e cossenos de frequências diferentes, com comprimentos de onda em progressão geométrica. A hipótese dos autores é que isso facilita aprender a atender por posições relativas, pois, para um deslocamento fixo, a codificação em pos + k é uma função linear da codificação em pos. Também testaram embeddings posicionais aprendidos e obtiveram resultados quase idênticos; preferiram a versão senoidal porque ela talvez permita extrapolar para sequências mais longas que as vistas no treino.
Por que self-attention: a tabela de custo por camada
A seção 'Why Self-Attention' compara camadas de self-attention, recorrentes e convolucionais por três critérios: complexidade computacional por camada, quantidade de computação paralelizável (medida pelo número mínimo de operações sequenciais) e comprimento do caminho máximo entre duas posições, já que caminhos curtos facilitam aprender dependências de longo alcance. Na Tabela 1, com n o comprimento da sequência, d a dimensão da representação, k o tamanho do kernel e r o tamanho da vizinhança na atenção restrita, os valores são estes. Self-attention: complexidade O(n² · d), operações sequenciais O(1), caminho máximo O(1). Recorrente: O(n · d²), O(n) e O(n). Convolucional: O(k · n · d²), O(1) e O(log_k n). Self-attention restrita a uma vizinhança r: O(r · n · d), O(1) e O(n/r). Os autores observam que a self-attention é mais rápida que a camada recorrente quando n é menor que d, o que ocorre com frequência nas representações de sentenças usadas em tradução, como word-piece e byte-pair. Para sequências muito longas, propõem como trabalho futuro restringir a atenção a uma vizinhança, ao custo de aumentar o caminho máximo. O trade-off é claro: a atenção completa conecta todas as posições com um número constante de operações sequenciais, mas o termo n² faz o custo crescer rápido com o comprimento da sequência. Como benefício colateral, os autores apontam modelos potencialmente mais interpretáveis, pois várias cabeças parecem aprender tarefas distintas.
Resultados em tradução e legado declarado pelos autores
Os modelos foram treinados em uma máquina com 8 GPUs NVIDIA P100. O modelo base levou cerca de 12 horas (100 mil passos) e o modelo grande 3,5 dias (300 mil passos). Na tradução inglês-alemão do WMT 2014, o modelo grande alcançou 28,4 BLEU, mais de 2 BLEU acima dos melhores resultados anteriores, incluindo ensembles; até o modelo base superou todos os modelos e ensembles publicados antes, com custo de treino bem menor. Na tradução inglês-francês, o resumo do paper reporta 41,8 BLEU como novo estado da arte para modelo único, após 3,5 dias de treino em oito GPUs, uma fração do custo dos melhores modelos da literatura (o corpo do texto cita 41,0 para esse mesmo modelo, uma divergência interna do paper). Os autores estimam o custo em FLOPs multiplicando tempo de treino, número de GPUs e uma estimativa da capacidade sustentada de cada GPU. Nas variações da Tabela 3, modelos maiores foram melhores e dropout ajudou contra overfitting. O Transformer também foi aplicado a análise sintática de constituintes em inglês, onde teve bom desempenho mesmo sem ajuste específico para a tarefa. Na conclusão, os autores se descrevem como autores do primeiro modelo de transdução baseado inteiramente em atenção e planejam estendê-lo a outras modalidades além de texto, investigar atenção local e restrita para entradas grandes e tornar a geração menos sequencial. O paper não descreve usos posteriores; qualquer legado além disso vem de fontes externas a ele.
Veja as ideias em ação
Pratique na escola de System Design os desenhos que usam este paper.