Predecir una secuencia simple con un transformer hecho a mano, sin entrenamiento

• Un experimento en el que se diseñaron manualmente solo los pesos de un transformer solo decodificador similar a GPT-2, sin entrenamiento, para predecir el patrón (aab)*
• Como para predecir "aabaabaabaab..." hay que mirar los dos tokens anteriores, esto permite mostrar mejor el funcionamiento de la atención que un patrón alternado simple
• El modelo se construyó pequeño, con N_CTX=5, N_VOCAB=2, N_EMBED=8, y usa tokenización a=0, b=1 junto con embeddings one-hot
• Un único head de atención reparte 0.5 de atención a cada uno de los dos tokens más recientes, y calcula el siguiente token usando la cancelación de la codificación a=1, b=-1
• En contextos no ambiguos mostró una precisión de 100.0% (27/27), pero requiere unas 4,000 FLOPs para una sola predicción con contexto de 5 tokens, por lo que es mucho menos eficiente que una regla implementada directamente

---

Mini GPT-2 con pesos fijados a mano, sin entrenamiento

• El objetivo es entender de forma intuitiva qué hace realmente cada componente de un transformer y de la atención
• El modelo no se entrena ni usa pesos preentrenados; fue construido asignando manualmente cada peso en una sola tarde
• La estructura es un transformer solo decodificador similar a GPT-2, y la implementación se simplificó a partir de la implementación picoGPT de jaymody
  • se eliminó layer norm
  • se usó un solo head en lugar de multi-head attention
  • se eliminó la capa feed-forward mlp del bloque transformer

Por qué se eligió la secuencia (aab)*

• Al principio se intentó predecir una secuencia como "ababababab", pero como el transformer predice una secuencia desplazada, la tarea resultó demasiado fácil
  • basta con predecir b si es a, y si no, a, así que no hace falta usar embeddings posicionales
• La tarea final quedó definida como la predicción de la secuencia "aabaabaabaab...", es decir, (aab)*
  • si los dos tokens anteriores son ab o ba, el siguiente token es a
  • si los dos tokens anteriores son aa, el siguiente token es b
  • bb se trata como un caso fuera del alcance de la tarea
• La tokenización es un esquema simple que solo maneja dos símbolos
  • a es 0
  • b es 1

Dimensiones del modelo y flujo de cálculo

• Los parámetros elegidos del modelo son tres
  • N_CTX = 5: longitud máxima de contexto que el modelo ve de una vez
  • N_VOCAB = 2: los dos tokens a, b
  • N_EMBED = 8: tamaño del embedding que contiene espacio para tokens, posiciones y cálculo
• Aunque la tarea real solo necesita los dos tokens anteriores, se fijó N_CTX=5 para incluir la situación en la que hay que ignorar tokens irrelevantes
• La función gpt funciona en este orden
  • suma embeddings de token y de posición con wte[inputs] + wpe[range(len(inputs))]
  • pasa por un bloque transformer
  • al final crea logits en el espacio de vocabulario con x @ wte.T

Embeddings que guardan posición y token como one-hot

• wpe es el embedding posicional, y usa las primeras 5 dimensiones del embedding como one-hot posicional
  • la posición 0 es [1, 0, 0, 0, 0]
  • la posición 4 es [0, 0, 0, 0, 1]
• wte es el embedding de tokens, y usa las siguientes 2 dimensiones como one-hot de token
  • el token a es [1, 0] en esas dimensiones de token
  • el token b es [0, 1]
• La octava posición del embedding no se usa al principio y se reserva como scratch space dentro del bloque transformer
• Por ejemplo, "aabaa" se representa como una matriz de embeddings 5 x 8, y cada fila contiene tanto el one-hot posicional como el one-hot de token

Cómo el head de atención selecciona los dos tokens más recientes

• El bloque transformer está compuesto por un solo head de atención y una capa lineal c_proj que devuelve el resultado de la atención al espacio de embeddings
• c_attn es una capa lineal de tamaño embed_size x (embed_size * 3) que convierte el embedding de entrada en una matriz qkv, y luego la divide en q, k, v
• k separa el embedding posicional y representa la información de posición de cada token
• q representa el rango de posiciones que cada posición quiere buscar, y mediante q @ k.T forma la matriz de scores de atención
• Después del softmax y la máscara causal, la matriz de atención tiene estas propiedades
  • la primera fila pone 100% de atención solo en el primer token
  • las filas siguientes ponen 0.5 de atención en cada uno de los dos tokens accesibles más recientes
