Visualización de la atención: el corazón del Transformer [Video]

• La atención en Transformer es el mecanismo que actualiza las incrustaciones de tokens según el contexto, moviendo incluso la misma palabra hacia vectores con distinto significado dependiendo de las palabras que la rodean
• Una cabeza de atención crea vectores query/key/value para cada token, y calcula el attention pattern —los pesos de relevancia entre palabras— mediante el producto punto key-query y softmax
• Los modelos autorregresivos estilo GPT aplican masking para que los tokens posteriores no influyan sobre los anteriores, y como el tamaño del attention pattern crece con el cuadrado de la longitud del contexto, ampliar mucho la context window resulta costoso
• En el ejemplo de GPT-3, las matrices key/query tienen 1,572,864 parámetros cada una, y el value map se divide en una transformación de bajo rango, dando alrededor de 6.3 millones de parámetros por cabeza
• Un Transformer que repite múltiples cabezas y bloques de atención aprende distintas formas de actualizar el contexto, y una parte clave de su éxito está en la capacidad de paralelización para procesar muchos cálculos rápidamente en GPU

---

El papel de la atención en Transformer

• Transformer usa el texto de entrada para predecir el siguiente token, y la entrada primero se tokeniza en palabras o fragmentos de palabras
• Cada token se convierte en una incrustación de alta dimensión
  • La dirección dentro de este espacio de incrustación puede corresponder a significado
  • Por ejemplo, moverse en cierta dirección podría llevar la incrustación de un sustantivo masculino a la incrustación correspondiente de un sustantivo femenino
• El objetivo de la atención es ajustar gradualmente las incrustaciones iniciales para que contengan no solo información individual de cada palabra, sino también un significado contextual más rico

Por qué la misma palabra cambia según el contexto

• En “American shrew mole”, “One mole of carbon dioxide” y “Take a biopsy of the mole”, mole tiene significados distintos
• En la etapa inicial de incrustación, el vector de mole es prácticamente el mismo en los tres casos, porque se parece a una lookup table que no mira el contexto
• En la siguiente etapa, el attention block, las incrustaciones vecinas pueden enviar información a la incrustación de mole para actualizar su valor
• Un modelo bien entrenado conecta los distintos significados de mole con direcciones diferentes dentro del espacio de incrustación, y calcula qué debe sumarse a la incrustación general según el contexto
• Como en “Eiffel tower” y “miniature Eiffel tower”, la incrustación de una palabra puede actualizarse no solo con información de palabras cercanas, sino también con información que llega desde tokens lejanos
• Como para predecir la siguiente palabra solo se usa el último vector, la incrustación de la última palabra en una entrada larga debe contener hasta cierto punto toda la información contextual necesaria para la predicción

Flujo de cálculo de una sola cabeza de atención

• La explicación básica se desarrolla tomando como referencia una single head of attention
• En la oración de ejemplo “A fluffy blue creature roamed the verdant forest.” se asume una situación en la que los adjetivos actualizan la incrustación inicial del sustantivo correspondiente
  • Este ejemplo sirve para mostrar el tipo de operación que puede realizar una cabeza de atención
  • El comportamiento real de una cabeza es difícil de interpretar, porque es el resultado de muchos parámetros ajustados para reducir la función de costo
• Las incrustaciones iniciales incluyen no solo información de la palabra, sino también información de posición, y se representan como \vec{E}
• El objetivo es crear una nueva incrustación \vec{E}' que refleje el contexto a partir de la incrustación existente

• Query

  • En el primer paso, se multiplica cada incrustación de token por la query matrix W_Q para producir el vector query \vec{Q}
  • Puede pensarse como si el sustantivo hiciera la pregunta “¿hay un adjetivo delante de mí?”
  • Los elementos de W_Q son parámetros del modelo que se aprenden, y en la práctica es difícil interpretar exactamente qué hace una cabeza específica
  • Como ejemplo intuitivo, puede verse como un mapeo de la incrustación del sustantivo hacia una dirección que “busca adjetivos en posiciones anteriores”

