Para comprender los cambios realizados aquí, primero debemos analizar la caché de valores clave. Dentro del transformador tenemos 3 vectores que son críticos para que la atención funcione: clave, valor y consulta. Desde un nivel alto, la atención es cómo pasamos información crítica sobre los tokens anteriores al token actual para que pueda predecir el siguiente token. En el ejemplo de autoatención con una cabeza, multiplicamos el vector de consulta en el token actual con los vectores clave de los tokens anteriores y luego normalizamos la matriz resultante (la matriz resultante la llamamos patrón de atención). Ahora multiplicamos los vectores de valor con el patrón de atención para obtener las actualizaciones de cada token. Luego, estos datos se agregan a los tokens actuales incorporados para que ahora tengan el contexto para determinar lo que viene a continuación.
Creamos el patrón de atención para cada token nuevo que creamos, de modo que si bien las consultas tienden a cambiar, las claves y los valores son constantes. En consecuencia, las arquitecturas actuales intentan reducir el tiempo de cálculo almacenando en caché los vectores de clave y valor a medida que se generan en cada ronda sucesiva de atención. Este caché se llama caché de valores-clave.
Si bien arquitecturas como los modelos de transformador solo codificador y codificador-decodificador han tenido éxito, los autores postulan que la autorregresión que se muestra arriba y la velocidad que permite a sus modelos es la razón por la cual los modelos solo decodificador son los más utilizados en la actualidad.
Para comprender la arquitectura YOCO, debemos comenzar por comprender cómo establece sus capas.
Para la mitad del modelo, utilizamos un tipo de atención para generar los vectores necesarios para llenar el caché KV. Una vez que cruce a la segunda mitad, utilizará KV Cache exclusivamente para los vectores clave y de valor respectivamente, generando ahora las incrustaciones de tokens de salida.
Esta nueva arquitectura requiere dos tipos de atención: autoatención eficiente y atención cruzada. Analizaremos cada uno a continuación.
La Autoatención Eficiente (ESA) está diseñada para lograr una memoria de inferencia constante. Dicho de otra manera, queremos que la complejidad de la caché dependa no de la longitud de entrada sino del número de capas de nuestro bloque. En la siguiente ecuación, los autores extrajeron ESA, pero el resto del autodecodificador es consistente como se muestra a continuación.
Repasemos la ecuación paso a paso. X^l es nuestra incorporación de token e Y^l es una variable intermediaria utilizada para generar la siguiente incorporación de token X^l+1. En la ecuación, ESA es Autoatención eficiente, LN es la función de normalización de capa, que aquí siempre fue la norma cuadrática media (RMSNorm
), y finalmente SwiGLU
. SwiGLU
se define por lo siguiente:
Aquí swish = x*sigmoid (Wg * x)
, donde Wg es un parámetro entrenable. Luego encontramos el producto por elementos (Producto Hadamard) entre ese resultado y X*W1 antes de multiplicar ese producto completo por W2. el objetivo con SwiGLU
es obtener una función de activación que pasará condicionalmente diferentes cantidades de información a través de la capa hasta el siguiente token.
Ahora que vemos cómo funciona el autodecodificador, analicemos las dos formas en que los autores consideraron implementar la ESA.
Primero, consideraron lo que se llama Retención Cerrada. Es cierto que la retención y la autoatención son muy similares, y los autores del artículo «Retentive Network: A Successor to Transformer for Large Language Models» dicen que la diferencia clave radica en la función de activación: la retención elimina softmax, lo que permite una formulación recurrente. Utilizan esta formulación recurrente junto con la paralelización para impulsar la eficiencia de la memoria.
Para profundizar en los detalles matemáticos:
Tenemos nuestras matrices típicas de Q, K y V, cada una de las cuales se multiplica por los pesos que se pueden aprender asociados con cada matriz. Luego encontramos el producto de Hadamard entre las matrices ponderadas y el escalar Θ. El objetivo al usar Θ es crear una caída exponencial, mientras que luego usamos la matriz D para ayudar con el enmascaramiento casual (evitando que los tokens futuros interactúen con los tokens actuales) y la activación.
La retención cerrada se diferencia de la retención a través del valor γ. Aquí se utiliza la matriz Wγ para permitir que nuestra ESA se base en datos.
Sliding Window ESA introduce la idea de limitar a cuántos tokens debe prestar atención la ventana de atención. Mientras que en la autoatención normal se atienden todos los tokens anteriores de alguna manera (incluso si su valor es 0), en la ventana deslizante ESA, elegimos algún valor constante C que limite el tamaño de estas matrices. Esto significa que durante el tiempo de inferencia, la caché KV se puede reducir a una complejidad constante.
Para sumergirnos nuevamente en las matemáticas:
Tenemos nuestras matrices escaladas según sus pesos correspondientes. A continuación, calculamos la cabeza de manera similar a cómo se calcula la atención de múltiples cabezas, donde B actúa como un mapa causal y también para garantizar que solo se atiendan los tokens C.