Explicando como os transformers funcionam: entendendo a matemática por trás

• Reduz o processo de inferência de um transformer usando o exemplo de tradução Hello World → Hola Mundo, para que seja possível acompanhar manualmente desde a tokenização até o encoder, decoder e o cálculo da probabilidade do próximo token
• Em vez da configuração grande do artigo original, usa embeddings de 4 dimensões, 2 cabeças de atenção e uma camada feedforward de 8 dimensões para tornar menores o fluxo de multiplicações de matrizes e do softmax
• O encoder soma codificação posicional aos embeddings dos tokens e então produz uma representação contextual da sequência de entrada passando por self-attention multi-head e por uma camada feedforward
• O decoder começa em SOS, usa em conjunto os tokens gerados anteriormente e a saída do encoder, e no encoder-decoder attention a query vem do decoder enquanto key/value são calculados a partir da saída do encoder
• O embedding final do decoder passa por uma camada linear e pelo softmax para virar a probabilidade do próximo token, mas o exemplo usa pesos aleatórios, então não se espera qualidade real de tradução

---

Objetivo e premissas

• Ver em um exemplo end-to-end como a matemática na inferência se encadeia dentro de um modelo transformer
• Reduz bastante o tamanho do modelo para facilitar acompanhar os cálculos à mão
  • Em vez da dimensão de embedding 512 do artigo original, o exemplo usa 4 dimensões
  • Em vez de 8 cabeças de atenção como no artigo original, usa 2
  • Em vez da dimensão feedforward 2048 do artigo original, usa 8 dimensões
• O pré-requisito necessário é álgebra linear básica, e a maior parte dos cálculos é feita com multiplicação de matrizes
• O foco está menos no que o transformer “é” e mais em como os cálculos realmente acontecem
• Para uma explicação mais intuitiva, vale ler junto The Illustrated Transformer, e o artigo original é Attention is all you need

Montando a entrada do encoder

• Tokenização

  • Como modelos de machine learning processam números, não texto, o texto de entrada é convertido em IDs de token
  • Para simplificar, o exemplo divide "Hello World" em dois tokens de palavra: "Hello" e "World"
  • Na prática, os métodos de tokenização podem ser baseados em palavra, caractere ou subword
  • O modelo baseado em palavras precisa de um vocabulary grande e trata "dog" e "dogs" como tokens diferentes
  • O modelo baseado em caracteres usa um vocabulary pequeno, mas pode ter menos informação semântica
  • A tokenização subword fica no meio do caminho entre palavras e caracteres, e o tokenizador é treinado por um processo estatístico

• Embeddings de tokens

  • Como o próprio ID do token não carrega significado, cada token é convertido em um vetor de tamanho fixo chamado embedding
  • Os embeddings do exemplo usam valores arbitrários
    • Hello -> [1, 2, 3, 4]
    • World -> [2, 3, 4, 5]
  • Em transformers reais, esse mapeamento de embedding também é aprendido, e o modelo aprende representações de tokens adequadas à tarefa
  • Os dois embeddings são reunidos em uma única matriz para uso nas multiplicações de matrizes seguintes

• Codificação posicional

  • Só com embeddings não dá para saber a posição da palavra na frase, então soma-se uma codificação posicional
  • O artigo original usa uma codificação posicional fixa com seno/cosseno, e o exemplo segue o mesmo método
  • A codificação posicional do exemplo é calculada assim
    • Hello -> [0, 1, 0, 1]
    • World -> [0.84, 0.99, 0, 1]
  • Somando o embedding do token com a codificação posicional, obtém-se a matriz de entrada do encoder
    • Hello -> [1, 3, 3, 5]
    • World -> [2.84, 3.99, 4, 6]

Cálculo de self-attention

• Criando Q, K, V

  • A self-attention calcula query (Q), key (K) e value (V) a partir dos embeddings de entrada
  • O exemplo usa 2 cabeças de atenção, e cada uma tem suas próprias matrizes WQ, WK, WV
  • Cada matriz de pesos transforma embeddings de 4 dimensões em query/key/value de 3 dimensões
  • Na primeira cabeça, multiplica-se a matriz de entrada por WK1, WV1 e WQ1 para obter K1, V1 e Q1