• Key

  • Al mismo tiempo, cada incrustación se multiplica por la key matrix W_k para crear el vector key \vec{K}
  • La key puede verse como una respuesta potencial a la query, y vive en el mismo espacio reducido de menor dimensión que la query
  • El grado de alineación entre key y query se mide con el producto punto
  • Cuanto mayor sea el producto punto, más fuertemente alineados están los dos vectores
  • Si las key de fluffy y blue encajan bien con la query de creature, tendrán valores positivos grandes
  • Al calcular el producto punto de todos los pares key-query, aparece una cuadrícula de puntajes que indica qué tan relevante es una palabra para actualizar el significado de otra

Attention pattern y softmax

• Como los puntajes del producto punto pueden tomar valores de -\infty a \infty, se aplica softmax a cada columna para normalizarlos a valores entre 0 y 1
• A la cuadrícula normalizada se le llama attention pattern
  • Cada columna puede verse como un conjunto de pesos sobre qué tan relevante es la palabra de la izquierda para actualizar la palabra de arriba
• El artículo original de Transformer lo escribe de una forma más compacta
  • Q y K son los arreglos completos de vectores query y key
  • K^TQ representa la cuadrícula de todos los productos punto posibles entre key y query
  • En la notación del paper, query y key se colocan por filas y se escribe como QK^T, por lo que queda invertido diagonalmente respecto al diagrama explicado aquí
• Se incluye un término que divide por \sqrt{d_k}, la raíz cuadrada de la dimensión del espacio key-query, por estabilidad numérica
• El softmax envuelve toda la expresión, pero conceptualmente se aplica a cada columna

Masking y limitaciones del tamaño de contexto

• Durante el entrenamiento, el modelo no predice solo un siguiente token para un texto dado, sino también los posibles siguientes tokens después de cada subsecuencia al mismo tiempo
  • Así, un solo ejemplo de texto funciona como varios ejemplos de entrenamiento y la eficiencia aumenta
• En el caso de GPT, si los tokens posteriores influyeran en los anteriores, se filtraría la respuesta correcta del siguiente token, así que se usa masking
  • Antes del softmax, los valores de esas posiciones se fijan en infinito negativo
  • Después del softmax, esas posiciones pasan a valer 0, y la columna sigue quedando normalizada
• No toda atención aplica masking siempre, pero en el ejemplo de GPT se usa de forma constante para evitar que los tokens posteriores influyan en los anteriores
• El tamaño del attention pattern es igual al cuadrado del tamaño de contexto
  • Por eso, el tamaño de contexto puede convertirse en una limitación importante en los modelos de lenguaje grandes
  • Han surgido variantes para hacer el mecanismo de atención más escalable y permitir una context window más grande, pero aquí solo se trata la forma básica

Cómo Value actualiza realmente las incrustaciones

• El attention pattern proporciona los pesos sobre qué palabra actualiza a cuál, y el siguiente paso es producir el cambio real en la incrustación
• Cada incrustación se multiplica por la value matrix W_V para crear el vector value
  • El vector value vive en un espacio de alta dimensión igual al de la incrustación
  • Representa qué cambio concreto debe añadirse cuando una palabra relevante ajusta el significado de otra
• En cada columna, se multiplica cada vector value por su peso correspondiente del attention pattern y luego se suman todos para obtener el cambio \Delta \vec{E}
• Si este cambio se suma a la incrustación original, se obtiene la nueva incrustación contextual \vec{E}'
  • En el ejemplo, creature absorbe la información de fluffy y blue y pasa a contener un significado más cercano a “fluffy blue creature”
• Si se aplica el mismo proceso a todas las columnas, del attention block salen incrustaciones refinadas para toda la secuencia de tokens
• Una sola cabeza de atención queda parametrizada por tres tipos de matrices entrenables: key matrix, query matrix y value matrix

Cálculo de parámetros en GPT-3

• En el ejemplo de GPT-3, las matrices key y query tienen cada una 12,288 columnas, correspondientes a la dimensión de incrustación, y 128 filas, correspondientes a la dimensión del espacio key-query
  • Cada matriz tiene 1,572,864 parámetros
