Auto-atención: la secuencia se mira a sí misma

La arquitectura terminó la lección anterior sin una sola pieza recurrente y con una única operación en el centro, la atención, a la que hay que pasarle tres cosas: consultas (queries), claves (keys) y valores (values). Las primeras se han quedado sin quien las haga. Era el decoder, un paso detrás de otro, y el decoder era recurrente. Lo único que sigue teniendo una fila por posición es la frase misma, ya convertida en vectores, así que para los tres papeles hay un candidato y no hay dos.

Usarlo tal cual, eso sí, no funciona, y falla por algo que se dice en una línea: a lo que más se parece un vector es a sí mismo. Cada posición se daría a sí misma el peso más alto, en todas las filas a la vez, y la capa se limitaría a devolverle a cada una lo que ya traía —sin un solo parámetro con el que hacer otra cosa—. Y hace falta que haga otra: en las llaves del coche están ahí, para conjugar están hay que haber mirado llaves, tres posiciones más atrás, y no coche, que está pegado y en singular. Esta lección monta la capa que sabe hacer eso, y lo que cuesta son tres matrices.

La rejilla de abajo es una frase contra sí misma: la fila es quien pregunta, la columna es quien contesta, y cada casilla es cuánto de esa columna entra en lo que lee esa fila. Escribe la frase que quieras y, sobre todo, apaga y enciende las proyecciones. Apagadas, la casilla más alta de cada fila es la que esa posición se da a sí misma, en las seis filas a la vez. Encendidas, la fila de están se va a buscar llaves por encima de todo lo demás.

Tres lecturas de la misma lista

Apila la frase por filas, como el bloque anterior apilaba los estados del encoder:

donde la fila $t$ es $\mathbf{x}_t^{\top}$, el vector del token que ocupa la posición $t$ —transpuesto, porque en este curso los vectores son columnas—. Eso es todo lo que entra en la capa: una sola matriz, ninguna segunda red.

Las tres listas salen de ahí, cada una por su propia matriz:

con $\mathbf{W}^Q, \mathbf{W}^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ y $\mathbf{W}^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$, de modo que $\mathbf{Q}, \mathbf{K} \in \mathbb{R}^{T \times d_k}$ y $\mathbf{V} \in \mathbb{R}^{T \times d_v}$. Los superíndices de estas tres matrices son etiquetas del papel que hace cada una, no índices de capa: es el único sitio del curso donde un superíndice no cuenta capas, y viene de la notación del artículo. Aquí hay una sola de estas tres parejas de proyecciones; una capa de verdad lleva varias en paralelo, y eso es una lección más adelante, sobre la atención multi-headComing up · Block 5 · Lesson 4Varias cabezas, varios repartosUna posición tiene una fila y en esa fila cabe una sola manera de mirar la frase. La atención multi-head monta h atenciones en paralelo sobre la misma X, cada una con sus tres proyecciones y su propio mapa, y las junta concatenando y proyectando. Lo que cambia no son las puntuaciones sino dónde se normaliza: cada grupo de coordenadas con su propio softmax, en vez de todas sumadas antes de uno solo, que es lo que deja a una regla tapar a otra sin dejar rastro. Y de repartir la anchura, d_k = d_model/h, sale que el coste no depende de h: los mismos parámetros y las mismas multiplicaciones con una cabeza que con dieciséis..

Ahí está el «auto» del nombre, y es lo único que separa esta lección de la lección sobre la consulta, la clave y el valorBlock 4 · Lesson 6Atención como consulta, clave y valorEl mecanismo del bloque está escrito con el vocabulario de una máquina concreta —un encoder que lee, un decoder que escribe, estados con barra y sin ella— y casi nada de eso le hace falta. Aquí se reescribe entero con tres papeles: una consulta, una lista de claves contra las que se compara y una lista de valores que se mezclan. No cambia ni una ecuación, y a cambio salen tres cosas: que la clave y el valor son el mismo vector por costumbre y no por necesidad, que el gradiente le llega a cada papel por una ruta y no por dos, y que la rejilla entera cabe en un producto de tres matrices con los nombres que usa el bloque siguiente.. Allí la consulta era el estado de una red y las claves y los valores eran los estados de otra: tres listas de dos sitios distintos. Aquí las tres salen de la misma $\mathbf{X}$, y lo que cambia no es la operación —es idéntica— sino quién le pasa los argumentos.

