Parte 8 do processo de desenvolvimento de LLM — técnica de self-attention treinável

• Em Transformers no estilo GPT, apenas com decodificador, a self-attention treinável calcula em quais tokens da entrada anterior cada token deve focar para criar vetores de contexto
• O ponto central é a scaled dot product attention, que usa três matrizes treináveis Wq, Wk e Wv para enviar os embeddings de entrada aos espaços de query, key e value
• A matriz de entrada X é transformada em Q=XWq, K=XWk e V=XWv, e depois de dividir Ω=QKᵀ por √c, obtém-se a matriz de pesos de atenção A com softmax aplicada por linha
• Os vetores de contexto são produzidos com uma única multiplicação de matrizes, C=AV, e todo o cálculo pode ser aplicado a todos os tokens com 5 multiplicações de matrizes e uma transposição
• Esta etapa vai além do exemplo didático em que se fazia dot product diretamente entre embeddings de entrada e leva a uma attention treinável que pode ser implementada com PyTorch nn.Module e nn.Linear

---

Posição da self-attention no fluxo de processamento de um LLM

• Um LLM baseado em Transformer apenas com decodificador, no estilo GPT, é uma arquitetura que prevê o próximo token olhando para os tokens vistos até então
• O fluxo de processamento segue a ordem: dividir a string em tokens, converter cada token em um embedding de token e depois somar um embedding posicional que representa a informação de posição para formar os embeddings de entrada
• A self-attention gera, para cada embedding de entrada, uma lista de pontuações de atenção que indica o quanto ele deve focar nos outros tokens
  • Na frase de exemplo "the fat cat sat on the mat", ao observar "cat", "fat" pode ser importante
  • Ao observar "mat", a importância de "fat" pode ser relativamente menor
• As pontuações de atenção passam por softmax e se tornam pesos de atenção cuja soma é 1; esses pesos são usados para fazer uma soma ponderada dos embeddings de entrada e criar um vetor de contexto
• O vetor de contexto é tratado como um vetor que representa o significado de cada token dentro do contexto de toda a entrada

Objetivo da self-attention treinável

• Até a etapa anterior, era usada uma self-attention didática que calculava diretamente o dot product entre embeddings de entrada
• O objetivo desta etapa é montar um mecanismo de atenção treinável capaz de produzir pontuações de atenção a partir dos vetores de entrada
• A seção 3.4 de Sebastian Raschka em Build a Large Language Model (from Scratch) implementa isso como scaled dot product attention
• O foco está menos em por que essa estrutura é eficaz e mais em como ela funciona em termos de cálculo

Matrizes Query, Key, Value e projeção em espaços

• Considere o comprimento da sequência de entrada como n, a dimensão do embedding de entrada como d e a dimensão do vetor de contexto como c
• A sequência de embeddings de entrada é representada como x1, x2, x3, ... xn, e cada embedding de entrada é um vetor de dimensão d
• Definem-se três matrizes de pesos treináveis
  • query weights matrix: Wq
  • key weights matrix: Wk
  • value weights matrix: Wv
• Cada matriz tem tamanho d×c e projeta um vetor de entrada de dimensão d para um espaço de dimensão c
• O cálculo que envia o vetor de entrada xm para o espaço de query é qm=xmWq
• Os espaços de key e value também projetam os embeddings de entrada para espaços c-dimensionais distintos da mesma forma

Entendendo matrizes como projeções

• Matrizes podem ser usadas em transformações geométricas, como rotacionar pontos
• Matrizes quadradas realizam transformações dentro da mesma dimensão, enquanto matrizes não quadradas podem enviar vetores para espaços de outra dimensão
• Por exemplo, uma matriz 3×2 pode transformar vetores 3D em vetores 2D
• Em gráficos 3D, a matriz de frustum que transforma pontos 3D em pontos 2D na tela também é um exemplo desse tipo de projeção
• A self-attention envia os embeddings de entrada para três espaços de projeção diferentes — query, key e value — e depois realiza os cálculos usando os vetores projetados
• Como essas matrizes de projeção são aprendidas durante o treinamento, surge uma indireção que não existia na atenção por dot product simples

Cálculo das pontuações de atenção

