Visualização de atenção: o coração do Transformer [Vídeo]

• A atenção no Transformer atualiza embeddings de tokens de acordo com o contexto, deslocando até a mesma palavra para vetores com significados diferentes dependendo das palavras ao redor
• Um head de atenção cria vetores de query/key/value para cada token e calcula o attention pattern, que são os pesos de relevância entre palavras, por meio do produto interno entre key e query e do softmax
• Modelos autorregressivos no estilo GPT aplicam masking para impedir que tokens posteriores influenciem tokens anteriores, e como o tamanho do attention pattern cresce com o quadrado do comprimento do contexto, expandir muito a context window se torna custoso
• No exemplo do GPT-3, as matrizes de key/query têm 1.572.864 parâmetros cada, e o value map é dividido em transformações de baixo posto, resultando em cerca de 6,3 milhões de parâmetros por head
• O Transformer, que repete vários heads e blocos de atenção, aprende diferentes formas de atualizar o contexto, e uma grande parte do seu sucesso vem da capacidade de paralelização para processar muitos cálculos rapidamente em GPUs

---

O papel da atenção no Transformer

• O Transformer usa o texto de entrada para prever o próximo token, e a entrada primeiro passa por tokenização, sendo dividida em palavras ou pedaços de palavras
• Cada token é convertido em um embedding, um vetor de alta dimensão
  • As direções nesse espaço de embeddings podem corresponder a significado
  • Por exemplo, mover-se em uma direção específica pode levar o embedding de um substantivo masculino ao embedding do substantivo feminino correspondente
• O objetivo da atenção é ajustar gradualmente os embeddings iniciais para que passem a conter não só a informação de cada palavra, mas também um significado contextual mais rico

Por que a mesma palavra muda conforme o contexto

• Em “American shrew mole”, “One mole of carbon dioxide” e “Take a biopsy of the mole”, mole tem significados diferentes
• Na etapa inicial de embedding, o vetor de mole é praticamente o mesmo nos três casos, pois funciona mais como uma lookup table que não vê o contexto
• Na etapa seguinte, o attention block, embeddings vizinhos podem transmitir informação ao embedding de mole e atualizar seu valor
• Um modelo bem treinado conecta os vários significados de mole a direções diferentes no espaço de embeddings e calcula, com base no contexto, o que deve ser somado ao embedding genérico
• Como em “Eiffel tower” e “miniature Eiffel tower”, o embedding de uma palavra pode ser atualizado não só por palavras próximas, mas também por informação vinda de tokens distantes
• Como a previsão da próxima palavra usa apenas o vetor final, o embedding da última palavra de uma entrada longa precisa conter, até certo ponto, toda a informação de contexto necessária para a previsão

Fluxo de cálculo de um único head de atenção

• A explicação básica segue um single head of attention
• Na frase de exemplo “A fluffy blue creature roamed the verdant forest.”, assume-se uma situação em que adjetivos atualizam os embeddings iniciais dos substantivos correspondentes
  • Esse exemplo serve para mostrar o tipo de operação que um head de atenção pode realizar
  • O comportamento real de um head é difícil de interpretar, porque resulta de muitos parâmetros ajustados para reduzir a função de custo
• Os embeddings iniciais incluem não só informação da palavra, mas também informação de posição, e são representados por \vec{E}
• O objetivo é produzir um novo embedding \vec{E}', que reflita o contexto, a partir do embedding existente

• Query

  • Na primeira etapa, multiplica-se cada embedding de token pela query matrix W_Q para criar o vetor query \vec{Q}
  • Dá para pensar nisso como se o substantivo estivesse fazendo a pergunta “há um adjetivo antes de mim?”
  • Os elementos de W_Q são parâmetros treináveis do modelo, e é difícil interpretar exatamente o que um head específico faz
  • Como exemplo intuitivo, dá para imaginar que ele mapeia embeddings de substantivos para uma direção de “buscar adjetivos em posições anteriores”