• La máscara causal impide ver tokens futuros sumando un valor muy pequeño a esas posiciones, como -1e10 en el código real
  • este modelo hecho a mano no está diseñado para espiar el futuro, pero la máscara se mantuvo para conservar una estructura más cercana a GPT-2
• La escala al dividir por np.sqrt(q.shape[-1]) ayuda a mejorar los gradientes en entrenamiento real, pero no afecta a este modelo artesanal

Codificación de v y predicción mediante suma con cancelación

• v convierte el one-hot del token a la codificación a=1, b=-1
• Como el resultado de la atención promedia con peso 0.5 los dos tokens más recientes, esta codificación calcula la siguiente regla
  • a, b → 0.5 * 1 + 0.5 * (-1) = 0
  • b, a → 0.5 * (-1) + 0.5 * 1 = 0
  • a, a → 0.5 * 1 + 0.5 * 1 = 1
• Como resultado, en la séptima posición de la fila aparece el siguiente valor
  • 0 cuando se debe predecir a
  • 1 cuando se debe predecir b
• En la entrada "aabaa", la primera fila puede producir una predicción b por falta de información, pero las predicciones posteriores sí coinciden con la regla de (aab)*

Cómo se devuelve la predicción al espacio de vocabulario

• c_proj convierte el valor de la séptima posición del resultado de atención de vuelta al formato one-hot de token
• No genera simplemente [..., 1, 0, ...] o [..., 0, 1, ...], sino un one-hot escalado por 1024
  • embedding[row, 5] = 1024 + (-1024) * prediction
  • embedding[row, 6] = 0 + 1024 * prediction
• El bloque transformer tiene una conexión residual, así que x = x + causal_self_attention(...) suma el embedding original
• Como esa señal residual queda ahí aunque no haga falta, se usa la escala 1024 para dominarla
• Por último se calcula x @ wte.T para obtener los logits y luego se aplica softmax
  • en el contexto "aabaa", la fila final de predicción apunta a b
  • durante el entrenamiento sirven las predicciones de todas las filas, pero en inferencia solo hace falta la última

Resultados de generación y precisión

• La función complete mete en el modelo los últimos 5 tokens como máximo y elige el siguiente token con argmax sobre la última fila del resultado de softmax
• Ejemplos de generación
  • complete("a") → a :: baabaabaab
  • complete("ba") → ba :: abaabaabaa
  • complete("abaab") → abaab :: aabaabaaba
• Incluso con entradas fuera del alcance, a veces recupera el patrón repetitivo
  • complete("ababa") → ababa :: abaabaabaa
  • complete("bbbbb") → bbbbb :: aabaabaaba
• En la prueba "aab" * 10, si se evalúan solo contextos no ambiguos, la precisión es de 100.0% (27/27)

La diferencia entre 4,000 FLOPs y 8 instrucciones

• Al usar todo el contexto de 5 tokens, este modelo necesita alrededor de 4,000 operaciones de punto flotante para predecir un solo token
  • la mayoría se usan en el cálculo de la atención
  • se puede reducir con una ventana de contexto más pequeña, fused multiply-add, kv caching, etc.
  • aun así, para predecir un solo token siguen haciendo falta cientos de instrucciones de máquina
• Un ensamblador x64 escrito directamente para la misma regla (aab)* calcula el siguiente token en 8 instrucciones
• Queda abierta la pregunta de si se puede entrenar un modelo de lenguaje 1000 veces más eficiente que los modelos actuales para generación de lenguaje natural

[GN⁺]

Opiniones en Hacker News