• Fórmula da attention

  • As pontuações de attention são calculadas em quatro etapas
    • Calcula-se o produto escalar entre a query e cada key
    • Divide-se pela raiz quadrada da dimensão da key
    • Converte-se em pesos positivos cuja soma é 1 usando softmax
    • Usa-se esses pesos para fazer uma soma ponderada dos vetores value
  • Esse processo é condensado na fórmula do artigo original
  • [
  • Attention(Q,K,V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d}}\right)V
  • ]
  • No exemplo, por causa da baixa dimensionalidade e dos valores iniciais arbitrários, o resultado do softmax fica quase todo concentrado em 0 e 1
  • Valores grandes no dot product podem ser amplificados ainda mais pelo softmax, por isso é necessário o escalonamento dividindo pela raiz quadrada da dimensão da key
  • Para fins de explicação, também se usa temporariamente uma variação dividindo por 30 em vez de sqrt(3), mas isso não é uma solução de longo prazo

• Saída da atenção multi-head

  • Os resultados de attention de cada cabeça são concatenados e então multiplicados por uma matriz de pesos aprendida para voltar à dimensão de embedding
  • No exemplo, os resultados das duas cabeças formam uma matriz de 6 dimensões, que é convertida em uma saída de 4 dimensões
  • Essa saída é enviada para a próxima etapa do bloco do encoder, a camada feedforward

Camada feedforward e bloco do encoder

• Camada feedforward

  • Depois da self-attention vem uma rede neural feedforward (FFN)
  • A FFN é composta por duas transformações lineares com uma ativação ReLU entre elas
  • A primeira camada linear expande a dimensão, e a segunda reduz a dimensão de volta ao tamanho original
  • A ReLU transforma valores negativos em 0 e mantém os positivos, adicionando não linearidade
  • No exemplo, a entrada de 4 dimensões é expandida para 8 e depois reduzida novamente para 4
  • [
  • \text{FFN}(x) = \text{ReLU}(xW_1 + b_1)W_2 + b_2
  • ]

• Bloco do encoder

  • Um bloco de encoder é composto por atenção multi-head e FFN
  • O artigo original empilha 6 encoders, e o código do exemplo também repete o encoder com n=6
  • Se vários blocos de encoder forem atravessados simplesmente em sequência, os valores podem crescer demais, causando overflow no cálculo do softmax e resultando em nan

Residual connection e layer normalization

• Problema de explosão dos valores

  • No exemplo, ao passar por 6 encoders, surgem os avisos overflow encountered in exp e invalid value encountered in divide, e a saída vira nan
  • Esse fenômeno em que os valores ficam grandes demais e crescem ainda mais na próxima camada é um problema comum em redes neurais profundas
  • Quando os gradientes ficam grandes demais durante o backpropagation, isso é chamado de gradient explosion

• Residual connection

  • Residual connection é a técnica de somar a entrada da camada à saída da própria camada
  • [
  • \text{Residual}(x) = x + \text{Layer}(x)
  • ]
  • No exemplo, residual connections são aplicadas tanto à saída da attention quanto à saída da FFN
  • Residual connections são usadas para aliviar o problema de vanishing gradient

• Layer normalization

  • Layer normalization normaliza cada dimensão do embedding para ter média 0 e desvio padrão 1
  • A fórmula é a seguinte
  • [
  • \text{LayerNorm}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \times \gamma + \beta
  • ]
  • (\epsilon) é um valor pequeno para evitar divisão por zero quando o desvio padrão é 0
  • (\gamma) e (\beta) são parâmetros aprendidos que controlam scaling e shifting
  • Depois de adicionar residual connection e layer normalization, mesmo passando por 6 encoders o modelo produz valores normais sem nan

Estrutura do decoder

• Entrada do decoder e forma de geração

  • O decoder recebe como entrada a saída do encoder e a sequência gerada até o momento
  • Durante a inferência, ele começa com o token SOS (start-of-sequence)
  • O decoder gera um token por vez de forma autoregressiva
    • Rodada 1: recebe SOS como entrada e gera "hola"
    • Rodada 2: recebe SOS + hola como entrada e gera "mundo"
    • Rodada 3: recebe SOS + hola + mundo como entrada e gera EOS
  • Quando o token EOS (end-of-sequence) é gerado, o decoding para
  • O encoder pode produzir sua representação em uma única forward pass, mas o decoder precisa fazer várias forward passes, então é mais lento

• Composição do bloco do decoder

  • O bloco do decoder é mais complexo que o do encoder e é composto na seguinte ordem
    • masked self-attention
    • residual connection e layer normalization
    • encoder-decoder attention
    • residual connection e layer normalization
    • camada feedforward
    • residual connection e layer normalization
  • No exemplo de inferência, soma-se codificação posicional ao embedding de SOS para usar [1, 1, 0, 1]
  • Durante o treinamento, usa-se masked self-attention que mascara a pontuação de attention com -inf para impedir que se veja tokens futuros