• Key

  • Ao mesmo tempo, multiplica-se cada embedding pela key matrix W_k para criar o vetor key \vec{K}
  • A key pode ser vista como uma resposta potencial à query, e fica em um espaço de menor dimensão, assim como a query
  • O grau de alinhamento entre key e query é medido pelo produto interno
  • Quanto maior o produto interno, mais fortemente alinhados estão os dois vetores
  • Se as keys de fluffy e blue combinarem bem com a query de creature, terão valores positivos altos
  • Ao calcular o produto interno para todos os pares key-query, surge uma grade de pontuações que indica o quanto cada palavra é relevante para atualizar o significado de outra palavra

Attention pattern e softmax

• Como as pontuações de produto interno podem variar de -\infty a \infty, aplica-se softmax a cada coluna para normalizá-las em valores entre 0 e 1
• A grade normalizada é chamada de attention pattern
  • Cada coluna pode ser vista como pesos de quão relevante é a palavra à esquerda para atualizar a palavra no topo
• O artigo original do Transformer escreve isso de forma mais compacta
  • Q e K são os arranjos completos de vetores query e key
  • K^TQ representa a grade de todos os produtos internos possíveis entre key e query
  • Na notação do artigo, query e key ficam nas linhas e aparece a forma QK^T, que fica invertida na diagonal em relação ao diagrama explicado aqui
• Para estabilidade numérica, entra um termo de divisão pela raiz quadrada da dimensão do espaço key-query, \sqrt{d_k}
• O softmax envolve a expressão inteira, mas em termos de significado é aplicado a cada coluna

Masking e limite de tamanho de contexto

• Durante o treinamento, o modelo não prevê apenas um único próximo token de um texto dado; ele também prevê simultaneamente os possíveis próximos tokens após cada subsequência
  • Assim, um único exemplo de texto funciona como vários exemplos de treinamento, o que aumenta a eficiência
• No exemplo do GPT, usa-se masking para impedir que tokens posteriores influenciem anteriores e acabem vazando a resposta correta do próximo token
  • Antes do softmax, os valores dessas posições são definidos como infinito negativo
  • Depois do softmax, essas posições viram 0, e a coluna continua normalizada
• Nem toda atenção usa masking o tempo todo, mas no caso do GPT isso é sempre usado para impedir que tokens futuros influenciem tokens passados
• O tamanho do attention pattern é igual ao quadrado do tamanho do contexto
  • Por isso, o context size pode se tornar uma limitação importante em modelos de linguagem grandes
  • Surgiram variações para tornar o mecanismo de atenção mais escalável para context windows maiores, mas aqui trata-se apenas da forma básica

Como o value atualiza de fato os embeddings

• O attention pattern fornece os pesos de quais palavras atualizam quais palavras; a etapa seguinte é produzir a variação real do embedding
• Multiplica-se cada embedding pela value matrix W_V para criar um vetor value
  • O vetor value fica no mesmo espaço de alta dimensão dos embeddings
  • Ele representa que mudança concreta deve ser somada quando uma palavra relevante ajusta o significado de outra
• Em cada coluna, multiplica-se cada vetor value pelo peso correspondente do attention pattern e soma-se tudo, obtendo a variação \Delta \vec{E}
• Ao somar essa variação ao embedding original, obtém-se um novo embedding \vec{E}' que reflete o contexto
  • No exemplo, creature absorve a informação de fluffy e blue e passa a carregar um significado mais próximo de “fluffy blue creature”
• Aplicando o mesmo processo a todas as colunas, saem do attention block embeddings refinados para toda a sequência de tokens
• Um único head de atenção é parametrizado por três tipos de matrizes treináveis: key matrix, query matrix e value matrix

Cálculo de parâmetros no GPT-3

• No exemplo do GPT-3, as matrizes de key e query têm, cada uma, 12.288 colunas, correspondentes à dimensão do embedding, e 128 linhas, correspondentes à dimensão do espaço key-query
  • Cada matriz tem 1.572.864 parâmetros