• Si la value matrix fuera una matriz cuadrada de 12,288×12,288, añadiría 150,994,944 parámetros, mucho más que key/query
• En la práctica, resulta más eficiente descomponer el value map en dos matrices más pequeñas para dejar el número de parámetros en un rango parecido al de key/query
  • La primera matriz reduce el gran espacio de incrustación a un espacio pequeño, como de 128 dimensiones
  • La segunda matriz vuelve a expandir ese espacio pequeño hacia el espacio de incrustación
  • Desde la perspectiva del álgebra lineal, esto restringe el value map total a una transformación de bajo rango
• En esta explicación, esas dos matrices se llaman Value_\downarrow y Value_\uparrow, aunque no son nombres convencionales
• Sumando las cuatro matrices, una cabeza de atención tiene alrededor de 6.3 millones de parámetros

Self-attention y cross-attention

• La estructura descrita hasta aquí corresponde más precisamente a una self-attention head
• Una cross-attention head aparece en modelos que procesan dos conjuntos de datos diferentes
  • Por ejemplo, en un modelo de traducción, la key puede venir de un idioma y la query de otro
  • El attention pattern puede representar cómo las palabras de un idioma se corresponden con las del otro
• En cross-attention, la diferencia frente a self-attention es que los mapas key y query actúan sobre conjuntos de datos distintos
• En configuraciones como traducción, normalmente no hay masking porque no existe la idea de que los tokens posteriores influyan sobre los anteriores

Multi-headed attention y bloques repetidos

• Un attention block real está compuesto por multi-headed attention, que ejecuta varias cabezas en paralelo
• GPT-3 usa 96 cabezas de atención dentro de cada bloque
  • 96 matrices key/query distintas producen 96 attention patterns distintos
  • Cada cabeza crea su propia secuencia de vectores value con su propia matriz value
  • En cada posición de token, se suman todos los cambios \Delta \vec{E} propuestos por todas las cabezas y se añaden a la incrustación original
• Ejecutar múltiples cabezas en paralelo le da al modelo capacidad para aprender varias formas en que el contexto cambia el significado
• En GPT-3, un bloque de multi-headed attention con 96 cabezas tiene alrededor de 600 millones de parámetros
• En los papers y en las implementaciones reales, las matrices equivalentes a Value_\uparrow de cada cabeza se agrupan en una gran output matrix conectada a todo el bloque de multi-headed attention
  • Cuando normalmente se habla de la value matrix de una cabeza concreta, suele referirse a la primera etapa de proyección, la que aquí se llamó Value_\downarrow

Cómo se acumula el significado en Transformers más profundos

• Los datos dentro de Transformer no pasan por un solo attention block, sino por varios attention blocks y por un multi-layer perceptron
• Incluso después de que la incrustación de una palabra haya absorbido parte del contexto, sigue teniendo oportunidades de recibir la influencia de incrustaciones vecinas ya más refinadas
• A medida que la red se hace más profunda, cada incrustación incorpora más significado de las demás, y gana capacidad para codificar características abstractas de mayor nivel, como emoción, tono o si algo es irónico
• GPT-3 incluye 96 capas, y se describe que los parámetros relacionados con key/query/value suman menos de 58 mil millones en total
• Eso representa cerca de un tercio de los parámetros de toda la red, y la mayoría del resto proviene de los bloques entre capas de atención
• Una gran parte del éxito del mecanismo de atención no está en una conducta específica, sino en su alta capacidad de paralelización, que permite ejecutar muchos cálculos en poco tiempo con GPU
• En deep learning, la lección de que escalar el tamaño puede producir grandes mejoras cualitativas en el rendimiento del modelo hace que una arquitectura paralelizable que permita escalar tenga una gran ventaja

[GN⁺]

Opiniones en Hacker News