Encoder-decoder attention

• Encoder-decoder attention é a etapa que faz o decoder focar nas partes relevantes da frase de entrada
• O modo de cálculo é o mesmo da self-attention, mas a origem de Q/K/V é diferente
  • A query é calculada a partir da saída da camada anterior do decoder
  • Key e value são calculados a partir da saída do encoder
• Graças a essa estrutura, cada posição do decoder pode consultar todas as posições da sequência de entrada
• Isso é útil em tarefas como tradução, em que os tokens de saída precisam depender de posições relevantes da frase de entrada

Geração do token de saída

• Camada linear e softmax

  • A saída do decoder ainda não é uma palavra, então o embedding final passa por uma camada linear para virar um vetor de logits com tamanho igual ao do vocabulary
  • O vocabulary do exemplo tem tamanho 10, e os candidatos a próximo token são os seguintes
    • hello, mundo, world, how, ?, EOS, SOS, a, hola, c
  • Os logits passam pelo softmax e viram uma distribuição de probabilidade sobre cada token
  • Na probabilidade do exemplo, "hola" tem a maior probabilidade e é escolhido como próximo token
  • Escolher sempre o token de maior probabilidade é chamado de greedy decoding, e nem sempre é a melhor estratégia
  • É possível ver mais detalhes sobre técnicas de geração neste texto da Hugging Face

• Loop completo de geração

  • O procedimento completo de geração segue este fluxo
    • Converte a sequência de entrada em embeddings
    • O encoder gera a representação contextual de toda a entrada
    • O decoder começa em SOS e usa em conjunto os tokens gerados anteriormente e a saída do encoder
    • Aplica-se a camada linear e o softmax ao embedding final do decoder
    • Escolhe-se o próximo token mais provável e ele é adicionado à sequência
    • Repete-se até aparecer EOS ou atingir o comprimento máximo
  • A execução do exemplo gera SOS hola mundo world para a entrada hello world
  • Como todos os pesos e embeddings são aleatórios, o resultado não é uma boa tradução, e esse é o comportamento esperado

Conclusão e escopo

• O exemplo conecta em um único fluxo os principais componentes do transformer: embeddings, codificação posicional, self-attention, atenção multi-head, FFN, residual connection, layer normalization, encoder-decoder attention e saída via softmax
• Arquiteturas transformer mais recentes adicionam várias técnicas, mas a matemática central se baseia na estrutura tratada neste exemplo
• A stack usada pode variar conforme o tipo de tarefa
  • Em tarefas focadas em compreensão, como classificação, pode-se usar uma camada linear sobre a stack de encoder
  • Em tarefas focadas em geração, como tradução, pode-se usar juntas as stacks de encoder e decoder
  • Em tarefas de geração livre, como ChatGPT ou Mistral, pode-se usar apenas a stack de decoder
• O processo de treinamento não é abordado; o foco está em entender a matemática da inferência ao usar um modelo já existente
• Para um material matemático mais formal, consulte este PDF

[GN⁺]

Comentários do Hacker News

• O “mistério” do Transformer está no fato de que, em vez de multiplicar pesos e valores estáticos em ordem linear em cada camada, ele cria 3 matrizes obtidas ao multiplicar a mesma entrada por pesos aprendidos e multiplica essas matrizes entre si.
  Isso funciona bem por aumentar o paralelismo, mas a própria fórmula de atenção é fixa, então é muito limitada.
  Para avançar mais, parece necessário encontrar uma forma de generalizar o próprio grafo computacional como parâmetros aprendíveis. Não sei se isso é possível com métodos tradicionais de gradiente, por causa de efeitos caóticos em que pequenas mudanças levam a grandes variações de desempenho; talvez seja preciso algo internamente como algoritmos genéticos ou otimização por enxame de partículas.

  • Essa explicação não está nada correta. O que há de especial no Transformer é permitir que cada elemento da sequência escolha, entre todos os outros elementos, as partes importantes para si e as extraia para fazer o cálculo.
    A grande vantagem teórica em relação a RNNs é que ele dá suporte a isso sem perdas. Isso porque cada elemento pode acessar todas as informações de todos os outros elementos da sequência, ou, na ordem temporal, todas as informações dos elementos anteriores.
    Já RNNs e “Transformers lineares” comprimem os valores passados, então normalmente é difícil para o último elemento de uma sequência longa acessar todas as informações do primeiro; é impossível a menos que o estado interno seja enorme e não descarte nenhuma informação.
  • Dizer que “é preciso generalizar o grafo computacional como parâmetros aprendíveis para avançar mais” é uma afirmação bem ousada. Afinal, já houve um progresso enorme sem aprender o grafo computacional.
  • Esse método basicamente pode aprender a ignorar alguns caminhos e amplificar os mais importantes, e depois é possível podar caminhos cuja perda de qualidade não seja grande.
    O problema é que não se ganha muito com isso. Operações que não sejam multiplicação de matrizes provavelmente serão mais lentas ou terão velocidade parecida.
  • Concordo com a ideia de generalizar o grafo computacional como parâmetros aprendíveis. Isso parece semelhante à forma como os processos mentais humanos resolvem os problemas que eu gostaria que LLMs resolvessem: problemas mais próximos de raciocínio real, além do “processamento de linguagem” em que Transformers são bons.
    Porém, ao adicionar controle de fluxo, há o risco de isso virar, na prática, uma máquina de Turing; nesse caso, como você disse, o treinamento vira o problema. Ainda assim, talvez não seja um problema totalmente intratável.
  • Ajuste de hiperparâmetros também já caminha, em alguma medida, na direção de aprender o grafo computacional. Mas é algo muito limitado e exige muito mais treinamento.