• Ao considerar uma entrada específica xm, a pontuação de atenção para outra entrada xp é definida como o dot product entre a projeção em query e a projeção em key
• As fórmulas são as seguintes
  • qm=xmWq
  • kp=xpWk
  • ωm,p=qm·kp
• É possível fazer esse cálculo com loops para todas as entradas, mas o uso de multiplicação de matrizes permite calcular tudo de uma vez
• Se tomarmos toda a entrada embutida como a matriz X, então X tem tamanho n×d
• A matriz key é calculada de uma vez como K=XWk
  • O resultado K tem tamanho n×c
  • Cada linha é o vetor correspondente ao embedding de entrada projetado no espaço de key
• A matriz query é calculada da mesma forma, como Q=XWq
• O dot product entre todas as queries e todas as keys é obtido com QKᵀ
  • Q é n×c
  • Kᵀ é c×n
  • O resultado Ω é n×n
• Ωm,p é a pontuação de atenção que indica o quanto se deve focar em xp ao criar o vetor de contexto de xm

Escalonamento e normalização com softmax

• Assim como no exemplo anterior, as pontuações de atenção passam por softmax para se transformar em pesos cuja soma é 1
• O softmax amplia valores altos, reduz valores baixos e ajusta a lista inteira para que a soma seja 1
• Em LLMs reais, d e c podem chegar à casa dos milhares, e usar apenas softmax pode gerar gradientes pequenos
• Nessa situação, o softmax pode se comportar “como uma step function”
  • Isso pode ser interpretado como um cenário em que o maior valor domina e os demais se tornam muito pequenos
• Para aliviar isso, divide-se a pontuação de atenção pela raiz quadrada da dimensão do espaço projetado, c, antes de aplicar o softmax
• Em notação matricial, isso fica assim
  • A=softmax(Ω/√c, axis=1)
• axis=1 é a notação no estilo PyTorch e significa aplicar o softmax por linha
• O resultado A é a matriz de pontuações de atenção normalizadas, ou seja, a matriz de pesos de atenção

Criação dos vetores de contexto

• A projeção no espaço de value é calculada como V=XWv
• A é uma matriz de pesos de atenção de tamanho n×n
  • Am,p é o peso de atenção aplicado à entrada p ao criar o vetor de contexto de xm
• V tem tamanho n×c, e cada linha é um vetor do embedding de entrada projetado no espaço de value
• A matriz de vetores de contexto é calculada como C=AV
  • O resultado C é n×c
  • A linha m de C é o vetor de contexto da entrada xm
• Esse cálculo faz, com uma única multiplicação de matrizes, o trabalho de multiplicar os vetores value pelos pesos de atenção e somá-los para cada token

Resumo do cálculo completo

• A matriz de entrada X contém os embeddings de entrada da sequência de tokens e tem tamanho n×d
• Três matrizes treináveis projetam a entrada, respectivamente, para os espaços de query, key e value
  • Q=XWq
  • K=XWk
  • V=XWv
• As pontuações de atenção são calculadas pelo dot product entre query e key
  • Ω=QKᵀ
• Depois de escalar as pontuações, aplica-se softmax por linha para gerar os pesos de atenção
  • A=softmax(Ω/√c, axis=1)
• Multiplicam-se a projeção em value e os pesos de atenção para produzir os vetores de contexto
  • C=AV
• Todo o mecanismo de self-attention pode produzir os vetores de contexto de todos os tokens de entrada com 5 multiplicações de matrizes e uma transposição

Implementação em PyTorch e próximos passos

• A seção 3.4 do livro implementa os cálculos acima em código PyTorch e cria uma subclasse simples de nn.Module que realiza as mesmas operações matriciais
• A primeira versão usa objetos nn.Parameter comuns para as três matrizes de pesos
• A segunda versão usa nn.Linear para um treinamento mais eficiente
• Depois disso, há dois temas a tratar
  • causal self-attention: uma abordagem em que, ao observar um determinado token, não se presta atenção aos tokens posteriores
  • multi-head attention: apresentada como um tema menos complexo do que parecia inicialmente
• O processamento em lote continua como uma questão separada
  • Mesmo com uma única sequência de entrada, usa-se uma matriz de pontuações de atenção
  • Para processar várias sequências de entrada em paralelo, pode ser necessário usar tensores de ordem superior às matrizes
• O próximo texto continua em Writing an LLM from scratch, part 9 -- causal attention