con el softmax por filas. El divisor $\sqrt{d_k}$ es lo único de esa línea que todavía no ha justificado nadie; queda anotado y vuelvo a él al final.

Mira la forma del mapa de pesos, que es donde se nota que la entrada y la salida son la misma frase. $\mathbf{Q}\mathbf{K}^{\top}$ multiplica una matriz $T \times d_k$ por una $d_k \times T$, así que $\mathbf{A} \in \mathbb{R}^{T \times T}$ es cuadrada: una fila y una columna por posición. En el bloque anterior era $T_y \times T_x$ y sus dos ejes eran dos frases distintas; aquí los dos ejes son la misma, y por eso la casilla $\alpha_{ii}$ existe y significa algo —cuánto de sí misma se queda una posición—. Cada fila sigue sumando $1$, porque el softmax sigue normalizando una fila cada vez, y ninguna posición está excluida de ninguna fila: todas se miran a todas.

Y cuenta los pesos, que es la otra propiedad de forma que decide el resto del bloque. Las tres matrices juntan

y $T$ no aparece en esa cuenta. Multiplican filas de $\mathbf{X}$, y les da igual cuántas filas haya. Ponlo al lado de la lección sobre por qué falla el MLP (multilayer perceptron) en secuenciasBlock 3 · Lesson 1Por qué el MLP falla con secuenciasPoner los vectores de los tokens en fila mete el orden dentro de la entrada, y cobra dos cosas por ello: una longitud fija que hay que truncar o rellenar, y una primera capa partida en tantos bloques como posiciones, cada uno ajustado sólo por lo que cae en la suya. Una red que aprende «no» en las cuatro primeras posiciones acierta el 0.500 cuando «no» aparece en la quinta.: allí cada posición tenía su propio bloque de columnas $\mathbf{W}^{(1)}_t$, así que alargar la entrada obligaba a añadir pesos y cada posición aprendía por su cuenta. Estas tres matrices son las mismas para todas las posiciones y para todas las longitudes.

Sin proyecciones, la capa se copia a sí misma

Queda por defender lo que parece un rodeo: si las tres listas son lecturas de $\mathbf{X}$, ¿por qué no leerla tal cual y ahorrarse las matrices? Pon $\mathbf{Q} = \mathbf{K} = \mathbf{V} = \mathbf{X}$ y desarrolla. La puntuación de la pareja $(i,j)$ pasa a ser el producto escalar de los dos vectores de token,

y de ahí salen dos cosas, ninguna buena. La primera: $\mathbf{x}_i^{\top}\mathbf{x}_j = \mathbf{x}_j^{\top}\mathbf{x}_i$, de modo que la rejilla de puntuaciones es simétrica. Los pesos no lo son —cada fila se divide entre la suma de la suya, y esas sumas no tienen por qué coincidir—, pero el orden dentro de cada fila lo fija un número que no distingue quién pregunta de quién contesta. Un idioma sí los distingue: están necesita a llaves mucho más de lo que llaves necesita a están.

La segunda es peor. Supongamos que todos los vectores de token miden lo mismo, $\lVert \mathbf{x}_t \rVert = r$ —lo supongo yo aquí para que la cuenta salga exacta; la desigualdad de Cauchy-Schwarz hace el resto—:

con igualdad sólo cuando $\mathbf{x}_j$ apunta en la dirección de $\mathbf{x}_i$. La casilla de la diagonal es la mayor de su fila, siempre, en todas las filas a la vez, así que el softmax le da el peso más alto a la propia posición y la mezcla se inclina hacia $\mathbf{x}_i$: la capa devuelve una versión desteñida de lo que le entró. Cuánto se inclina depende de lo que midan los vectores —cuanto mayor es $r$, más se cierra la fila sobre su diagonal—, y con longitudes distintas entre sí la conclusión no cambia de forma: la fila se la lleva entonces la posición cuyo vector sea más largo, que sigue siendo una propiedad de los vectores y no de lo que la frase necesita.

Lo que remata el argumento no es ninguna de las dos, sino lo que falta en las dos. Esa capa no tiene un solo parámetro. No hay nada que ajustar, así que no hay forma de que aprenda que un verbo en plural tiene que ir a buscar un sustantivo en plural. Las tres matrices son lo que separa los tres papeles —comparar no es ser comparado, y ninguno de los dos es aportar contenido— y, de paso, lo único que la capa puede aprender.

La secuencia en los tres papeles, en NumPy

La capa entera son cinco líneas. Míralas con las formas delante, que es donde se ve que la frase entra una vez y sale por tres sitios.

Cambia T a $9$ y ejecuta otra vez: el mapa pasa a ser $9 \times 9$, la salida gana tres filas y la línea de los parámetros no se mueve. Ésa es la propiedad de forma de esta capa en una frase.

Ahora la comprobación que sostiene la sección anterior. La celda de abajo quita las proyecciones, mide lo que pasa y las vuelve a poner.

Los dos primeros bloques son la sección anterior con números. Sin proyectar, las seis filas ponen su máximo en su propia diagonal, y la línea de los vectores tres veces más largos dice de qué depende que ese máximo sea una preferencia o un monopolio: con longitud $1$ la fila reparte casi por igual, y estirando los mismos vectores el peso propio se dispara. Con dos lecturas distintas de la misma matriz eso se acaba, y las posiciones empiezan a mirarse entre ellas. El tercer bloque es otra cosa y conviene leerlo despacio, porque es de lo que trata la lección siguiente.

Comprueba tu intuición

Cinco preguntas: de dónde salen las tres listas, qué tamaño tiene el mapa, qué calcula la capa sin sus matrices, qué es simétrico y qué no, y qué se mueve cuando la frase se alarga.

Queda el divisor, y ya no se puede dejar pasar. El tercer bloque de la última celda lo mide sin tocar nada más: con consultas y claves de $d_k = 8$ el mayor peso de una fila se queda en $0.558$, y con los mismos vectores aleatorios a $d_k = 512$ sube a $0.983$. Una fila así ha dejado de ser una mezcla —es quedarse con una posición y tirar las otras cinco—, y eso tiene nombre desde la lección sobre las funciones de activaciónBlock 2 · Lesson 2Funciones de activación y no linealidadEl escalón deja pasar el signo de la preactivación y tira lo demás, y su derivada vale cero allí donde existe. Sigmoide, tangente hiperbólica y ReLU devuelven la magnitud y admiten derivada; ninguna puede ser una recta, porque encadenar funciones afines vuelve a dar una función afín.: está saturada. Dividiendo entre $\sqrt{d_k}$ el máximo también crece, pero despacio, de $0.330$ a $0.486$, y la fila sigue repartiendo.

Esa raíz no es una constante que alguien fuera probando hasta que las cosas salieron bien: se calcula, y lo que hay que calcular es cuánto crece un producto escalar a medida que se le añaden coordenadas. La cuenta ocupa entera la lección siguiente, sobre el producto interno escalado, y merece hacerse despacio, porque es donde vuelven a encontrarse dos cosas que este curso dejó en bloques distintos: la saturación que apaga un reparto y el gradiente que se desvanece cuando lo atraviesa. Ninguna de las dos necesitaba una sigmoide ni una recurrencia para aparecer.