• Para uma explicação mais seca, formal e concisa, há “The Transformer Model in Equations”, de John Thickstun [0].
  Tudo cabe em uma página usando notação matemática padrão.
  [0] https://johnthickstun.com/docs/transformers.pdf

  • Finalmente é bom ver uma explicação assim. Acho 7 linhas de notação matemática muito melhores do que várias páginas de conversa qualitativa.
    Muitas vezes parece que pesquisadores de machine learning nunca estudaram matemática.
  • Ao ler artigos, precisei montar esse tipo de conteúdo nas minhas anotações pessoais várias vezes, e sempre ficava sem certeza se não tinha deixado algo de fora.

• A explicação de que “aparece NaN, os valores ficam grandes demais e explodem ao passar para o próximo encoder; isso é explosão de gradiente” está errada, pelo que entendi.
  Aqui não se calcula gradiente em nenhum ponto, então não é explosão de gradiente.
  O problema parece estar na implementação do softmax, e uma forma numericamente estável de implementar softmax é explicada aqui [0].
  [0]: https://jaykmody.com/blog/stable-softmax/

  • Correto. Eu tentei conectar os problemas comuns de treinamento, explosão e desaparecimento de gradientes, com a sensibilidade do softmax a valores grandes, mas concordo que isso é enganoso ou impreciso, então vou reescrever essa parte.
    Ainda assim, a rede neural como um todo é sensível a valores grandes, portanto um softmax numericamente estável por si só não resolve. Para a rede funcionar, normalização é essencial.

• Tutoriais de Transformer talvez virem os novos tutoriais de Monad. É um conceito difícil de entender, mas é do tipo que só se entende depois de lutar com exemplos e praticar.
  Como muitas coisas em ciência da computação.

  • No momento em que você entende Transformers, deixa de conseguir explicá-los.
  • Monad é um conceito difícil? Monad é só um monoide na categoria dos endofuntores; qual é o problema?
  • Estou esperando um post de blog intitulado “You could have invented transformers”.

• Li só seis parágrafos e já tenho uma pergunta.
  Em Hello -> [1,2,3,4] World -> [2,3,4,5], dizem que os vetores são aleatórios, mas parece haver um padrão. Fico me perguntando se o 2 presente nos dois vetores significa alguma coisa, ou se é o conjunto inteiro que cria a unicidade.

  • A reutilização dos números é mais um caso de o autor ter sido um pouco preguiçoso. Dá para ver se os dois vetores apontam em direções parecidas ou calcular o ângulo para estimar a similaridade.
    Aqui eles estão separados por cerca de 60 graus e apontam mais ou menos na mesma direção; como o exemplo tenta evitar números negativos, os vetores ficam mais parecidos do que seriam na prática.
    O fato de números terem sido reutilizados não significa nada por si só. O 1 na primeira posição quase não tem relação com o 1 na segunda posição. Afinal, não se está fazendo convolução sobre esse vetor.
  • Não é um bom exemplo. O vetor de cada token é inicializado aleatoriamente, com cada elemento sorteado de uma distribuição normal.
    Depois do treinamento, palavras semelhantes terão alguma similaridade de cosseno, mas quase nunca uma similaridade de cosseno tão alta quanto a de [1,2,3,4] e [2,3,4,5].