• Un trabajo relacionado es "Thinking Like Transformers".
  Presenta un lenguaje de programación primitivo llamado RASP, compuesto por operaciones que pueden modelarse con componentes de Transformer, y muestra que se pueden escribir programas como histogramas u ordenamientos.
  También hay una excelente entrada de blog de Sasha Rush y Gail Weiss, y trabajos posteriores mostraron que los programas tipo RASP pueden compilarse en pesos de modelo reales sin entrenamiento.
  [1] https://arxiv.org/abs/2106.06981
  [2] https://srush.github.io/raspy/
  [3] https://arxiv.org/abs/2301.05062

  • Me gustan mucho los enfoques de la familia RASP.
    Si este campo te parece divertido, también vale la pena echarle un vistazo a mi trabajo HandCrafted Transformers, donde elegí manualmente los pesos de un modelo Transformer para que hiciera sumas largas de forma parecida a como las personas aprenden en la primaria.
    [1] https://colab.research.google.com/github/newhouseb/handcraft...
  • Creo que para este tipo de trabajo encajaría bien un lenguaje funcional como Haskell.
    La dirección de ir de redes neuronales a código también parece muy interesante desde el punto de vista de la explicabilidad.

• Pensaba que entendía bien los Transformers, pero nunca había implementado uno por mi cuenta.
  Un día probé implementarlo yo mismo y no funcionaba ni entrenaba tan bien como el Transformer estándar de PyTorch; al final me di cuenta de que la causa era que había ignorado el dropout.
  Lo entrené para sumar números, y como nunca le mostraba el mismo par dos veces, pensaba que el sobreajuste era imposible, pero el papel del dropout fue mucho más grande de lo que esperaba.
  En resumen, conviene simplemente implementar un Transformer uno mismo, y cuanto más desde cero, mejor.
  Todos los que lo hicieron aprendieron algo inesperado, y lo que entendió cada persona fue distinto: desde la paralelización del entrenamiento a nivel de tokens hasta cómo funciona realmente la retropropagación.

  • Me pregunto si habrá referencias que ayuden a abordar este trabajo.

• También me gustan los materiales de Karpathy, pero el video que finalmente hizo que entendiera los Transformers fue este: https://youtu.be/kWLed8o5M2Y?si=SJT5_lCJ0hSR7Z_k

• He tenido una idea similar desde hace un tiempo.
  ¿No se podría crear una interfaz intuitiva sobre los pesos del modelo que permita a un experto de dominio ajustarlos manualmente para acelerar el entrenamiento?
  Por ejemplo, en un modelo de visión, al detectar conos de tráfico, subir un conjunto de pesos correspondiente a "lo anaranjado".
  Así una persona podría acelerarlo, en vez de exigir miles o millones de ejemplos adicionales para calibrar correctamente "lo anaranjado".
  Claro que la dificultad está en que esta interfaz tendría que mapearse a conjuntos de pesos con significados diferentes, y me pregunto si hay alguna razón técnica por la que eso sea imposible.

  • Eso de que "un experto de dominio ajuste manualmente los pesos del modelo" suena parecido a cómo funcionaba el reconocimiento de imágenes antes del deep learning.
    [1] https://www.youtube.com/watch?v=8SF_h3xF3cE&t=1358s
  • La razón que buscas se llama The Bitter Lesson.
    En pocas palabras, las formas de ayuda humana en IA casi siempre son menos rentables que simplemente correr más poder de cómputo.
    Mientras una persona calibra una capa de pesos para detectar conos de tráfico anaranjados, un clúster de GPU ya habrá entrenado a la IA para detectar conos de tráfico, semáforos, árboles, otros autos e incluso conos de tráfico de un naranja ligeramente distinto.
  • La cantidad de capas y pesos no está en una escala que una persona pueda actualizar manualmente, e incluso si fuera posible, sería demasiado difícil gestionar los efectos aguas abajo de cambiar los pesos.
    Aunque ajustes la imagen para que vea mejor el naranja, si no puedes monitorear al mismo tiempo la precisión para todos los demás colores, es probable que sin darte cuenta generes problemas con otros colores.
  • La razón por la que es técnicamente imposible, o muy difícil, es que los pesos normalmente son muy difíciles de interpretar.
    No es que un clúster específico de neuronas corresponda a un concepto específico; más bien, en general todo hace un poco de todo.
  • El mecanismo de atención de los Transformers no parece mapearse fácilmente a una semántica comprensible para humanos.
    Hay demasiados parámetros involucrados.

• El paper de los Transformers era tan técnico que siempre quise entenderlo aunque fuera superficialmente, pero me resultaba difícil.
  Este artículo me ayudó muchísimo a entender cómo funciona, y al menos los ejemplos eran muy claros.
  Gracias a eso también pude recordar las matrices que había aprendido en la universidad.

