Desarrollo de LLM, parte 8 - técnica de self-attention entrenable

• En un Transformer decoder-only de estilo GPT, la self-attention entrenable calcula a qué tokens de la entrada previa debe prestar atención cada token para crear vectores de contexto
• La clave es la scaled dot product attention, que usa tres matrices entrenables Wq, Wk, Wv para enviar los embeddings de entrada a los espacios de query, key y value
• La matriz de entrada X se transforma como Q=XWq, K=XWk, V=XWv, y tras dividir Ω=QKᵀ entre √c se obtiene la matriz de pesos de atención A con softmax aplicada por filas
• Los vectores de contexto se crean con una sola multiplicación matricial, C=AV, y todo el cálculo puede aplicarse a todos los tokens con 5 multiplicaciones de matrices y una transposición
• Esta etapa va más allá del ejemplo de juguete que hacía dot product directamente entre embeddings de entrada, y lleva a una attention entrenable que puede implementarse con PyTorch nn.Module y nn.Linear

---

Dónde se ubica la self-attention en el flujo de procesamiento de un LLM

• Un LLM basado en un Transformer decoder-only de estilo GPT es una arquitectura que predice el siguiente token viendo los tokens anteriores
• El flujo de procesamiento consiste en dividir una cadena en tokens, convertir cada token en un token embedding, y luego sumar un positional embedding que representa la información de posición para crear los embeddings de entrada
• La self-attention genera, para cada embedding de entrada, una lista de attention scores que indica cuánto debe fijarse en los demás tokens
  • En la oración de ejemplo "the fat cat sat on the mat", al mirar "cat", "fat" puede ser importante
  • Al mirar "mat", la importancia de "fat" puede ser relativamente menor
• Los attention scores pasan por softmax para convertirse en pesos de atención cuya suma es 1, y con esos pesos se hace una suma ponderada de los embeddings de entrada para formar un vector de contexto
• El vector de contexto se trata como un vector que representa el significado de cada token dentro del contexto de toda la entrada

Objetivo de la self-attention entrenable

• Hasta la etapa anterior se usaba una self-attention de juguete que calculaba directamente el dot product entre embeddings de entrada
• El objetivo de esta etapa es construir un mecanismo de attention entrenable que pueda generar attention scores a partir de los vectores de entrada
• La sección 3.4 de Sebastian Raschka, Build a Large Language Model (from Scratch), lo implementa como scaled dot product attention
• El foco no está tanto en por qué esta estructura es efectiva, sino en con qué cálculos funciona

Matrices Query, Key, Value y proyección de espacios

• Se define n como la longitud de la secuencia de entrada, d como la dimensión del embedding de entrada y c como la dimensión del vector de contexto
• La secuencia de embeddings de entrada se representa como x1, x2, x3, ... xn, y cada embedding de entrada es un vector de dimensión d
• Se definen tres matrices de pesos entrenables
  • query weights matrix: Wq
  • key weights matrix: Wk
  • value weights matrix: Wv
• Cada matriz tiene tamaño d×c y proyecta un vector de entrada de dimensión d a un espacio de dimensión c
• El cálculo para enviar el vector de entrada xm al espacio de query es qm=xmWq
• Los espacios key y value también proyectan los embeddings de entrada a distintos espacios de dimensión c del mismo modo

Ver las matrices como proyecciones

• Las matrices pueden usarse para transformaciones geométricas, como rotar puntos
• Las matrices cuadradas realizan transformaciones dentro de la misma dimensión, y las matrices no cuadradas pueden enviar vectores a espacios de otra dimensión
• Por ejemplo, una matriz 3×2 puede transformar vectores de 3 dimensiones en vectores de 2 dimensiones
• En gráficos 3D, una matriz de frustum que transforma puntos 3D en puntos de una pantalla 2D también es un ejemplo de este tipo de proyección
• La self-attention envía los embeddings de entrada a tres espacios de proyección distintos llamados query, key y value, y luego hace los cálculos con esos vectores proyectados
• Como estas matrices de proyección se aprenden durante el entrenamiento, aparece una indirección que no existía en la atención por dot product simple

Cálculo de los attention scores

