Más allá de la autoatención: cómo los modelos de lenguaje pequeños predicen el siguiente token
(shyam.blog)- Se entrenó un Transformer decoder-only de unos 10 millones de parámetros con TinyShakespeare y luego se rastreó, mediante estados internos y experimentos, cómo el resultado del cálculo después de la autoatención se convierte en probabilidades del siguiente token
- La hipótesis central es que cada bloque Transformer conecta el prompt con un conjunto específico de cadenas del corpus de entrenamiento, y que la distribución de frecuencias de los tokens que realmente aparecieron después de esas cadenas se aproxima a la predicción del siguiente token del bloque
- El procedimiento de aproximación implementado busca cadenas de entrenamiento similares a la salida de la feed-forward network de cada bloque, suma la distribución de caracteres que las siguieron y reconstruye la salida del modelo mediante una suma ponderada por bloque
- En 20,000 prompts de longitud 10, la Hellinger distance promedio entre la salida aproximada y la salida real del Transformer fue de alrededor de 0.17, mientras que la distancia entre el modelo original y modelos alternativos con la misma estructura entrenados con otros seeds fue de alrededor de 0.11 ± 0.08
- El análisis respalda la interpretación de que la salida feed-forward tiene un peso importante en la transformación de los bloques y que la attention mapea el prompt a la clase correcta de cadenas de entrenamiento, pero no se sabe si esto generaliza más allá de modelos pequeños
Modelo experimental y pregunta
- Se implementó y entrenó un Transformer de unos 10 millones de parámetros siguiendo Let’s build GPT: from scratch, in code, spelled out, de Andrej Karpathy
- El código central del modelo de lenguaje es obra de Karpathy; el análisis interno y el código auxiliar son contribuciones separadas
- El código relacionado está en transformer-experiments
- El modelo es un Transformer decoder-only compuesto por 6 bloques
- Los datos de entrenamiento son TinyShakespeare, que contiene 40,000 líneas de obras de Shakespeare
- Después de entrenarse durante aproximadamente 1 hora en una GPU RTX 4000, genera texto falso de Shakespeare estructuralmente verosímil
- Los tokens no son palabras, sino caracteres
- El tamaño del vocabulario es de 65, y el modelo predice la distribución de probabilidades del carácter que sigue al prompt de entrada
- Por ejemplo, después de
'my most gr'prediceacon 0.819,econ 0.081 eicon 0.059
- La pregunta central no es la autoatención en sí, sino cómo el resultado después de la attention se convierte en la probabilidad del siguiente token final
Estructura de los bloques y peso de la feed-forward network
- Cada bloque Transformer incluye multi-head self-attention y una feed-forward network
- En la estructura de PyTorch, la salida del bloque tiene la forma de sumar
x + self.sa(self.ln1(x))y luegox + self.ffwd(self.ln2(x)) - Si se expresa con variables intermedias, la salida del bloque puede verse como
x + sa_out + ffwd_out
- En la estructura de PyTorch, la salida del bloque tiene la forma de sumar
- En el modelo analizado, la feed-forward network representa el 65.71% de todos los parámetros entrenables
- Parámetros feed-forward: 7,089,408
- Parámetros entrenables totales: 10,788,929
- La salida de la feed-forward network parece ser el factor principal que convierte la entrada del bloque en su salida
- En experimentos posteriores, la salida feed-forward también mostró una influencia mayor sobre la distribución de probabilidades final que la salida de self-attention por sí sola
Buscar cadenas mediante salidas feed-forward similares
- El procedimiento de aproximación usa la similitud de las salidas feed-forward entre el prompt y las cadenas del corpus de entrenamiento
- Se introduce el prompt en el modelo y se guarda la salida de la feed-forward network de cada bloque
- Para cada bloque, se buscan en el corpus de entrenamiento cadenas cuya salida feed-forward sea similar
- Se suma la distribución de frecuencias de los tokens que aparecieron después de esas cadenas
- Se combinan las distribuciones por bloque con una suma ponderada y se normalizan para crear una distribución de probabilidades
- En el prompt de demo
'And only l', se encontraron 94 cadenas de longitud 10 con cosine similarity de 0.95 o más respecto de la salida feed-forward del bloque 1- Las cadenas encontradas, como
hat only l,\nMy only l,I dearly l,ng Henry lyAnd only l, en general terminaban eny loly l - La distribución de frecuencias normalizada de los caracteres que siguieron a esas cadenas fue muy similar a la predicción real del siguiente carácter del modelo
- Las cadenas encontradas, como
- En el ejemplo
'And only l', la aproximación y la predicción real del modelo coinciden bastante de cerca- Modelo:
i0.437,o0.204,a0.195,e0.160 - Aproximación:
i0.389,o0.250,a0.222,e0.139 - La Hellinger distance es 0.0711
- Modelo:
- También hay casos en los que un solo bloque falla
- Con el prompt
'hing tremb', si se usa solo el bloque 1, el modelo predicelcon 0.999, pero la aproximación dae0.543 yl0.343, con una Hellinger distance de 0.6305 - Al reflejar todos los bloques, en el mismo ejemplo la aproximación predice
lcon 0.997 y la Hellinger distance baja a 0.0233
- Con el prompt
Escala de cálculo y ajuste manual
- Por eficiencia, todo el análisis se realizó principalmente con cadenas de longitud 10
- El corpus de entrenamiento tiene 1,115,394 caracteres y 858,923 subcadenas únicas de longitud 10
- La salida feed-forward es un vector
float32de 384 dimensiones, y hay una por cada uno de los 6 bloques
- Se eligieron al azar 20,000 cadenas de longitud 10 para usarlas como prompts experimentales
- Para cada prompt y las 858,923 subcadenas únicas, se calculó previamente la cosine similarity por bloque y se guardó en disco
- Como los matches de interés no aparecían por debajo de cosine similarity 0.7, solo se prefiltraron los valores de 0.7 o más
- El similarity threshold por bloque se definió mediante ajuste manual
- Bloque 0: 0.95
- Bloque 1: 0.94
- Bloque 2: 0.85
- Bloque 3: 0.76
- Bloque 4: 0.81
- Bloque 5: 0.89
- También se usaron pesos manuales al combinar las distribuciones de frecuencias por bloque
- Los pesos usados fueron
[0.01, 0.01, 0.1, 1.5, 6, 0.01] - También se intentó buscar pesos con un enfoque de deep learning, pero no se obtuvieron mejores resultados que con el ajuste manual
- Los pesos usados fueron
Evaluación con 20,000 prompts
- Se calculó la Hellinger distance entre la distribución aproximada y la distribución de salida del modelo para los 20,000 prompts
- Promedio: 0.1677
- Desviación estándar: 0.1215
- Mínimo: 0.0013
- Máximo: 0.9994
- La Hellinger distance mide el solapamiento entre distribuciones de probabilidad: 0 significa que son idénticas y 1 que no se solapan
- Para comparar la calidad de la aproximación, se entrenó 3 veces más la misma arquitectura con distintos random seeds
- Seed del modelo original: 1337, training loss estimado 0.9334, validation loss 1.5063
- Alternate 1 seed: 1442, training loss 0.9293, validation loss 1.5038
- Alternate 2 seed: 88, training loss 0.9294, validation loss 1.4991
- Alternate 3 seed: 99999, training loss 0.9339, validation loss 1.4941
- La Hellinger distance promedio entre el modelo original y los modelos alternativos es de alrededor de 0.11 ± 0.08
- Original vs Alternate 1: 0.1064 ± 0.0823
- Original vs Alternate 2: 0.1057 ± 0.0817
- Original vs Alternate 3: 0.1053 ± 0.0828
- La distancia promedio de 0.17 entre la aproximación y el modelo es mayor que la distancia de 0.11 entre modelos alternativos, pero está dentro del rango de la desviación estándar, lo que sirve como base para considerar que la aproximación es bastante buena
Interpretación del funcionamiento interno del Transformer
- Después del embedding, el modelo puede verse como una serie de transformaciones espaciales
- Los 6 bloques Transformer transforman el embedding de entrada en un embedding de salida dentro de un espacio de embeddings de 384 dimensiones
- La LayerNorm final y la linear layer convierten el espacio de embeddings en un espacio de logits de 65 dimensiones, y el softmax produce las probabilidades del siguiente token
- La transformación interna de los bloques puede verse como suma de vectores
- La salida del bloque es la suma de la entrada
x, la salida de self-attentionsa_outy la salida feed-forwardffwd_out - En el análisis se observó que los vectores de salida feed-forward en general tienen mayor norm que la salida de self-attention, y que ambas salidas suelen apuntar en direcciones parecidas según el bloque
- La salida del bloque es la suma de la entrada
- Incluso si solo se suma la salida feed-forward, el vector de salida final suele tener una dirección parecida
- La salida que usa solo feed-forward tiene menor norm, pero su dirección es cercana a la salida original
- Debido a la LayerNorm final, la diferencia de norm no es muy importante en la entrada de la linear layer final
- Esto no significa que el cálculo de self-attention pueda eliminarse
- La feed-forward network incluye la salida de self-attention en su entrada, como en
ffwd_out = self.ffwd(self.ln2(x + sa_out)), por lo que si se elimina self-attention, la propia salida feed-forward cambia
- La feed-forward network incluye la salida de self-attention en su entrada, como en
Hipótesis del subespacio de tokens
- Un embedding específico puede hacer que el modelo prediga un token determinado casi con certeza
- Por ejemplo, se puede aprender un embedding de entrada para la LayerNorm y la linear layer después del bloque final, de modo que la probabilidad del token
asea casi 1 - Los pesos del Transformer se mantienen fijos y solo se optimiza el tensor de embedding
- Por ejemplo, se puede aprender un embedding de entrada para la LayerNorm y la linear layer después del bloque final, de modo que la probabilidad del token
- Para cada token no existe un único embedding, sino múltiples embeddings
- Esto se interpreta como un subespacio de embeddings no lineal complejo correspondiente a cada token
- Fue posible aprender fácilmente miles de embeddings únicos para cada token
- Aunque no se pudo obtener el subespacio exacto de forma matemática, en la etapa final una aproximación lineal funciona bien
- Se apilan los embeddings aprendidos para cada token y se realiza SVD
- El primer right singular vector funciona bien como aproximación lineal unidimensional del subespacio de ese token
- Al combinar linealmente vectores aproximados de subespacios de varios tokens, se puede crear un embedding que distribuya la probabilidad entre varios tokens
- Si se suman los vectores aproximados de
ayb, se obtiene una distribución donde la probabilidad se reparte entre esos dos tokens - Sin embargo, por el error de aproximación, la no ortogonalidad de los vectores de subespacios y las diferencias en la magnitud de los logits por token, no se obtiene una distribución 50:50 perfecta
- Si se suman los vectores aproximados de
Correspondencia entre la salida feed-forward y los subespacios de tokens
- La salida feed-forward se interpreta como un ajuste de la salida del bloque hacia el subespacio de un token específico
- Esos tokens coinciden con los tokens predichos por el procedimiento de aproximación, es decir, los tokens que aparecieron después de las cadenas de entrenamiento que producen salidas feed-forward similares
- En el ejemplo
med me Aut, si se mira solo el último bloque, la aproximación prediceocomo siguiente token más probable yhcomo el segundo- Al proyectar el mismo vector de salida feed-forward sobre la aproximación del subespacio de tokens, los tokens más similares también fueron
o,h,i,u,y, en ese orden
- Al proyectar el mismo vector de salida feed-forward sobre la aproximación del subespacio de tokens, los tokens más similares también fueron
- En el ejemplo
if and thyaparece una correspondencia similar- La aproximación toma space,
sy newline como principales tokens predichos - Los subespacios más alineados con la salida feed-forward también se acercaron al orden space,
sy newline
- La aproximación toma space,
- Se estableció un criterio arbitrario para agregarlo sobre los 20,000 prompts
- El criterio es: “¿los subespacios de los tokens que acumulan el 90% de la masa de probabilidad en la predicción aproximada entran en la mitad superior del ranking de cosine similarity con la salida feed-forward?”
- Tomando el bloque final como referencia, 16,357 casos, es decir 81.78%, cumplieron el criterio
Resultados agregados por bloque y comparación con azar
- Al usar la aproximación de subespacio aprendida en cada posición de bloque, la tasa de cumplimiento del criterio es mayor en los bloques posteriores
- Bloque 6: 16,357 casos, 81.78%
- Bloque 5: 10,142 casos, 50.71%
- Bloque 4: 7,760 casos, 38.80%
- Al aplicar la aproximación de subespacio de la etapa final a todos los bloques, los resultados mejoran
- Bloque 6: 81.78%
- Bloque 5: 68.26%
- Bloque 4: 58.15%
- Bloque 3: 57.34%
- Bloque 2: 52.02%
- Bloque 1: 49.71%
- También se simuló la probabilidad de que el criterio se cumpliera por casualidad con cosine similarity aleatoria
- Bloque 6: 20.76% ± 0.25%
- Bloque 5: 20.55% ± 0.26%
- Bloque 4: 18.37% ± 0.24%
- Bloque 3: 18.20% ± 0.24%
- Bloque 2: 17.04% ± 0.23%
- Bloque 1: 16.31% ± 0.23%
- Los resultados están muy por encima del azar, pero por el ruido de medición y las limitaciones de la aproximación de subespacios, es difícil considerarlos prueba definitiva
Rol de self-attention
- En esta interpretación, una buena predicción depende de mapear el prompt a la clase correcta de cadenas del corpus de entrenamiento
- El eje encargado de ese mapeo es self-attention
- Las attention layers identifican patrones entre los tokens del prompt
- El patrón puede ser algo simple como la cadena final
y l, o una clase de token más general, como una vocal o una mayúscula en una posición específica
- Los pesos aprendidos por cada attention head determinan a qué patrones responde
- Cuando la salida de una attention head pasa por la feed-forward network, se convierte en una representación en el espacio de embeddings que contiene información sobre la distribución de tokens aparecidos después de cadenas de entrenamiento similares
- El modelo analizado tiene 6 bloques y 6 attention heads por bloque, por lo que puede evaluar un prompt frente a múltiples patrones latentes
Conclusión y límites
- La evidencia de que el procedimiento de aproximación y la salida real del Transformer son similares es fuerte
- La evidencia de que el procedimiento de aproximación corresponde al cálculo interno real del Transformer es menos clara, pero sugiere que al menos en parte podría ser cierto
- No se sabe si los resultados obtenidos en un Transformer pequeño se generalizan a modelos más grandes o a otros datasets
- Este proyecto es un proceso de hacer preguntas sobre el interior del modelo y diseñar experimentos; más que reducir la “magia” de los modelos de lenguaje a un mecanismo simple, ayuda a ver mejor su complejidad
1 comentarios
Opiniones de Hacker News
Varios de los temas del post padre no deberían sorprender demasiado a quien haya leído https://people.math.harvard.edu/~ctm/home/text/others/shanno...
Si no lees la literatura fundamental de tu propio campo, hasta los fenómenos que surgen como consecuencia natural de trabajos ya establecidos terminan pareciendo misterios inexplicables
Aun así, a primera vista los experimentos se ven bastante exhaustivos, y valoro mucho la cantidad de trabajo detallado que hubo detrás
Hay un equilibrio difícil entre aprender la teoría existente y derivarlo todo de nuevo desde cero. Sin una base tradicional se pueden hacer nuevos descubrimientos, pero con esa base se puede entender más a fondo un fenómeno específico
En los comentarios aquí, se ve a varias personas sorprendidas de que un modelo que maximiza la log-verosimilitud de una secuencia dados los datos no se desvíe mágicamente de ese comportamiento durante la inferencia. Es un modelo de estimación de densidad; ¿esperaban que recitara Shakespeare de la nada?
Volver a lo básico hace que estos experimentos se entiendan mucho mejor. Ya existe una base matemática muy clara que explica esto, y también los llamados fenómenos emergentes
Más concretamente, hay varios niveles, y el tratamiento de Shannon de los sistemas ergódicos es un buen punto de partida. Aquí hay partes que no encajan del todo, pero parece una correspondencia lo bastante cercana para entender la dinámica general
No es raro que personas inteligentes digan “esta estructura matemática se parece a aquella idea; si agregamos o quitamos unas cuantas estructuras más, todo queda explicado”. Pero en realidad hay mucho que no sabemos
No he visto que los teóricos de este campo lleguen con una teoría y construyan algo nuevo o hagan predicciones útiles. Normalmente prueban de todo, y si funciona luego le ponen una explicación plausible de por qué funcionó; si no funciona, lo entierran
Hace poco también hubo un post que veía los transformers como suavizado por kernel: https://arxiv.org/abs/1908.11775
El redescubrimiento repetido no necesariamente es un desperdicio; también puede ser un proceso de confirmación y verificación de una verdad profunda a la que se puede llegar por varios caminos
Muchos libros de texto de teoría de la información ya señalan la naturaleza direccionable por contenido (content-addressable) de estas redes[1], y por eso también se usan en aplicaciones como la compresión[2][3]
Así que tampoco sorprende que, cuando el NYT le puso como prompt unos párrafos de sus propios artículos al modelo de OpenAI, este los reprodujera casi tal cual
[1] https://www.inference.org.uk/itprnn/book.pdf
[2] https://bellard.org/nncp/
[3] https://pub.towardsai.net/stable-diffusion-based-image-compr...
Cuando vi que Google señaló que se podía hacer que ChatGPT escupiera sus datos de entrenamiento tal cual si se le pedía repetir la misma palabra una y otra vez[0][1], se me ocurrió exactamente la misma idea. Me alegra que alguien la haya implementado de verdad
Aquí me surgen dos preguntas adicionales
[0] Antes era posible. OpenAI ahora bloquea la generación si llenas la ventana de contexto con una sola palabra
[1] https://arxiv.org/abs/2311.17035
No se saltan muchos pasos; al contrario, se agrega bastante trabajo extra
Creo que se obtendrían resultados casi iguales incluso si se separaran por completo el corpus de entrenamiento del modelo y el corpus usado para buscar cadenas con activaciones similares. La parte difícil es, para empezar, lograr que cadenas con estadísticas de siguiente token similares produzcan activaciones similares
En los pesos por capa [0.01, 0.01, 0.1, 1.5, 6, 0.01], lo más importante es la penúltima capa, y para entonces la entrada ya está muy transformada. Por eso no hay que esperar que esto permita reemplazar un transformer por un simple grep sobre los datos de entrenamiento
La razón por la que el peso de la penúltima capa es mucho mayor que el de la capa final probablemente sean los induction heads. Como en https://transformer-circuits.pub/2021/framework/index.html, podría tratarse de una estructura que implementa la función de copiar cadenas repetidas de la entrada, donde la penúltima capa decide qué buscar y la capa final realiza la copia
La contribución de este artículo, como dice el autor, está en mostrarle ese punto al lector técnico que construye GPT directamente, a diferencia de los artículos centrados en “¿cómo se implementa?” dentro del transformer en sí
Si uno miraba con detalle, eran README de GitHub que no existían y no tenían coherencia, folletos informativos vacíos, conversaciones aleatorias, cosas así
Pero jamás llaman ataque a raspar todo Internet para entrenar estos modelos
Entonces, ¿por qué no deberíamos considerar que también se tienen derechos sobre la distribución de probabilidad de caracteres dentro de esa obra?
Apenas conocí NanoGPT de Andrej Karpathy, lo entrené con la versión rusa de War and Peace, y me pareció interesante que, aun siendo un modelo de apenas 3 MB, captara bastante bien la gramática rusa
El ruso tiene una estructura sintética y flexiva compleja. Por ejemplo, después de la preposición “na” (“upon”) debe venir un sustantivo en acusativo, lo que se manifiesta como -a en sustantivos masculinos animados, sin terminación en sustantivos inanimados, -ia en sustantivos que terminan en “consonante suave”, -u en sustantivos femeninos, etc.
Además, el verbo “usar” exige el caso instrumental cuando va seguido de un sustantivo que funciona como herramienta
No era perfecto y cometía errores, pero me pareció interesante que NanoGPT infiriera ciertas reglas complejas con solo 3 minutos de entrenamiento. Busqué en el texto original las frases exactas que generó, pero no encontré ninguna coincidencia literal
Eso sí, aunque entendía la gramática hasta cierto punto, semánticamente era un disparate total
Había una buena visualización 3D que muestra el mismo sistema, y creo que leerlo junto con eso puede ser muy útil
LLM Visualization (https://bbycroft.net/llm)
https://news.ycombinator.com/item?id=38505211
Creo que es porque intentó mostrarlo todo a la vez y no relegó nada a conceptos abstractos, aunque no estoy seguro
Es un buen proyecto, pero el modelo analizado es, en la práctica, un modelo de juguete tanto por tamaño como por cantidad de datos de entrenamiento
Por eso este modelo probablemente pueda aproximarse con un modelo más simple, quizá un modelo de lenguaje n-gram, pero es difícil verlo como representativo de cómo funcionan los modelos de lenguaje más grandes
Me pregunto si el autor está afirmando que los LLM son generadores de texto basados en cadenas de Markov.
Es decir, ¿quiere decir que la distribución de probabilidad del siguiente token generado es igual a la probabilidad de esa secuencia de tokens en los datos de entrenamiento?
Si fuera así, ¿significaría que si “simplemente” construyéramos una cadena de Markov con los datos de entrenamiento originales podríamos obtener un rendimiento similar al de un LLM?
el estado es un vector de tokens del largo del contexto, y el modelo describe la matriz de transición. Para un vector de tokens dado, del tamaño del largo de contexto, produce las probabilidades del siguiente vector de tokens del mismo tamaño.
La parte difícil es la clasificación: saber qué textos de los datos de entrenamiento son “similares” al texto del prompt.
Con el ejemplo del post del blog, sería así:
Original string: 'And only l'
Similar strings: 'hat only l' 's sickly l' ' as\nthey l' 'r kingly l'
Eso sugiere que podría haber una forma de evitar el transformer y obtener el mismo resultado. Me pregunto si sería más eficiente.
Por ejemplo, dado un modelo base, quizá sería posible entrenar otra cosa y ejecutarla en dispositivos mucho más pequeños.
Es bastante difícil entender qué afirma exactamente el autor que demostró.
Leí varias veces la sección “Interpretation: Why Does the Approximation Work?”, pero me suena a una explicación mecánica de las etapas del transformer. No sé cuál es la afirmación central.
Me pregunto si se sabe que los desplazamientos de attention y FF suelen apuntar en la misma dirección.
Ya me sorprende un poco que estén en el mismo espacio latente a través de las capas. ¿La red FF no podría hacer una rotación arbitraria? Creo que estoy entendiendo algo mal.
Inevitablemente algo se pierde, y una de las cosas que se pierde es describir con precisión una rotación arbitraria en un espacio de alta dimensión.