• Se a value matrix fosse uma matriz quadrada 12.288×12.288, isso adicionaria 150.994.944 parâmetros, muito mais do que key/query
• Na prática, é mais eficiente decompor o value map em duas matrizes menores para manter o número de parâmetros próximo ao de key/query
  • A primeira matriz reduz o grande espaço de embeddings para um espaço menor, como 128 dimensões
  • A segunda matriz projeta esse espaço pequeno de volta ao espaço de embeddings
  • Do ponto de vista da álgebra linear, isso restringe o value map completo a uma transformação de baixo posto
• Nesta explicação, essas duas matrizes são chamadas de Value_\downarrow e Value_\uparrow, mas esses não são nomes padronizados
• Somando as quatro matrizes, um head de atenção tem cerca de 6,3 milhões de parâmetros

Self-attention e cross-attention

• A estrutura descrita até aqui corresponde mais precisamente a um self-attention head
• Um cross-attention head aparece em modelos que processam dois conjuntos de dados diferentes
  • Por exemplo, em um modelo de tradução, as keys podem vir de uma língua e as queries de outra
  • O attention pattern pode mostrar como palavras de uma língua se relacionam com palavras da outra
• No cross-attention, a diferença em relação ao self-attention é que os mapas de key e query atuam sobre conjuntos de dados diferentes
• Em cenários como tradução, normalmente não há masking, porque não faz sentido pensar em tokens posteriores influenciando anteriores

Multi-headed attention e blocos repetidos

• Na prática, um attention block é composto por multi-headed attention, em que vários heads rodam em paralelo
• O GPT-3 usa 96 heads de atenção em cada bloco
  • 96 matrizes diferentes de key/query produzem 96 attention patterns diferentes
  • Cada head cria sua própria sequência de vetores value com sua matriz value correspondente
  • Em cada posição de token, somam-se as variações \Delta \vec{E} propostas por todos os heads e adiciona-se o resultado ao embedding original
• Executar vários heads em paralelo dá ao modelo capacidade para aprender várias formas pelas quais o contexto altera o significado
• No GPT-3, um único bloco de multi-headed attention com 96 heads tem cerca de 600 milhões de parâmetros
• Em artigos e implementações reais, as matrizes equivalentes a Value_\uparrow de cada head são agrupadas em uma grande output matrix conectada ao bloco inteiro de multi-headed attention
  • Em geral, quando se fala na value matrix de um head específico, costuma-se estar se referindo à primeira projeção, aqui chamada de Value_\downarrow

Como o significado se acumula em Transformers mais profundos

• Os dados dentro do Transformer não passam por apenas um attention block, mas por vários attention blocks e por multi-layer perceptrons
• Mesmo depois que o embedding de uma palavra absorve parte do contexto, ele continua tendo oportunidades de ser influenciado por embeddings vizinhos já mais refinados
• À medida que a rede se aprofunda, cada embedding passa a absorver mais significado de outros embeddings e ganha capacidade de codificar características abstratas de nível mais alto, como sentimento, tom e até se algo é ou não irônico
• O GPT-3 inclui 96 camadas, e os parâmetros relacionados a key/query/value são descritos como totalizando menos de 58 bilhões
• Isso representa cerca de um terço dos parâmetros totais da rede, e a maior parte do restante vem dos blocos entre as atenções
• Uma parte importante do sucesso do mecanismo de atenção não está em um comportamento específico, mas na alta capacidade de paralelização que permite executar muitos cálculos em pouco tempo com GPUs
• Como o aprendizado profundo mostrou que ampliar escala pode gerar grandes melhorias qualitativas no desempenho dos modelos, arquiteturas paralelizáveis que permitem essa expansão têm uma grande vantagem

[GN⁺]

Comentários do Hacker News

