¿Claves, consultas, valores y atención al producto punto?
"Se necesitan claves y consultas de tokens para comunicarse entre sí", dice la popular analogía de Internet. ¿Pero por qué? Hay muchas analogías excelentes sobre cómo funcionan, pero mucho menos material sobre por qué realmente las necesitamos. ¿Existen alternativas o estos conceptos abstractos son inevitables?
Esto podría parecer una pregunta tonta dado el éxito total de la arquitectura Transformer en 2026. Pero si la historia sirve de indicación, los Transformers eventualmente serán reemplazados por algo mejor. Cuanto más comprendamos por qué funcionan, más rápido podremos superarlos.
¡Increíblemente, podemos demostrar que la forma general del Transformador es difícil de evitar! Comience con algunas presiones clave de diseño y la sopa de letras de matrices comenzará a parecer mucho menos arbitraria. Sin depender de ninguna analogía de "el token hace una pregunta", veremos que las consultas surgen de un problema de simetría, mientras que los valores y las cabezas de atención aparecen cuando reemplazamos una matriz de peso dinámica difícil de manejar con un pequeño conjunto de transformaciones reutilizables.
Y a medida que conciliamos nuestro modelo de juguete con el Transformer original (autoregresivo), terminaremos con una última conexión fascinante: el MLP (el bloque de avance a menudo pasado por alto) puede verse como un almacén de valores clave propio.
¿Listo para inventar el Transformer por nosotros mismos? En primer lugar, debemos retroceder hasta 2014 y preguntarnos por qué las redes neuronales recurrentes no son suficientes.
Por qué falla la memoria fija
Hay muchos problemas con las redes neuronales recurrentes estándar (RNN). Algunos problemas, como el problema de la "desaparición del gradiente", fueron resueltos por la extremadamente popular LSTM (red de memoria a corto plazo) de Hochreiter y Schmidhuber.[1]que es una versión más avanzada de RNN. Pero queda un problema central para todos los tipos: los RNN "aplastan" las entradas pasadas en una memoria fija.
Para ver por qué, recuerde que un RNN tiene un estado de memoria en el que la entrada actual y el estado anterior "escriben". De esta conexión entre estados en el tiempo es de donde proviene la “recurrencia”. Considere el siguiente RNN de dos capas desenrollado en el tiempo:
Aquí las entradas son verdes, las salidas son azules y las dos capas intermedias de la red que calculan las activaciones de redes neuronales estándar son grises. Todos los componentes son vectores y las flechas dirigidas son pesos matriciales que multiplican sus entradas. Para ver la red en un momento dado, simplemente mire una única porción vertical del diagrama; en ese segmento, hay dos unidades grises fijas que colectivamente representan la "memoria" de la red en ese momento.
Usemos una cruda analogía para mostrar por qué esta “memoria fija” es un problema. Supongamos que codifica la oración "Tengo cinco dólares" en la memoria del RNN, que representaremos con el rectángulo gris a continuación:
Ahora ampliemos esa frase a “tengo cinco dólares y cuarenta centavos en mi bolsillo”:
Vaya. Después de agregar más información a la memoria finita, hay una mayor “competencia” por los bienes inmuebles y parte de la memoria se sobrescribe. Esto es desastroso cuando es necesario recordar hechos específicos o seguir instrucciones muy específicas.
Quizás esté pensando: "¿Por qué no utilizar una memoria dinámica que crece con la secuencia de entrada?"
¡Gran intuición! Esto es exactamente lo que Bahdanau et al.[2]Lo intenté en 2014 cuando popularizaron la idea de "atención" dentro de RNN (sí, ¡los RNN usaron atención antes que Transformers!)
La idea, a alto nivel, es mantener toda la historia previa de los estados RNN como nuestra "memoria en expansión", que naturalmente tiene la propiedad de crecer con la secuencia de entrada. Considere el siguiente diagrama actualizado (esta no es la arquitectura exacta del artículo de Bahdanau, sino una analogía):
Observe las nuevas conexiones en rojo; estos conectan los estados grises en las capas RNN con cada entrada anterior en el tiempo. Antes de agregar estas conexiones, cada estado se vio obligado a comprimir toda la información histórica y transmitirla en conexiones “recurrentes” de izquierda a derecha. Pero esta compresión ya no es necesaria porque ahora cada estado tiene acceso directo a todo el historial de entradas, ¡también conocido como nuestra memoria en crecimiento!
Queda un gran desafío pendiente: entrenar la velocidad. En el momento del entrenamiento, tenemos cada secuencia de entrada disponible desde el principio. Pero generar el enésimo resultado final requiere N pasos secuenciales en el tiempo dadas las dependencias recurrentes; con secuencias de entrada largas, tenemos cálculos secuenciales largos que las GPU no pueden paralelizar.
La idea del artículo histórico de Vaswani et al.[3]es esto: ¿Qué pasa si podemos eliminar esas conexiones recurrentes de izquierda a derecha? ¿Qué pasa si las conexiones rojas son todo lo que necesitas? Vea el siguiente diagrama que representa esta idea a un alto nivel:
Apilemos las arquitecturas recurrentes y no recurrentes una al lado de la otra y comparemos su número total de pasos de cálculo colocando etiquetas ordenadas para cada paso:
Tenga en cuenta que el modelo no recurrente solo necesita 2 pasos de cálculo, ya que puede calcular cada capa completamente en paralelo una vez que se calcula la capa anterior. Por otro lado, el modelo recurrente necesitaba 5 pasos de cálculo debido a las dependencias recurrentes dentro de cada capa. A medida que la secuencia se hace más larga, el modelo no recurrente permanecería en 2 pasos, mientras que los pasos del modelo recurrente crecerían para siempre con la secuencia. ¡No pinta bien la recurrencia!
Pasemos al modelo no recurrente de la derecha. Ahora nos topamos con nuestro siguiente desafío: ¿Cómo diablos deberíamos elegir los pesos para estas conexiones rojas?
Transformadores y pesos dinámicos
Si observa el diagrama de nuestra red no recurrente, parece una red neuronal ordinaria con dos capas y cuatro unidades por capa. Con una red así, podríamos preguntarnos: ¿Por qué no aprender pesos fijos como lo hacemos con cualquier otra red?
Pero a diferencia de este diagrama fijo, las secuencias no son fijas durante el entrenamiento o la predicción. La red podría encontrar una secuencia de entrada de tamaño 2 o podría encontrar una secuencia de entrada de tamaño 2000, y nuestro diagrama podría crecer indefinidamente hacia la derecha con cada vez más unidades y pesos. Entonces, ¿cómo configuramos esos nuevos pesos entrantes?
Lo que necesitamos es una función para generar nuevos pesos sobre la marcha, con parámetros que podamos configurar durante el entrenamiento. ¡Necesitamos pesos dinámicos!
Para identificar una buena función, acerquémonos a una unidad particular con una secuencia de longitud 3:
Primero, una nota importante sobre la adición de un nuevo diagrama: para completar el diagrama, necesitábamos agregar conexiones de salto. A través de estos, las entradas x1, x2 y x3 se vuelven a agregar a las salidas de O1, O2 y O3 respectivamente (y este proceso se repite nuevamente en la siguiente capa cuando O1, O2 y O3 se convierten en entradas). Estos "saltos" liberan a las transformaciones intermedias de tener que preservar las entradas y les permiten concentrarse en la tarea mucho más fácil de ajustar aditivamente esas entradas. Las conexiones omitidas son una optimización del rendimiento crítica que fue pionera en la famosa arquitectura ResNet de He et al.[4].
Ahora volvamos a nuestra función para generar pesos dinámicamente. Primero, debemos decidir de qué deberían ser función los pesos.
Podríamos comenzar definiendo un peso determinado como una función de la salida de la unidad fuente del peso, así como la posición de esa unidad fuente en la secuencia izquierda-derecha. Para simplificar las cosas, combinaré "entrada" y "posición" y solo diré "entrada" en el futuro (esta combinación se puede implementar codificando la posición en la entrada directamente; para ver cómo se puede hacer esto, lea sobre codificaciones sinusoidales en la arquitectura original de Transformer). Entonces, nuestro diagrama se parece a este:
Aquí hay una desventaja obvia: si x1 necesita ser "importante" para O3 a través de un peso "grande", entonces se ve obligado a ser también importante para O2 y O1, ya que todos comparten el mismo valor para sus pesos x1. Esto significa que las unidades azules cercanas dentro de una capa calcularán cosas muy similares, lo que anula la flexibilidad de esta arquitectura para modelar conceptos únicos. Para solucionar este problema, necesitamos romper la "simetría" y hacer que cada uno de los pesos de O1, O2 y O3 sea único entre sí.
Podríamos romper la simetría haciendo que un peso dado sea función tanto de la producción de la unidad fuente como de la producción de la unidad final; sin embargo, usar la salida de la unidad final directamente es circular ya que primero necesitamos el peso para calcular la salida final.
Para sortear esta circularidad, tenga en cuenta que el propósito único de cada unidad es modificar su flujo de entrada; por ejemplo, O3 es el único responsable de modificar x3 a través de la conexión de salto. Dado que x3 y O3 están entrelazados de manera única, x3 es un candidato natural para el segundo argumento que rompe la simetría para generar los pesos de O3.
Cuando actualizamos el diagrama, queda claro que ahora todos los pesos son únicos:
Si tiene buen ojo, ¡puede que ya empiece a detectar la "clave" y la "consulta" de la arquitectura Transformer! Pero si no, no te preocupes; Nos acercaremos a eso de manera más formal.
Una última advertencia: también podríamos romper la simetría haciendo que los pesos sean cero para cada entrada excepto la que modifica una unidad, pero luego obtenemos el siguiente diagrama que muestra que se trata efectivamente de una red sin estado:
En otras palabras, la ruptura de la simetría es una condición necesaria pero no suficiente. También necesitamos cierta interactividad distinta de cero entre unidades y otras unidades de diferentes pasos de tiempo.
Las claves, consultas y valores surgen de forma natural
Ahora que hemos decidido generar cada peso mediante una función de dos argumentos, debemos decidir cómo se ve realmente esta función.
Recuerde que cada peso es una matriz que multiplica los vectores de entrada en vectores de salida, por lo que nuestra función en realidad necesita tener un valor matricial. Pero eso introduce un nuevo problema: ¿Cómo hacemos que estas matrices sean dinámicas sin aumentar el número de parámetros en nuestro modelo?
Para ver por qué la explosión de parámetros es real, escribamos nuestra función en forma matricial, usando el peso entre x1 y O3 como ejemplo:
Aquí, cada posición (i, j) de la matriz (dxd) tiene su propia función. Si parametrizamos cada función por separado, son conjuntos de parámetros separados de d cuadrado que deben ajustarse. Cuando la dimensión d está en el rango de cientos, son >100.000 conjuntos de parámetros que deben aprenderse. ¡Pase duro!
Podríamos reducir enormemente ese número ajustando únicamente una matriz diagonal:
Sin embargo, todavía quedan cientos de conjuntos de parámetros de funciones que ajustar, y ni siquiera hemos llegado al tamaño de esos conjuntos. Además, tenga en cuenta que las matrices diagonales realizan una multiplicación de elementos en sus entradas; Si partes del vector final necesitan ser cero por razones como la escasez, entonces las funciones en esos elementos diagonales deben ser muy cercanas a cero o exactamente cero, lo que supone una gran carga para las funciones mismas.
Para ver una alternativa, reescribamos esa matriz diagonal como una combinación lineal de matrices one-hot:
¿Qué pasa si en lugar de usar esas matrices únicas en la suma, usamos las matrices que quisiéramos? Podríamos hacer de cada matriz un parámetro para aprender y luego elegir el número final de matrices que se pueden aprender para que el recuento total de parámetros anidados sea razonable. Entonces podríamos reescribir nuestro peso entre x1 y O3 como una nueva suma matricial:
Impresionante; ¡Acabamos de reemplazar una dolorosa matriz dinámica con una pequeña cantidad de matrices estáticas y coeficientes dinámicos! Aquí, las V son las matrices estáticas que aprendemos como parámetros libres y las funciones en la suma son nuestros coeficientes dinámicos con valores escalares. Tenga en cuenta las V; aparecerán en los cálculos del “valor” de atención más adelante.
Finalmente, necesitamos definir las funciones con valores escalares dentro de esa suma anterior. Para hacer eso, descompongamos aditivamente una de las funciones en dos funciones de "no interacción" de un solo argumento intercalando una función de interacción pura:
La desventaja de tener funciones sin interacción es la misma que encontramos con las simetrías de peso en la última sección: si el término de la izquierda es grande, entonces será grande para todas las unidades conectadas a x1 ya que x1 es la única dependencia en el término. Queremos romper esta simetría, así que mantengamos solo la función de interacción intermedia v, a la que llamaré función de “atención” de ahora en adelante.
Un candidato para la función de atención fue presentado en Bahdanau et al.[2]:
Puede reconocer esto como una red neuronal simple de una capa con activación tangente hiperbólica, donde el resultado final se reduce a un valor escalar mediante un producto escalar.
Hay una desventaja de usar la función tanh aquí: solo interactúa débilmente debido a su no linealidad. Los contornos de la función se aplastan de forma no lineal, pero su estructura básica todavía parece una suma lineal. Para ver esto visualmente, compare la similitud en los gráficos de contorno para tanh de x1 más x3 versus la suma directa de x1 y x3 (donde haremos entradas 1D para x1 y x3 para facilitar la visualización):
Destaquemos realmente el problema con un ejemplo: supongamos que todas las unidades de nuestra red usan una de las funciones anteriores y una unidad tiene entradas x1 y x3.
Si esta unidad necesita generar un valor positivo grande, entonces solo hay una región que lo satisface en cualquiera de los gráficos: la mitad superior derecha. Pero si otra unidad comparte la misma entrada x1 y requiere que x1 esté en el lado izquierdo (es decir, negativo) para lograr su objetivo, entonces estará atrapado en el requisito de la primera unidad a menos que x3 sea un número positivo enorme para compensar. Los requisitos de ambas unidades son, hasta cierto punto, incompatibles. Esta red es menos flexible en lo que sus diferentes unidades pueden modelar cuando comparten algunos insumos.
Por otro lado, observe el gráfico de contorno de la función producto a continuación:
Ahora, hay dos regiones distintas donde el resultado final es grande, ¡no solo una! Si otra unidad comparte la misma entrada x1 y necesita que x1 sea negativo, no hay problema; la primera unidad aún puede generar un valor positivo grande si x3 es negativo. Los requisitos de ambas unidades ya no son tan incompatibles como antes. Este es un argumento matemático sutil sobre cómo la “interactividad” (por ejemplo, a través de un producto) es importante para permitir que las unidades modelen cosas únicas a pesar de que sus entradas se superpongan.
Quizás se pregunte: "¿Por qué no utilizar una red neuronal más profunda para modelar una función de atención más interactiva?" Esto es bastante costoso en la práctica porque el número de cálculos de atención aumenta cuadráticamente con la longitud de la secuencia. Y considere que tampoco complicamos cada capa de una red neuronal tradicional; A menudo utilizamos funciones no lineales simples como ReLU (unidades lineales rectificadas) y dejamos que capas adicionales acumulen complejidad de forma iterativa.
¿Podríamos hacer algo similar aquí, es decir, utilizar una función de atención simple pero interactiva y desarrollar complejidad a través de capas? Si es así, la función del producto es una gran candidata, ¡y a las GPU también les encanta!
Hay un ajuste que debemos hacer, ya que nuestras entradas de atención son vectores multidimensionales, no entradas 1D. En lugar de un producto escalar, necesitamos un producto escalar. De manera más general, es posible que desee calcular el producto escalar en un subespacio particular, lo que da como resultado una forma bilineal similar a esta:
Tenga en cuenta que esto se reduce a un producto escalar simple cuando A es la matriz identidad, por lo que esta forma es más general.
Hay otra consideración que debemos comenzar a discutir aquí: el “caché” de cálculo (lo que quizás haya escuchado como el “caché de valores clave” en las arquitecturas modernas). Podemos evitar una cantidad significativa de cálculos de atención almacenando en caché multiplicaciones de matrices (por ejemplo, A multiplicado por x1 arriba) para entradas pasadas, ya que reutilizamos estos cálculos pasados para cada nuevo paso de tiempo en cada capa.
Sin embargo, el tamaño de este caché se convierte en un problema bastante importante. Ejemplo: si tenemos una longitud de secuencia de 5000 (bastante más pequeña que muchas conversaciones LLM estándar), 50 capas, 20 matrices V diferentes -> 20 funciones de producto escalar diferentes (consulte nuestra definición de suma de matrices anteriormente), productos matriz-vector almacenados en caché de dimensión 1000 y 2 bytes por número de punto flotante, entonces el tamaño final de nuestro caché es: 5000 x 50 x 20 x 1000 x 2 = 10 GB. Ay.
Cachés más grandes significan más costo de GPU y sobrecarga de memoria. Una optimización es simplemente reducir la dimensión de los productos almacenados en caché de 1000 a una dimensión inferior r, siempre que no afecte significativamente la precisión del modelo. Pero esto requiere que el producto escalar final esté ahora en un espacio de r dimensiones en lugar de uno de 1000 dimensiones.
Podemos reducir esa dimensión del producto escalar factorizando la matriz A en dos matrices que "proyectan" los vectores x3 y x1 en el espacio r-dimensional:
Aquí, Wq y Wk son matrices (rx 1000), y la expresión más a la derecha es el producto escalar deseado entre dos vectores de r dimensiones. Si r es solo 200, entonces el tamaño final de la caché pasa de 10 GB a 2 GB: ¡una reducción de 5 veces!
Ahora, la gran revelación: los términos izquierdo y derecho en el nuevo producto escalar no son más que la “consulta” y la “clave” en la arquitectura Transformer, ¡y las matrices de proyección son las mismas matrices que se pueden aprender en el artículo original! (Una advertencia es que la arquitectura Transformer agrega escalamiento para la estabilidad computacional, de ahí el término "atención al producto escalado". Pero durante el resto de este artículo, me centraré en la forma de la arquitectura en lugar de en optimizaciones de entrenamiento como el escalado).
Ahora podemos sustituir esta función de atención en los coeficientes de nuestra suma matricial anterior, con diferentes parámetros para cada j:
Resumen rápido de dónde estamos: La función “f” a la izquierda escupe la matriz de peso entre la entrada x1 y la unidad O3; esta matriz es igual a una suma que involucra H funciones de atención escalares como coeficientes y, según nuestra discusión, hemos elegido utilizar la atención del producto escalar para cada función. Entonces podemos representar la suma total de todas las entradas ponderadas en O3 con una nueva suma s3:
El término WQ * x3 es la consulta para x3, el término WK * xi es la clave para xi y el término V * xi es el valor de xi. El subíndice j en las matrices denota una “cabeza” de atención particular entre las cabezas H; cada encabezado tiene una consulta única que se puede aprender, una clave y una matriz de proyección de valores. ¡Esto está empezando a parecerse a un Transformer!
De nuestra atención a la atención del transformador
En nuestro ejemplo anterior, la suma exterior se extiende sobre una secuencia insignificante de tamaño 3. Pero esta suma aumenta de tamaño cuando la secuencia se vuelve muy larga.
Una forma de abordar esto es un truco estándar de aprendizaje automático llamado normalización: tomar los coeficientes existentes en la suma ponderada y transformarlos para que los coeficientes transformados sumen 1, independientemente del número y tamaño de los coeficientes anteriores.
También podríamos querer escasez en los coeficientes transformados; En el dominio del lenguaje, hay cientos de palabras irrelevantes que pueden ganar peso rápidamente, por lo que queremos aplastar agresivamente todo excepto unos pocos coeficientes.
Para ver cómo debemos agregar normalización a nuestra fórmula, tenga en cuenta que debemos hacer esta normalización en la suma sobre la longitud de la secuencia, no en la suma sobre las cabezas. En ese caso, necesitamos invertir el orden de las sumas de la siguiente manera:
Ahora, el término dentro del paréntesis exterior parece una suma sobre la longitud de la secuencia L con ponderaciones de coeficientes (productos escalares de consulta clave) en los vectores de valores finales, tal como queríamos. Luego aplicamos una transformación de normalización a esos coeficientes del producto escalar. La elección natural de transformación para lograr una escasez suave es la transformación softmax, que exponencia cada coeficiente y divide el resultado por la suma de todos los coeficientes exponenciados:
Esto se llama "softmax" porque empuja agresivamente el coeficiente más grande (máximo) a 1 y aplasta todos los demás coeficientes hacia 0. El parámetro tau en el exponente controla qué tan agresiva es esta transformación.
Bien, todo esto es fantástico, pero ¿dónde están las matrices Q, K y V que nos prometió el título del artículo?
Para hacer coincidir nuestra expresión con la ecuación icónica de Transformer, ignoremos temporalmente la suma externa sobre las cabezas eligiendo una cabeza j en particular y reescribamos la suma normalizada en la posición L para esa cabeza:
Aquí, q es el vector de consulta para la unidad L-ésima, y las filas de las matrices K y V son los vectores clave y de valor respectivamente. El subíndice 1:L indica que solo usamos las claves y valores para las posiciones 1 a L, lo que refleja el hecho de que la unidad L-ésima solo está conectada a unidades anteriores en el tiempo. El producto entre q y K^T crea un vector que contiene cada producto escalar entre q y una clave en K, y el softmax en la parte superior normaliza las puntuaciones finales del producto escalar. Luego, hacemos una suma ponderada de vectores de valor en V usando las puntuaciones softmax resultantes.
Podemos extender esto a una única expresión matricial consolidada para la suma en cada posición, no solo en la posición L:
¡Esta es exactamente la ecuación icónica de Transformer (menos el escalado dimensional), pero para un Transformer autorregresivo! Aquí, cada fila de S es la suma final en esa posición, y cada fila de Q, K y V son los respectivos vectores de consulta, clave y valor en esa posición. Dado que estamos usando las matrices completas Q, K y V, necesitamos una matriz de enmascaramiento M para poner a cero los productos escalares no válidos (como el producto escalar entre una consulta en la posición p y una clave en una posición posterior) agregando infinito negativo a esas entradas; el softmax luego reduce esos resultados a cero.
¿Por qué escribir esto con matrices en lugar de nuestra suma original si las dos expresiones son idénticas? La respuesta es simple: a las GPU les encanta la multiplicación de matrices. A menudo es más eficiente reescribir un cálculo como una multiplicación de matrices, incluso si eso significa hacer algún trabajo innecesario como calcular productos escalares que de todos modos quedarán enmascarados. Dicho esto, encuentro que la notación de suma es más fácil de usar para explicaciones, así que la usaré durante el resto del artículo.
Cerremos la sección volviendo a la suma de las cabezas de atención. Es posible que haya notado que, si bien sumamos las cabezas de atención H, Vaswani et al.[3]e implementaciones como “MultiheadAttention” de PyTorch concatenan las salidas en cada cabezal y multiplican la concatenación por una “matriz de mezcla” Wo que colapsa la concatenación a la dimensión del modelo anterior.
Esto podría parecer una diferencia cosmética; ¿Por qué concatenar si simplemente vas a colapsar?
El truco es el almacenamiento en caché, pero ahora centrándonos en los vectores de valores: necesitamos la matriz de mezcla que nos permita comprimir esos vectores de valores (V * x) para que nuestro caché no explote. Veamos qué sucede con nuestra expresión de doble suma si concatenamos la salida de cada cabeza j en un vector de bloque en lugar de sumar las cabezas:
En aras de la legibilidad, he sustituido los coeficientes normalizados anteriores con alfas para que no haya un montón de exponenciales por todas partes.
Ahora multipliquemos a la izquierda este vector de bloque largo por la matriz de mezcla Wo. Primero reescribamos Wo en una forma matricial de bloques correspondiente y luego realicemos la multiplicación:
¡Observe cuán similar es el resultado final a nuestra doble suma original! La única diferencia es la matriz de bloques que queda multiplicando V; cuando cada matriz de bloque es la matriz identidad, la nueva expresión es totalmente idéntica a nuestra expresión original.
Quizás se pregunte: "¿Por qué no simplemente absorber las matrices de bloques Wo en las matrices V?" A primera vista, las dos matrices parecen redundantes si se pueden agrupar en una sola. Pero la razón por la que queremos mantenerlos separados es la misma razón por la que no utilizamos una única matriz A para nuestro producto punto de atención: reducir el tamaño de nuestra caché.
Digamos que absorbimos todo en una sola matriz V. Si la dimensión de nuestro modelo es 1000 y usamos las mismas condiciones de nuestro ejemplo anterior de caché basado en claves, entonces el tamaño de caché acumulativo para nuestros vectores de valores es de 10 GB, difíciles de manejar. Pero si V tiene la forma (rx 1000) donde r < 1000, entonces la dimensión final de nuestros vectores de valores será r < 1000 y nuestro tamaño de caché puede ser considerablemente menor.
Aquí radica la necesidad de las matrices Wo; Una vez que hayamos terminado de divertirnos en dimensiones inferiores, debemos multiplicar los valores finales por una matriz de forma (1000 xr) para "proyectar hacia arriba" de regreso a la dimensión más grande del modelo de 1000.
The Transformer MLP: la otra tienda de valor clave
Todo lo que hicimos antes fue definir los pesos en nuestra red. Si aislamos una entrada xi en la suma de las entradas ponderadas de nuestra última expresión, entonces su peso en la unidad L es igual a:
Aquí, cada coeficiente alfa es el escalar normalizado softmax de la unidad L para la entrada i y una cabeza dada j.
Una vez que tenga todos los pesos, todo lo que necesita hacer es sumar las entradas ponderadas en cada unidad, calcular la no linealidad de la posición y volver a agregarla al flujo residual a través de la conexión de salto para actualizar la representación posicional. ¡Listo, ahí está tu red neuronal!
Estos pasos se pueden representar visualmente mediante el siguiente diagrama recapitulativo, utilizando la misma secuencia de longitud 3 para simplificar:
Quizás hayas notado que Vaswani et al.[3]introduzca dos nuevas matrices de peso (W1 y W2) que intercalen la no linealidad en las unidades azules. A primera vista, estas adiciones parecen complementos arbitrarios. Si ya hemos hecho todo este trabajo para definir pesos dinámicos mediante la atención, ¿realmente necesitamos dos nuevas matrices en torno a la no linealidad?
Resulta que W2 no está ahí porque sí. Geva et al.[5]Bríndenos una perspectiva importante para comprender por qué: el Transformer MLP (bloque de retroalimentación) actúa como un almacén de valores clave por derecho propio.
Para ver esto, mostremos dónde residirían las "claves" y los "valores" en una red MLP de ejemplo:
En esta red, la capa intermedia tiene unidades rojas, verdes y azules, cada una de las cuales modela patrones diferentes. Los pesos entrantes (primera capa) del mismo color representan los vectores "clave" para cada patrón, y los pesos salientes (última capa) del mismo color representan los vectores "valor" para cada patrón. Geva et al.[5]demostrar cómo podrían ser algunos de estos “patrones”; Para el modelo que entrenaron, un ejemplo de patrón aprendido era cuando la oración de entrada terminaba en la palabra "sustitutos", y otro era cuando la oración de entrada contenía una relación "parte de".
Ahora veamos qué hacen realmente aquí las “claves” y los “valores”. Cuando la entrada se alinea estrechamente con el vector clave (pesos entrantes) para un patrón determinado, el producto escalar entre la clave y la entrada es alto y la unidad asociada con ese patrón genera un número grande. Luego, ese número multiplica el vector de valor (ponderaciones salientes) para el patrón y este vector de valor ponderado finalmente se agrega a la capa de salida.
Volvamos a nuestros parámetros adicionales: ¿Qué nos dice esto sobre por qué necesitamos W2?
Recuerde que W2 es solo la matriz multiplicada por la salida de la no linealidad en el transformador original. En nuestro ejemplo de MLP, W2 es exactamente igual a la matriz de pesos en la capa final, es decir, una matriz donde cada columna es un vector de valores.
Si nos deshacemos de W2, nuestro diagrama se vería así:
Cada unidad en esta red escribe en una única coordenada del flujo residual (ya que agregamos la capa de salida directamente al flujo residual). Pero hay dos razones por las que esto es malo:
Si queremos que nuestras unidades aprendan patrones arbitrarios, entonces es posible que cada uno de esos patrones arbitrarios no se asigne claramente a coordenadas únicas del "espacio semántico" en esa capa. En ese caso, es mucho más expresivo permitir que cada unidad escriba vectores (valores) arbitrarios en el flujo residual en lugar de acoplarlos a coordenadas de salida individuales. Queremos que nuestros MLP modelen tantos patrones como sea posible, lo que nos empuja a tener más unidades en la red MLP en relación con la dimensión del modelo en el flujo residual (donde este último se mantiene más pequeño para evitar cálculos de atención complicados que no se benefician tanto de las dimensiones altas). En ese caso, una correspondencia de coordenadas 1:1 ni siquiera es posible; De todos modos, necesitamos una matriz para “reducir el proyecto” a la dimensión del modelo más pequeña.
Hemos pasado todo este tiempo discutiendo W2; ¿Qué pasa con W1, la matriz adicional que se multiplica dentro de la no linealidad?
En realidad, esta puede ser la matriz más arbitraria. En nuestra configuración, W1 podría absorberse en nuestras matrices de peso dinámico (específicamente, Wo de la sección anterior) y eliminarse como un parámetro adicional. Pero hay una razón por la que esto no se puede hacer en el Transformer original: en esa configuración, hay dos pasos intermedios entre la salida de atención y la no linealidad que impiden la consolidación de matrices.
Primero, el Transformador divide nuestra única conexión de salto grande en una conexión de salto entre la entrada y la salida suma, así como un segundo salto entre la salida suma y la salida de no linealidad. Vea el diagrama modificado con la conexión de salto dividido:
Una vez que tenga esta conexión de salto adicional, necesitará una matriz de ponderación separada aplicada a las sumas de la salida de los bloques de atención (sigma mayúscula) y los residuos de los primeros saltos. La matriz de ponderación aplicada a esa suma de resultados ya no se puede absorber limpiamente, por lo que necesita un parámetro separado.
En segundo lugar, existe una transformación de escala adicional entre la suma de las producciones de atención y la no linealidad que impide la absorción. Si bien estas transformaciones de escala son importantes, son optimizaciones de entrenamiento, por lo que las omití para centrarme aquí exclusivamente en la forma arquitectónica.
Recapitulando la lógica
Date una palmadita en la espalda; ¡acabas de inventar el transformador! Aquí hay un resumen de los pasos que tomamos para llegar aquí:
Necesitábamos acceder a todos los estados pasados de la secuencia sin pérdida de compresión. Esto requería conexiones directas con estados pasados, también conocidos como "atención". Necesitábamos paralelismo para un entrenamiento rápido de GPU. Esto requirió abandonar la recurrencia. Necesitábamos una forma de establecer nuevos pesos para entradas futuras sin introducir simetrías de peso en toda la red. Esto requirió hacer que los pesos fueran una función de la unidad de origen y de la unidad final (ambas únicas por peso), donde usamos el valor del flujo residual como sustituto de la "unidad final". Estos dos argumentos de función son la "clave" y la "consulta" antes de la proyección. Necesitábamos hacer que la función generadora de peso evitara simetrías codificando interacciones puras entre los argumentos. Una opción eficiente en la GPU para esta función de interacción fue el producto escalar entre los argumentos clave y de consulta. Necesitábamos reducir el tamaño de nuestra caché para las multiplicaciones de matriz-vector reutilizables en el producto escalar, lo que requería calcular el producto escalar en dimensiones más bajas. Esto requirió matrices de proyección Wk y Wq para proyectar la clave y consultar en ese espacio dimensional inferior. Necesitábamos reducir nuestra búsqueda en todo el espacio d^2 de operaciones lineales (matrices) a una búsqueda en un número menor H de operaciones posibles entre las que podemos seleccionar a través de nuestros pesos de producto escalar. Esto está representado por H proyecciones de diferentes valores V que representan nuestras H cabezas de atención. Necesitábamos que las ponderaciones finales de las proyecciones de valores fueran escasas para que el ruido no explotara en secuencias más grandes. Esto requirió la normalización softmax de los coeficientes del producto escalar sobre la longitud de la secuencia, por cabeza. Necesitábamos reducir el tamaño de nuestro caché para las multiplicaciones del vector matriz reutilizable entre las proyecciones de valor V y las entradas. Esto requirió que Vs "proyectara hacia abajo" los vectores de valor final en un espacio dimensional inferior, y una matriz de mezcla separada Wo para "proyectar hacia arriba" nuevamente a la dimensión del modelo. Necesitábamos un "espacio de características" más grande para que nuestras no linealidades aprendieran muchos patrones, al tiempo que permitíamos que estas características escribieran vectores arbitrarios en el flujo residual de menor dimensión. Esto requirió una nueva matriz W2 que multiplica la salida de la no linealidad.
Por qué los transformadores no son inevitables
Lo único inevitable en AI/ML es que una arquitectura reemplace a otra, y los Transformers no son una excepción a la regla.
Pero, ¿por qué se reemplazarán los Transformers si hoy funcionan tan bien?
Para empezar, los Transformers tienen una gran desventaja: los cálculos escalan cuadráticamente con la longitud de la secuencia. Podría implementar la atención de ventana deslizante para intentar solucionar este problema, pero luego perdería la capacidad de recordar hechos concretos o instrucciones ocultas en un contexto más amplio.
Quizás se pregunte: "Pero la mayoría de las puntuaciones de atención son básicamente cero debido a la normalización softmax. ¿Necesitamos calcular todas las puntuaciones n^2?"
Se trata de una gran idea y, de hecho, se han hecho intentos de crear mecanismos de atención escasos. Sin embargo, a pesar de la importante reducción en los cálculos, estos mecanismos, de manera contradictoria, funcionan más lento que la atención cuadrática completa debido a que los cuellos de botella en la memoria de la GPU se convierten en el mayor problema, como lo analizan Dao et al.[6].
Esto nos lleva a otro punto: los Transformers se están beneficiando enormemente al estar hechos a medida para GPU. Muchas de sus decisiones de diseño giran en torno a la memoria peculiar y los patrones de cómputo de las GPU, y cualquier cosa que pueda destronar a los Transformers tiene que pasar en última instancia la “prueba del olfato de la GPU”. En otras palabras, estamos atrapados en un mínimo local de hardware y debido a ello nos estamos perdiendo algoritmos mucho mejores. La inteligencia humana no funciona con GPU, entonces, ¿por qué deberíamos conformarnos con Transformers?
Referencias
[1]S. Hochreiter y J. Schmidhuber, “Memoria a corto plazo”, Neural Computation, vol. 9, núm. 8, págs. 1735-1780, 1997.
[2]D. Bahdanau, K. Cho e Y. Bengio, “Traducción automática neuronal mediante el aprendizaje conjunto de alineación y traducción”, Conferencia internacional sobre representaciones de aprendizaje (ICLR), 2015.
[3]A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, AN Gomez, Ł. Kaiser e I. Polosukhin, “La atención es todo lo que necesita”, Avances en los sistemas de procesamiento de información neuronal 30 (NIPS), págs. 5998-6008, 2017.
[4]K. He, X. Zhang, S. Ren y J. Sun, “Aprendizaje residual profundo para el reconocimiento de imágenes”, Conferencia IEEE sobre visión por computadora y reconocimiento de patrones (CVPR), págs. 770-778, 2016.
[5]M. Geva, R. Schuster, J. Berant y O. Levy, “Transformer Feed-Forward Layers Are Key-Value Memories”, Conferencia sobre métodos empíricos en el procesamiento del lenguaje natural (EMNLP), págs. 5484-5495, 2021.
[6]T. Dao, DY Fu, S. Ermon, A. Rudra y C. Ré, “FlashAttention: atención exacta rápida y eficiente en memoria con IO-Awareness”, Avances en sistemas de procesamiento de información neuronal 35 (NeurIPS), págs. 16344-16359, 2022.