Griffin
Griffin : combinación de recurrencias lineales cerradas con atención local para modelos de lenguaje eficientes
arXiv
Arquitectura modelo
Todos nuestros modelos contienen los siguientes componentes: (i) un bloque residual, (ii) un bloque MLP y (iii) un bloque de mezcla temporal. Si bien (i) y (ii) son iguales en todos los modelos, consideramos tres bloques de mezcla temporal: atención multiconsulta (MQA) global, MQA local (ventana deslizante) y nuestro bloque recurrente propuesto. Como parte del bloque recurrente utilizamos la Unidad Recurrente Lineal Real-Gated (RG-LRU), una novedosa capa recurrente inspirada en la Unidad Recurrente Lineal Orvieto et al., 2023b.
El bloque residual, como se muestra en la Figura 2 (a), define la estructura global de nuestros modelos y está inspirado en Transformers anteriores a la norma (Xiong et al., 2020). Después de incrustar la secuencia de entrada, la pasamos $N$ tales bloques ( $N$ que denota la profundidad del modelo), y luego aplicamos RMSNorm Zhang y Sennrich, 2019 para producir las activaciones finales. Para calcular las probabilidades de los tokens aplicamos una capa lineal final seguida de un softmax. Los pesos de esta capa se comparten con la capa de incrustación de entrada.
bloque residual
Figura 2: a) La columna vertebral principal de nuestra arquitectura de modo es el bloque residual, que está apilado $N$ veces. b) El bloque MLP cerrado que utilizamos. c) El bloque recurrente que proponemos como alternativa a Multi Query Attention (MQA). Utiliza nuestra capa RG-LRU propuesta, definida en la Sección 2.4.
El bloque residual contiene dos componentes, aplicados en orden. El primer componente toma el estado oculto. $chi$ y aplica una RMSNorm Zhang y Sennrich, 2019, seguida del bloque de mezcla temporal. Luego fusionamos la salida con una conexión de salto desde $chi$ mediante la suma. De manera similar, el segundo componente aplica RMSNorm, seguido por el bloque MLP y luego fusiona su salida con una conexión de salto desde la entrada de RMSNorm. Este bloque se ilustra en la Figura 2 (a).
bloque MLP
Usamos un bloque MLP cerrado Dauphin et al., 2017 (ilustrado en la Figura 2(b)), que crea dos ramas a partir de su entrada de dimensión $D$ . Aplicamos una capa lineal con dimensión de salida. $MD$ en cada rama, donde $M$ denota el factor de expansión. Por simplicidad usamos $M=3$ a lo largo de este trabajo. Aplicamos una no linealidad de GeLU Hendrycks y Gimpel, 2016 en una de las ramas antes de fusionarlas mediante multiplicación de elementos, similar a GeGeLU Shazeer, 2020. Sin embargo, en nuestro bloque MLP, aplicamos una capa lineal final con dimensión de salida. $D$ en las salidas de la capa GeGeLU.
Bloques de mezcla temporal
El bloque de mezcla temporal es el componente de nuestro modelo que agrega activaciones de capas ocultas en diferentes ubicaciones temporales de la secuencia. Consideramos tres bloques de mezcla temporal: MQA global Shazeer, 2019, MQA local Beltagy et al., 2020 y nuestro bloque recurrente propuesto.
Atención global multiconsulta
A menos que se indique lo contrario, utilizamos MQA en lugar de MHA para mejorar las velocidades de inferencia de nuestras líneas base de Transformer Shazeer, 2019. Utilizamos una dimensión de cabeza fija $D_{cabeza}=128$ , y fijamos el número de cabezas de atención. $H$ tal que $HD_{cabeza}=D$ . Esto requiere la dimensión del modelo. $D$ para ser un múltiplo de 128. No utilizamos ninguna incrustación posicional absoluta, pero utilizamos la incrustación de posición rotativa (RoPE) Su et al., 2021 como una incrustación posicional relativa.
Atención de ventanas correderas locales
Una de las principales desventajas de utilizar la atención global es que su complejidad computacional crece cuadráticamente en la longitud de la secuencia. Para abordar esto, varios trabajos han comenzado a adoptar la atención local Beltagy et al., 2020, también conocida como atención de ventana corredera. Permite que cada posición atienda solo una cantidad fija de tokens en el pasado. Esto no solo reduce los FLOP computacionales sino que también limita el tamaño de la caché KV al tamaño de la ventana, por lo que ya no es cuadrático en la longitud de la secuencia. Todos los demás detalles son los mismos que los del MQA global.
Bloqueo recurrente
Nuestro bloque recurrente (Figura 2 (c)) es similar al bloque GSS Mehta et al., 2022 y al bloque utilizado por Mamba Gu y Dao, 2023. Tomamos la entrada de dimensión $D$ y aplicar dos capas lineales con dimensión de salida $D_{RNN}$ en paralelo, creando dos ramas. En la primera rama, aplicamos una pequeña capa Conv1D separable, inspirada en Shift-SSM en H3 Dao et al., 2022b, con una dimensión de filtro temporal de 4. Tenga en cuenta que esta capa Conv1D es muy pequeña, con solo $4D$ parámetros. Seguimos la capa Conv1D con nuestra capa RG-LRU propuesta (definida a continuación). En la segunda rama aplicamos una no linealidad GeLU y luego fusionamos las ramas mediante multiplicación de elementos. Luego aplicamos una capa lineal final con dimensión de salida. $D$ .
Unidad recurrente lineal con puerta real (RG-LRU)
Nuestra capa RG-LRU propuesta tiene una recurrencia simple inspirada en la Unidad Lineal Recurrente (LRU) Orvieto et al., 2023b, pero incorpora un mecanismo de activación motivado por la literatura sobre RNN no lineales, en particular LSTM Hochreiter y Schmidhuber, 1997 y GRUs Chung et al., 2014. Las ecuaciones que describen la capa son las siguientes:
$$begin{align} r_t &= sigma(W_{a} x_t + b_a), & text{puerta de recurrencia} \ i_t &= sigma(W_{x} x_t + b_x), & text{ puerta de entrada} \ a_t &= a^{cr_t}, & text{} \ h_t &= a_t odot h_{t-1} + sqrt{1 - a_t^2} odot (i_t odot x_t). & text{} end{align}$$
La salida de la capa es $y_t=h_t$ , y la no linealidad $sigma$ en las ecuaciones está la función sigmoidea. El peso recurrente $a$ en la ecuación (4) es diagonal. Por tanto, todas las operaciones son por elementos. parametrizamos $a$ en la Ecuación (3) como $a=sigma(Lambda)$ , dónde $lambda$ es un parámetro que se puede aprender. Esto garantiza que $0 <= a <= 1$ , asegurando que la recurrencia sea estable. la variable $c$ es una constante de valor escalar establecida en 8. Para la estabilidad numérica, en la práctica calculamos $a^{cr_t}$ en el espacio de registro (ver Apéndice A). La capa tiene puertas tanto en la entrada $x$ y el peso recurrente $a$ . Sin embargo, ninguna de las puertas depende del estado recurrente. $ h_ {t-1} $ , lo que garantiza que el cálculo se pueda ejecutar de manera eficiente en el dispositivo. Inicializamos ambos $W_{a}$ y $W_{b}$ usando LeCun init LeCun et al., 2002. Inicializamos $lambda$ tal que $a^c$ se distribuye uniformemente entre $0.9$ y $0.999$ al inicio del entrenamiento, similar a (Orvieto et al., 2023b.). A diferencia de muchos trabajos recientes en la literatura SSM, el RG-LRU no utiliza inicialización inspirada en la teoría de polinomios ortogonales Gu et al., 2020, y tampoco se define como la discretización de un sistema continuo subyacente Gu et al., 2021a. A diferencia de la capa LRU original, no utilizamos álgebra compleja en la recurrencia. Si bien el uso de recurrencias complejas conduciría a una capa más expresiva, Orvieto et al., 2023a encontramos que las recurrencias complejas no eran beneficiosas para el modelado del lenguaje en la práctica, como también observaron Gu y Dao, 2023. (ver Apéndice B)
Comportamiento de la puerta
La puerta de entrada $yo_t$ es similar al de LSTM, que puede filtrar (o reducir) la entrada $x_t$ . Sin embargo, hasta donde sabemos, nuestra puerta de recurrencia $r_t$ es diferente de otros mecanismos de activación en la literatura. Por ejemplo, el mecanismo de selección propuesto en Mamba Gu y Dao, 2023 es comparable a la puerta de actualización de GRU que interpola $x_t$ . Su efecto sobre el estado oculto le permite restablecer su estado y olvidar cualquier información que contenga del pasado, similar a la puerta de olvido en el LSTM. Por el contrario, nuestra puerta de recurrencia puede interpolar aproximadamente entre la actualización LRU estándar de Orvieto et al., 2023a y el estado oculto anterior, lo que le permite descartar efectivamente la entrada y preservar toda la información del historial anterior (consulte el Apéndice A para obtener más detalles). ). Creemos que la función clave de esta puerta es permitir que el modelo logre una memoria superexponencial al reducir la influencia de entradas no informativas.