• Do ponto de vista de alguém que trabalhou com química quântica e um pouco de machine learning, ao ver este vídeo, a semelhança entre os modelos Transformer e a mecânica quântica chamou bastante a atenção
  Na mecânica quântica, o estado de todo o sistema físico é codificado como um vetor normalizado de altíssima dimensão, isto é, uma semirreta no espaço de Hilbert, e sua evolução no tempo é descrita por um operador de translação temporal que pode ser visto aproximadamente como uma matriz unitária U = exp(-iHt)
  No vídeo, diz-se que a previsão do próximo token é determinada calculando o próximo vetor de embedding sensível ao contexto apenas a partir do último vetor de embedding sensível ao contexto, o que parece o resultado de aplicar uma função de estado linear a um vetor de alta dimensão
  Parece algo semelhante a gerar offline o Hamiltoniano de todo o sistema a partir dos dados de treinamento, depois reparametrizar um subsistema específico, a janela de contexto, em uma base adequada a esse Hamiltoniano, aplicar um passo de evolução temporal e então voltar à base original
  Mas, para quem pesquisou uma área específica, todo problema pode acabar parecendo um prego para o martelo daquela área, então fico curioso se essa semelhança também parece válida para outras pessoas ou se estou forçando demais a barra

  • Acho que essa analogia não se sustenta muito bem. Mesmo ignorando todas as etapas não lineares anteriores, o que sobra é só um sistema dinâmico linear, sem a natureza complexa nem a propriedade unitária que são centrais na mecânica quântica
  • Parece mais que você está apenas descrevendo uma máquina de estados. Codificar estados como vetores e avançar etapas com matrizes não seria mais um detalhe de implementação?
  • Andei pensando um pouco nisso recentemente. Se o tempo não for contínuo, talvez seja possível modelar a evolução temporal do universo aplicando recursivamente algum operador ao estado quântico do universo
    Se uma aplicação do operador avançasse o estado do universo em um tempo de Planck, fico curioso se conseguiríamos observar a diferença entre esse universo e um universo com tempo contínuo
  • Eu tinha um estagiário de doutorado em matemática que dizia que álgebra linear de alta dimensão já era uma área extremamente avançada para os padrões dos anos 1900, e que ainda havia muito espaço para descobertas novas na ciência da computação
    Só agora me ocorreu a conexão com o que estava acontecendo na física naquela época
  • No fim das contas, isso quer dizer que o modelo computacional mais sofisticado que já criamos está começando a se aproximar do algoritmo que define o universo em que vivemos? Seria, por assim dizer, a simulação reaparecendo

• O vídeo do CodeEmporium no YouTube foi mais fácil de acompanhar: https://www.youtube.com/watch?v=Nw_PJdmydZY
  É difícil explicar Transformers por analogia e, na verdade, nem existe uma boa explicação do porquê eles funcionam, então talvez seja melhor apenas mostrar o mecanismo e deixar a interpretação para quem assiste
  Além disso, é mais simples explicar o produto escalar como a projeção de vetores uns sobre os outros

  • A explicação é simplesmente que a rede neural é um algoritmo de ajuste estatístico que aprende a distribuição de probabilidade condicional P(next_word|previous_words). Os pesos são o modelo dessa distribuição, e os LLMs estão mais para uma inovação de hardware que permitiu às GPUs calcular isso em larga escala sobre dados na faixa de terabytes
    A razão de “the cat sat on the ...” ser seguido por “mat” é que essa é a palavra que apareceu com mais frequência no dataset, e a rede neural é um modelo dessas frequências
    A razão de ela parecer saber “London in UK”, mas não “London in France”, é que “UK” aparece com muito mais frequência no dataset
    O algoritmo em si não faz nada especialmente interessante além de alinhar os cálculos ao hardware. O valor vem da estrutura de probabilidade condicional presente nos dados, e essa estrutura é resultado de pessoas organizarem palavras de forma útil para transmitir informações umas às outras
  • Do ponto de vista da ciência da computação, a interpretação como uma tabela hash diferenciável fez bastante sentido para mim. O artigo AIAYN também sugere isso ao usar os nomes query/key/value, mas não menciona explicitamente “hash table”. Talvez isso tenha sido introduzido em outro artigo
  • Meu entendimento pessoal sobre attention é que a saída do Transformer é uma sequência de novos vetores de tokens, e cada vetor de token de saída contém a informação de contexto dos vetores de tokens de entrada ao seu redor
    É uma explicação incompleta, claro, mas ainda melhor do que não ter explicação nenhuma