• ¿No es una especie de máquina abstracta, como una máquina de Turing o una máquina que parsea expresiones regulares?

  • Simplificando un poco, es una "máquina" que mapea un conjunto de entradas a un conjunto de probabilidades para la siguiente salida.
    Primero se define una lista de tokens; por ejemplo, digamos que son 24 caracteres para hacerlo fácil.
    Esta máquina toma una secuencia de tokens de entrada, realiza operaciones matriciales deterministas y luego devuelve una lista de probabilidades para todos los tokens.
    El "entrenamiento" no es más que el proceso de fijar algunos de los números dentro de las matrices usadas en esas operaciones.
    Vale la pena notar que en el código final hay un solo if, y ese es para evaluar la precisión del resultado.
    Toda la "lógica" sale del resultado de las operaciones matriciales.
  • Es bastante difícil interpretar estas cosas como autómatas en el sentido en que normalmente los imaginamos.
    En las redes neuronales, en general todo es un poco difuso y casi no hay cosas como if/else, aunque sí hay casos, como en los ejemplos de Transformers, donde los valores se "enmascaran" con 0 o -∞.
    La salida también casi siempre es un conjunto de puntajes o probabilidades; si un modelo que distingue fotos de gatos y perros devuelve algo como dog:0.95 cat:0.05, decimos que predijo perro porque el puntaje de perro es más alto.
    El mecanismo de atención, que es el núcleo de Transformer, se basa en una especie de operación de búsqueda suave.
    En un sistema no difuso, se recorrería cada token de la secuencia, se verificaría si está relacionado con el token actual y, si lo está, se haría alguna acción; en Transformer, en cambio, la relación no es una decisión binaria.
    En su lugar, calcula puntajes continuos de relación entre todos los pares de tokens de la secuencia y usa esos puntajes para decidir la siguiente acción.
    Pero algunas cosas no son fáciles de generalizar directamente desde un sistema basado en decisiones binarias.
    Por ejemplo, esos puntajes de relación se usan como pesos para calcular un promedio ponderado sobre los tokens del vocabulario y obtener un "token promedio" para la posición actual.
    No parece haber una forma sencilla de interpretar esto como una extensión de un proceso basado en lógica de ramificación.
  • ¿No es como si AllSpark hubiera tocado un montón de álgebra lineal?
  • Sí.
    Vale la pena ver este paper, que explica que los Linear Transformers en realidad son Fast Weight Programmers: https://arxiv.org/abs/2102.11174
  • Las redes neuronales son máquinas de Turing.
    Si se configuran los pesos con cuidado, se puede hacer que realicen cualquier cálculo.
    Aunque estaría bueno tener un compilador que no se base en aproximaciones.

• Me pregunto para qué serviría eso de "quizá te den ganas de crear tu propio modelo", aparte de como ejercicio de aprendizaje para satisfacer la curiosidad.
  Empieza a dar la impresión de que los modelos de machine learning complejos son poco realistas para alguien que lee blogs en casa.

  • En nanoGPT, si preentrenas un modelo con Shakespeare, en 3 minutos alcanza sobre el material original un nivel de fidelidad comparable al Jabberwocky de Lewis Carroll.
    Genera muchas palabras de inglés antiguo que parecen plausibles, y aprende las bases de la gramática inglesa y el formato de las obras de teatro, entre otras cosas.
    Me sorprendió bastante que llegara tan lejos en tan poco tiempo.
    Entrenar varios modelos localmente hasta un nivel de fidelidad tipo Shakespeare-from-Wish.com podría servir para evaluar si encontraste una buena arquitectura y si ya es momento de intentar escalarla.
  • El objetivo aparece en el primer párrafo del artículo.
    Dice que quería entender mejor Transformer y la atención, y que aunque había leído The Illustrated Transformer, no le quedaba intuitivamente claro qué hacía en la práctica cada una de las piezas de la atención.
    Cosas como la diferencia entre q y k, y ni hablar de v.
  • Es un excelente ejercicio de aprendizaje.
    Más allá de satisfacer la curiosidad, ayuda a construir y profundizar la comprensión.
  • ¿No será que simplemente disfrutan hackear este tipo de proyectos? Es raro, pero puede ser.

• Estaría bueno que el título pudiera incluir alguna expresión como neural network.
  Esto trata sobre la arquitectura "Transformer" de machine learning, no sobre un conjunto de bobinas que acoplan electromagnéticamente dos circuitos.