• Como alguien que ha trabajado en química cuántica y algo de machine learning, al ver este video me llamó bastante la atención la similitud entre los modelos Transformer y la mecánica cuántica.
  En mecánica cuántica, el estado de todo un sistema físico se codifica como un vector normalizado de muy alta dimensión, es decir, como un rayo en el espacio de Hilbert, y la evolución en el tiempo está a cargo de un operador de traslación temporal que puede verse aproximadamente como una matriz unitaria U = exp(-iHt).
  En el video dicen que la predicción del siguiente token se determina calculando el siguiente vector de embedding consciente del contexto únicamente a partir del último vector de embedding consciente del contexto, así que parece el resultado de aplicar una función de estado lineal a un vector de alta dimensión.
  Se siente parecido a generar offline el Hamiltoniano de todo el sistema a partir de los datos de entrenamiento, luego reparametrizar un subsistema específico, la ventana de contexto, en una base adecuada para ese Hamiltoniano, aplicar un paso de traslación temporal y después volver a la base original.
  Aunque para alguien que ha investigado en un área específica, todos los problemas pueden parecer clavos para el martillo de esa área, así que me da curiosidad si otros también ven esta similitud o si es demasiado forzada.

  • No creo que esta analogía encaje bien. Incluso si olvidamos todos los pasos no lineales anteriores, lo que queda no es más que un sistema dinámico lineal, y no tiene la naturaleza compleja ni la propiedad unitaria, que son características clave de la mecánica cuántica.
  • Me parece que simplemente está describiendo una máquina de estados. Codificar estados como vectores y avanzar pasos con matrices me parece más bien un detalle de implementación.
  • Últimamente he pensado un poco en esto. Si el tiempo no fuera continuo, ¿no se podría modelar la evolución temporal del universo aplicando recursivamente algún operador al estado cuántico del universo?
    Si una aplicación del operador hiciera avanzar el estado del universo en un tiempo de Planck, me pregunto si podríamos observar la diferencia entre un universo así y uno con tiempo continuo.
  • Hace tiempo tuve un pasante con doctorado en matemáticas, y decía que el álgebra lineal de alta dimensión era un campo extremadamente avanzado incluso para los estándares de los años 1900, y que en ciencias de la computación todavía había mucho por descubrir ahí.
    Recién ahora se me ocurre la conexión con lo que pasaba en la física de esa época.
  • Al final, me pregunto si esto significa que los modelos computacionales más sofisticados que hemos creado están empezando a acercarse al algoritmo que define el universo en el que vivimos. Es decir, ¿la simulación vuelve a asomar?

• El video de YouTube de CodeEmporium me resultó más fácil de seguir: https://www.youtube.com/watch?v=Nw_PJdmydZY
  Los transformers son difíciles de explicar con analogías y, de hecho, tampoco hay una buena explicación de por qué funcionan, así que quizá sea mejor simplemente mostrar el mecanismo y dejar la interpretación al espectador.
  Además, el producto punto es más sencillo de explicar como vectores proyectándose unos sobre otros.

  • La explicación es simplemente que una red neuronal es un algoritmo de ajuste estadístico que aprende una distribución de probabilidad condicional, P(next_word|previous_words). Los pesos son el modelo de esa distribución, y los LLM son más bien una innovación de hardware que permitió que las GPU calcularan esto a gran escala sobre datos del orden de terabytes.
    La razón por la que después de “the cat sat on the ...” viene “mat” es que es la palabra que aparece con mayor frecuencia en el conjunto de datos, y la red neuronal es un modelo de esas frecuencias.
    La razón por la que parece saber “London in UK” pero no “London in France” también es que “UK” aparece con mucha más frecuencia en el conjunto de datos.
    El algoritmo en sí no hace nada particularmente interesante aparte de alinear los cálculos con el hardware. El valor viene de la estructura de probabilidad condicional en los datos, y esa estructura es resultado de que las personas organizaron palabras de forma útil para transmitirse información.
  • Desde el punto de vista de un científico de la computación, la interpretación de tabla hash diferenciable me pareció acertada. El paper AIAYN también sugiere esa dirección al usar los nombres query/key/value, aunque no dice explícitamente “hash table”. Tal vez se haya introducido en otro paper.
  • Mi comprensión personal de la atención es que la salida del transformer es una secuencia de nuevos vectores de tokens, y que cada vector de token de salida contiene información de contexto de los vectores de tokens de entrada cercanos.
    Sé que es una explicación incompleta, pero creo que es mejor que nada.