• Não é uma pergunta totalmente relacionada, mas estou procurando algum artigo ou paper que explique por que um Transformer consegue lidar com perguntas como as abaixo, mesmo funcionando simplesmente como um “preditor do próximo token”

  1. Quando processa palavras, subpalavras ou tokens desconhecidos que não apareceram no conjunto de dados de treinamento. Ex.: criar no pandas uma tabela com "sdsfs_ff", "fsdf_value" como colunas
  2. Quando criamos um exemplo que não estava nos dados de treinamento e pedimos ao LLM uma saída parecida
     Parece ser uma pergunta comum, mas não consigo encontrar as palavras-chave para pesquisar. Também seria bom ter links que tratem de embeddings posicionais em profundidade; ainda não obtive uma resposta satisfatória sobre por que se usam seno/cosseno e sobre multiplicação versus soma
  • Meu palpite é que até caracteres individuais podem ser codificados como tokens, mas dentro do modelo isso usa mais largura de banda, e eles carregam inerentemente menos significado do que tokens de palavras específicas
    Se o modelo julgar necessário, pode reproduzir uma sequência desconhecida copiando tokens de caracteres individuais ou, se fizer sentido no contexto, pode criá-la do zero
  • P(X_1=x_1, X_2=x_2, X_3=x_3) = P(X_3=x_3 | X_1=X_1, X_2=x_2) • P(X_1=x_1, X_2=x_2)
= P(X_3=x_3 | X_1=X_1, X_2=x_2) • P(X_2=x_2 | X_1=x_1) • P(X_1=x_1)
    Ou seja, se houver a distribuição de probabilidade condicional correta para o próximo token dados os tokens anteriores, também se obtém a distribuição de probabilidade correta para toda a sequência de tokens
    “A distribuição de probabilidade correta para uma sequência de tokens”, ou a distribuição de probabilidade condicional correta de uma sequência de tokens dadas certas condições, é algo que, na prática, permite descrever quase qualquer tipo de comportamento de entrada/saída nesses termos
    Portanto, “funciona prevendo o próximo token” não é, em princípio, uma grande restrição sobre quais comportamentos de entrada/saída ele pode ter
    Mesmo que faça algo impressionante, isso não entra em conflito com o fato de que a saída vem de P(X_{n+1}=x_{n+1} | X_1=x_1, ..., X_n=x_n), isto é, de “previsão do próximo token”
  • Ele não está reproduzindo strings exatas dos dados de treinamento, mas sim padrões e padrões de padrões
    Prever o próximo token é uma tarefa mais inteligente do que parece

• Concordo com a frase “a complexidade vem do número de etapas e do número de parâmetros”
  Modelos Transformer simples o bastante para entendermos não fazem coisas interessantes, e Transformers complexos o bastante para fazer coisas interessantes parecem complexos demais para entendermos
  Gostaria de estudar modelos de porte intermediário, simples o bastante para serem compreendidos e complexos o bastante para fazerem coisas interessantes

  • Se você ainda não conhece, talvez ache interessante a pesquisa na área de interpretabilidade mecanística. Neel Nanda publicou bastante material acessível sobre o tema: https://www.neelnanda.io/mechanistic-interpretability
  • Acho que, na prática, a fronteira entre esses dois lados deve ser assim. É bem provável que a zona intermediária já seja complexa demais para humanos entenderem direito, mas ao mesmo tempo ainda pequena demais para fazer coisas interessantes

• Quando conceitos são usados sem serem definidos ou apresentados, fica difícil entender. A seção Encoder começa diretamente, sem explicar o que ele é nem onde se encaixa no processo geral
  Dá para perceber o que o autor quer fazer, mas falta a estrutura básica de texto: primeiro apresentar a ideia, explicá-la e só então usá-la
  Para quem já não está estudando o assunto e não o entende pela metade, o texto inteiro acaba parecendo confuso

• Eu já escrevi uma ANN do zero e não usei TensorFlow, mas ainda assim esta explicação continua confusa
  Pedi ao ChatGPT que explicasse como modificar uma ANN básica para implementar self-attention sem usar as palavras Matrix ou Vector, e ele deu uma explicação bem simples. Ainda não tentei implementar
  Prefiro pensar em tudo em termos de nós, pesos e camadas. Matrizes e vetores tornam mais difícil conectar a explicação ao que realmente acontece dentro de uma ANN
  Na forma familiar de escrever uma ANN, cada nó de entrada é um escalar, mas como o algoritmo de forward pass multiplica e soma pesos para todos os nós de entrada, ele parece uma multiplicação vetor-matriz. Tenho a sensação de que estou abordando essas explicações com a mentalidade errada, e talvez me falte o conhecimento de base necessário