• Al considerar una entrada específica xm, el attention score respecto a otra entrada xp se define como el dot product entre la proyección query y la proyección key
• La fórmula es la siguiente
  • qm=xmWq
  • kp=xpWk
  • ωm,p=qm·kp
• Este cálculo podría hacerse con bucles para todas las entradas, pero usando multiplicación matricial puede calcularse todo de una vez
• Si se toma la totalidad de los embeddings de entrada como la matriz X, X tiene tamaño n×d
• La matriz key se calcula de una vez como K=XWk
  • El resultado K tiene tamaño n×c
  • Cada fila es el vector que proyecta el embedding de entrada correspondiente al espacio key
• La matriz query se calcula del mismo modo como Q=XWq
• El dot product entre todas las query y todas las key se obtiene con QKᵀ
  • Q es n×c
  • Kᵀ es c×n
  • El resultado Ω es n×n
• Ωm,p es el attention score que indica cuánto debe atenderse a xp al construir el vector de contexto de xm

Escalado y normalización con softmax

• Los attention scores se convierten, como en el ejemplo anterior, en pesos cuya suma es 1 mediante softmax
• Softmax amplifica los valores grandes, reduce los pequeños y ajusta la lista completa para que la suma total sea 1
• En los LLM reales, d y c pueden llegar a estar en el orden de los miles, por lo que usar solo softmax puede producir gradientes pequeños
• En ese caso, softmax puede comportarse “como una step function”
  • Puede interpretarse como una situación en la que domina el valor más grande y los demás quedan muy pequeños
• Para aliviar esto, los attention scores se dividen por la raíz cuadrada de la dimensión c del espacio proyectado y luego se aplica softmax
• La expresión matricial es la siguiente
  • A=softmax(Ω/√c, axis=1)
• axis=1 es notación al estilo PyTorch y significa que softmax se aplica por filas
• El resultado A es la matriz de attention scores normalizados, es decir, la matriz de pesos de atención

Creación de vectores de contexto

• La proyección al espacio value se calcula como V=XWv
• A es una matriz de pesos de atención de tamaño n×n
  • Am,p es el peso de atención que se aplica a la entrada p al construir el vector de contexto de xm
• V tiene tamaño n×c, y cada fila es el vector que proyecta un embedding de entrada al espacio value
• La matriz de vectores de contexto se calcula como C=AV
  • El resultado C es n×c
  • La fila m de C es el vector de contexto para la entrada xm
• Este cálculo realiza, con una sola multiplicación matricial, la operación de multiplicar los vectores value por los pesos de atención y sumarlos para cada token

Resumen del cálculo completo

• La matriz de entrada X contiene los embeddings de entrada de la secuencia de tokens y su tamaño es n×d
• Con tres matrices entrenables, la entrada se proyecta por separado a los espacios query, key y value
  • Q=XWq
  • K=XWk
  • V=XWv
• Los attention scores se calculan con el dot product entre query y key
  • Ω=QKᵀ
• Tras escalar los puntajes, se aplica softmax por filas para crear los pesos de atención
  • A=softmax(Ω/√c, axis=1)
• Al multiplicar la proyección value por los pesos de atención, se generan los vectores de contexto
  • C=AV
• Todo el mecanismo de self-attention puede crear los vectores de contexto de todos los tokens de entrada con 5 multiplicaciones de matrices y una transposición

Implementación en PyTorch y siguiente paso

• La sección 3.4 del libro implementa el cálculo anterior en código PyTorch y crea una subclase simple de nn.Module que realiza las mismas operaciones matriciales
• La primera versión usa objetos nn.Parameter comunes para las tres matrices de pesos
• La segunda versión usa nn.Linear para un entrenamiento más efectivo
• Después se tratan dos temas
  • causal self-attention: una forma en la que, al mirar un token, no se presta atención a los tokens posteriores
  • multi-head attention: se adelanta como un tema no tan complicado como se pensaba al principio
• El procesamiento por lotes queda como una cuestión aparte
  • Incluso con una sola secuencia de entrada se usa una matriz de attention scores
  • Para procesar varias secuencias de entrada en paralelo, podrían necesitarse tensores de orden superior a una matriz
• El siguiente artículo continúa en Writing an LLM from scratch, part 9 -- causal attention