• Hay una visualización convincente que muestra cómo funciona un LLM al procesar una solicitud simple: https://bbycroft.net/llm
  Complementa muy bien la explicación detallada de 3blue1brown.

  • Al visualizarlo así, se siente lo absurdamente grande que es la escala de GPT-3. Ni siquiera puedo imaginar cómo se vería GPT-4 aquí.

• Excelente video. Muestra bien por qué la multiplicación de matrices Q*K es un cuello de botella. Si la secuencia, es decir, la longitud de la ventana de contexto, es S, hay que almacenar en memoria una matriz de tamaño SxS con los resultados de todos los queries contra todas las keys.
  Una idea relativamente nueva para mejorar este cuello de botella es Ring Attention, y este artículo la explica bien: https://learnandburn.ai/p/how-to-build-a-10m-token-context
  Yo edité ese artículo.

  • Con Flash Attention no hace falta crear la matriz (S, S) en absoluto. Como la fórmula tiene la forma softmax(Q @ K^T / sqrt(d)) @ V, la salida final se puede construir por tiles.
    En Unsloth, gracias a Flash Attention, el uso de memoria crece linealmente en lugar de cuadráticamente, el fine-tuning se vuelve 2 veces más rápido, el uso de VRAM baja 80% y la inferencia también se acelera 2 veces. Sin embargo, la cantidad de cómputo sigue siendo O(N^2).
    Para contextos largos, la versión más reciente de Unsloth puede manejar un contexto 4 veces más largo que HF+FA2 con un overhead de +1.9%, llegando hasta un contexto de 228K en H100.
  • En el video también enumeran Ring Attention y varias otras técnicas, pero dicen que quedan fuera del alcance de este video: https://youtu.be/eMlx5fFNoYc?t=784

• El artículo anterior, “But what is a GPT?”, también es realmente bueno: https://www.3blue1brown.com/lessons/gpt

• Gracias a este video me di cuenta de que el mecanismo de atención se parece más a una especie de metafunción que a una función específica.
  Si lo entendí bien, Attention + los pesos aprendidos permiten que el transformer aprenda una función más o menos arbitraria, y esa función incluye un mecanismo de matching como el scaled dot-product.

  • Correcto. El poder de la atención está en explorar el espacio de funciones y dar con la mejor función dentro de las restricciones.
    Por eso creo que la atención lineal difícilmente podrá acercarse a la capacidad de la atención estándar. El término cuadrático que explora todos los pares entrada-salida es una característica esencial.

• Gran parte de que este video fuera fácil de digerir se debió a las animaciones. La forma en que se expanden, contraen y despliegan sincronizadas con lo que se dice está muy bien hecha.

  • Sin duda, eso es algo que él hace mejor que la mayoría. Incluso tiene una biblioteca de animación personalizada que creó para animaciones matemáticas: https://github.com/3b1b/manim

• Trabajo en un campo estrechamente relacionado, y este video entró de inmediato en la documentación de onboarding de nuestro equipo.
  También es importante que buena parte del código de visualización esté en GitHub: https://github.com/3b1b/videos/tree/master/_2024/transformers

  • Interesante; me da curiosidad qué más incluye esa documentación de onboarding.