• Há uma visualização convincente de como um LLM funciona ao lidar com uma solicitação simples: https://bbycroft.net/llm
  Ela complementa muito bem a explicação detalhada do 3blue1brown

  • Ao visualizar assim, dá para sentir o quão absurdamente grande é a escala do GPT-3. Mal consigo imaginar como o GPT-4 pareceria aqui

• Excelente vídeo. Ele mostra muito bem por que a multiplicação de matrizes Q*K é o gargalo. Se a sequência, ou seja, o tamanho da janela de contexto, é S, é preciso armazenar em memória uma matriz SxS com os resultados de todas as queries contra todas as keys
  Uma ideia relativamente nova para melhorar esse gargalo é Ring Attention, e este texto explica bem: https://learnandburn.ai/p/how-to-build-a-10m-token-context
  Eu editei esse texto

  • Com Flash Attention, não é necessário criar a matriz (S, S) de forma alguma. Como a expressão tem a forma softmax(Q @ K^T / sqrt(d)) @ V, é possível produzir a saída final em blocos
    No Unsloth, graças ao Flash Attention, o uso de memória cresce linearmente em vez de quadraticamente, o fine-tuning fica 2x mais rápido, o uso de VRAM cai 80% e a inferência também fica 2x mais rápida. No entanto, a quantidade de computação ainda é O(N^2)
    Em contextos longos, a versão mais recente do Unsloth consegue comportar um contexto 4x maior que HF+FA2 com overhead de +1,9%, chegando a 228K de contexto em uma H100
  • O vídeo também lista Ring Attention e várias outras técnicas, mas diz que isso foge do escopo deste vídeo: https://youtu.be/eMlx5fFNoYc?t=784

• O texto anterior, “But what is a GPT?”, também é muito bom: https://www.3blue1brown.com/lessons/gpt

• Este vídeo me fez perceber que o mecanismo de attention está mais perto de uma espécie de meta-função do que de uma função específica
  Se entendi direito, Attention + pesos aprendidos permitem que o transformer aprenda uma função até certo ponto arbitrária, e essa função inclui um mecanismo de correspondência como o scaled dot-product

  • Isso mesmo. A força da attention está em explorar o espaço de funções e encontrar a melhor função dentro das restrições
    Por isso acho que a linear attention dificilmente conseguirá se aproximar da capacidade da attention padrão. O termo quadrático é uma característica essencial, porque ele explora todos os pares de entrada-saída

• O que tornou este vídeo fácil de digerir foram em grande parte as animações. O jeito como elas expandem, contraem e se desdobram no timing da fala ficou muito bem feito

  • Isso com certeza é algo que ele faz melhor do que a maioria. Ele também tem uma biblioteca de animação customizada que criou para animações matemáticas: https://github.com/3b1b/manim

• Trabalho em uma área bem relacionada, e este vídeo foi direto para a documentação de onboarding do nosso time
  Também é importante que uma boa parte do código de visualização esteja no GitHub: https://github.com/3b1b/videos/tree/master/_2024/transformers

  • Interessante, fiquei curioso sobre o que mais entra nessa documentação de onboarding

• Finalmente entendi. Não sei por que os outros vídeos conseguiam deixar isso tão confuso

  • É um tema que já é confuso por natureza, e o 3b1b é simplesmente muito bom nisso
  • Pela minha experiência, com exceção de raridades como Feynmann, pesquisadores em geral costumam ser péssimos em explicar com clareza aos outros o que fazem
    Isso me faz pensar se habilidade de ensino e habilidade de pesquisa não são, em grande medida, competências mutuamente exclusivas
  • Tenho curiosidade porque quero fazer vídeos e conteúdos educacionais melhores. Queria entender em que exatamente os outros vídeos falharam mais em comparação com o 3b1b
  • Grant tem um talento para explicar coisas complexas com muita clareza. Existe um motivo para o canal dele ser popular
  • Não sei se foi uma pergunta retórica, mas é uma pergunta interessante. Acho que há pelo menos três motivos para a maioria das pessoas se confundir com transformers
    Primeiro, a terminologia padrão não é boa. “attention” já é só vagamente intuitivo, “self-attention” é pior ainda, e “key” e “value” nem se fala
    Segundo, os artigos centrais, como Attention is All You Need e o artigo do BERT, não foram bem escritos. Não estou tentando diminuir o mérito deles, mas mesmo artigos influentes que trazem avanços enormes podem explicar mal, e acho que foi o caso
    Terceiro, essas arquiteturas em geral foram descobertas jogando várias coisas na parede para ver o que grudava. Não houve primeiro um processo de reflexão que levasse à previsão de que essa arquitetura funcionaria bem e depois uma validação experimental; foi empírico do começo ao fim
    Por isso não entendemos completamente por que isso funciona tão bem, e toda explicação fica perto de uma racionalização posterior. Mais recentemente, também há trabalhos sugerindo que outras arquiteturas, se ajustadas o suficiente, podem funcionar quase tão bem. É difícil explicar bem algo que você não entende completamente

• Tenho curiosidade se existe algum material de referência que explique como a arquitetura atual evoluiu. Queria ver a linha que vai de ideias centrais bem simples até o famoso artigo “all you need”
  Caso contrário, muitos dos componentes parecem simplesmente surgir do nada, com muita conta e pouca intuição
  Jeremy Howard disse no Twitter que viu várias versões dessa ideia ao longo do tempo, o que soa como um sinal de que era uma ideia natural. Acho que ver exemplos de como essa ideia apareceu em outros lugares ajudaria a construir intuição

  • O fluxo foi mais ou menos assim. As primeiras abordagens seq-2-seq usavam LSTM: uma rede codificava a sequência de entrada e outra decodificava a sequência de saída. Já era surpreendente que funcionasse codificar frases de comprimento variável em um vetor de tamanho fixo e depois decodificá-las de volta em outra sequência, normalmente de comprimento diferente
    Essa abordagem com RNN/LSTM tinha a fraqueza da representação de tamanho fixo e também a dificuldade de decidir qual parte da sequência de entrada usar ao gerar uma parte específica da saída. Bahdanau e colegas resolveram isso com uma arquitetura que combinava um mecanismo de attention com uma RNN encoder-decoder, permitindo olhar não só para o estado final, mas para todos os estados passados da RNN
    As RNNs eram ineficientes para treinar, então Jakob Uszkoreit queria encontrar uma forma de aproveitar melhor hardware massivamente paralelo, e chamou atenção para o fato de que a linguagem não é apenas sequencial, mas também hierárquica. Ele propôs uma estrutura em camadas em que os tokens de subsequências eram processados em paralelo em cada camada, mantendo a attention no estilo Bahdanau para prever a próxima camada com self-attention, na qual os tokens fazem referência uns aos outros
    As implementações iniciais funcionaram, mas não eram melhores do que outras abordagens da época, como convoluções. Depois, Noam Shazeer desenvolveu mais a ideia e criou uma arquitetura que funcionava muito melhor e, após experimentos removendo componentes desnecessários, isso virou o transformer original, pelo que sei. Não tenho certeza de quem concebeu a forma final de attention baseada em key
    O transformer original do artigo Attention is All You Need, seguindo as abordagens anteriores baseadas em RNN, tinha encoder e decoder separados, e isso também foi usado em modelos iniciais como o BERT do Google. Mas para language models isso não é essencial, então o GPT da OpenAI usou apenas a parte decoder, e hoje praticamente todo mundo segue essa abordagem. Em transformers decoder-only, a frase de entrada entra na camada mais baixa, depois vai sendo transformada camada por camada até sair no topo. No fim da sequência de entrada é anexado um token de término, e ele é transformado no próximo token da sequência de saída, isto é, o último token
  • O Karpathy resumiu bem a história da arquitetura transformer em uma aula de Stanford: https://youtu.be/XfpMkf4rD6E?si=MDICNzZ_Mq9uzRo9&t=618