• Por fin lo entendí. No sé por qué otros videos lo hicieron tan confuso.

  • Es un tema confuso de por sí, y 3b1b es así de bueno.
  • Según mi experiencia, salvo excepciones muy raras como Feynmann, los investigadores suelen ser pésimos para explicar con claridad a otros lo que hacen.
    Me hace pensar que la habilidad para enseñar y la habilidad para investigar quizá sean, en general, capacidades mutuamente excluyentes.
  • Como alguien que quiere hacer mejores videos o contenidos educativos, me da curiosidad. Comparados con 3b1b, ¿en qué aspectos otros videos lo hicieron peor?
  • Grant tiene talento para explicar cosas complejas con mucha claridad. Hay una razón por la que su canal es popular.
  • No sé si era una pregunta retórica, pero es una pregunta interesante. Creo que hay al menos tres razones por las que la mayoría se confunde con los transformers.
    Primero, la terminología estándar no es buena. “attention” apenas resulta intuitivo, “self-attention” es peor, y ni hablar de “key” y “value”.
    Segundo, los artículos clave, como Attention is All You Need y el paper de BERT, no están bien escritos. No intento restarles mérito: un paper influyente con un avance enorme puede explicar mal las cosas, y creo que eso fue lo que pasó.
    Tercero, estas arquitecturas se descubrieron en gran medida probando cosas y viendo qué se pegaba. No hubo primero un proceso de reflexión que llevara a predecir que esta estructura funcionaría bien y luego se verificara con experimentos; fue empírico de principio a fin.
    Por eso no entendemos del todo por qué funciona tan bien, todas las explicaciones se parecen a racionalizaciones a posteriori, y últimamente también hay trabajos que sugieren que, con suficiente ajuste, otras estructuras podrían funcionar casi igual de bien. Es difícil explicar algo que no se entiende por completo.

• Me gustaría saber si hay alguna referencia que explique cómo evolucionó la arquitectura actual. Quisiera ver el recorrido desde una idea central muy simple hasta el famoso paper “all you need”.
  Si no, muchos componentes se sienten como si aparecieran de la nada, con mucho cálculo y poca intuición.
  Jeremy Howard dijo en Twitter que había visto varias versiones de esta idea muchas veces, lo que suena a que era una idea natural. Creo que ver ejemplos de cómo apareció esta idea en otros lugares ayudaría a construir intuición.

  • A grandes rasgos, el flujo fue así. Los primeros enfoques seq-2-seq usaban LSTM: uno codificaba la secuencia de entrada y otro decodificaba la secuencia de salida. Ya de por sí es sorprendente que funcionara codificar oraciones de longitud variable en un vector de tamaño fijo y luego decodificarlas de nuevo en otra secuencia, normalmente de longitud distinta.
    Este enfoque con RNN/LSTM tenía la debilidad de usar una representación de tamaño fijo, y también hacía difícil decidir qué parte de la secuencia de entrada debía usarse al generar una parte específica de la salida. Bahdanau y otros resolvieron esto con una arquitectura que combinaba un mecanismo de atención con una RNN encoder-decoder, permitiendo mirar no solo el estado final sino todos los estados pasados de la RNN.
    Como entrenar RNN era ineficiente, Jakob Uszkoreit quiso encontrar una forma de aprovechar mejor el hardware masivamente paralelo, y observó que el lenguaje no solo es secuencial, sino también jerárquico. Propuso una arquitectura por capas que procesaba en paralelo los tokens de subsecuencias en cada capa, mientras mantenía la atención al estilo Bahdanau para predecir la siguiente capa mediante self-attention, de modo que los tokens pudieran referirse entre sí.
    Las primeras implementaciones funcionaban, pero no eran mejores que otros enfoques de la época, como los convolucionales. Luego, según entiendo, Noam Shazeer desarrolló más esa idea y creó una arquitectura que funcionaba mucho mejor, y tras experimentos para eliminar componentes innecesarios, terminó convirtiéndose en el transformer original. No sé bien quién ideó la forma final de atención basada en keys.
    El transformer original del paper Attention is All You Need seguía los enfoques previos basados en RNN y tenía un encoder y un decoder separados, y también se usó en modelos tempranos como BERT de Google. Pero para modelos de lenguaje eso no es indispensable, así que GPT de OpenAI usó solo la parte del decoder, y hoy casi todos usan más o menos ese enfoque. En un transformer solo con decoder, la oración de entrada entra por la capa inferior, luego se transforma paso a paso al pasar por cada capa y sale por arriba. Al final de la secuencia de entrada se añade un token de fin, y este se transforma en el siguiente token de la secuencia de salida, es decir, en el último token.
  • Karpathy resumió bien la historia de la arquitectura transformer en una clase de Stanford: https://youtu.be/XfpMkf4rD6E?si=MDICNzZ_Mq9uzRo9&t=618