Do Value Vectors in Deep Layers Need Context from the Residual Stream?

Resumen Técnico: Bank of Values (BoV)

Planteamiento del Problema
La arquitectura transformer, la columna vertebral de los modelos de lenguaje extensos (LLMs) modernos, se basa en capas de atención donde los vectores de consulta (query), clave (key) y valor (value) se computan a partir del flujo residual. Si bien el flujo residual acumula eficazmente el contexto para las consultas y las claves, la necesidad de contexto para los vectores de valor en las capas profundas sigue sin ser examinada exhaustivamente. La atención estándar computa los vectores de valor ($V = x_i W_V$) utilizando el flujo residual ($x_i$), lo que mezcla la información original del token con el contexto acumulado. Trabajos recientes que intentan inyectar información de tokens libre de contexto en los vectores de valor (por ejemplo, Value Residual Learning, SVFormer) han mostrado resultados mixtos, fallando a menudo al intentar aislar la contribución específica de los componentes libres de contexto o aplicándolos de manera ineficiente en todas las capas. El problema central abordado es si las capas profundas realmente requieren vectores de valor dependientes del contexto, o si se beneficiarían más preservando la información original y libre de contexto de los tokens sin la carga computacional de proyectar el flujo residual.

Metodología
Los autores proponen una investigación sistemática sobre la composición de los vectores de valor, seguida de un nuevo diseño arquitectónico llamado Bank of Values (BoV).

1. Estudio de Ablación Sistemática:
   Utilizando un modelo de 12 capas y 135 millones de parámetros, los autores realizaron experimentos a través de cuatro ejes para determinar la composición óptima de los vectores de valor en las capas profundas:

   • Aditivo vs. Sustitutivo: Comparando el reemplazo total del vector de valor estándar por una fuente libre de contexto frente a la combinación lineal de ambos.
   • Coeficientes: Probando coeficientes aprendibles (no acotados) contra coeficientes fijos.
   • Fuente de Información Libre de Contexto: Comparando una proyección de la incrustación inicial específica de la capa ($x_0 W_V$) contra un vector de valor compartido de la primera capa ($V_1$).
   • Capas Objetivo: Evaluando el impacto de aplicar estos cambios a todas las capas, a cada otra capa, o solo al último tercio de las capas.

   Hallazgos Clave de la Ablación:

   • Las capas profundas prefieren fuertemente los componentes libres de contexto ($\tilde{V}$) sobre los dependientes del contexto ($V$). En las variantes aditivas, los coeficientes aprendidos para $\tilde{V}$ fueron significativamente mayores (a menudo >10) que para $V$ (a menudo llevados a 0).
   • Sustituir $V$ por $\tilde{V}$ produce ganancias de rendimiento significativas solo en el último tercio de las capas; aplicar esto a las capas superficiales degrada el rendimiento.
   • Una vez que las capas profundas tienen acceso a un componente libre de contexto de alta calidad, añadir el componente dependiente del contexto proporciona un beneficio insignificante.
   • Los vectores libres de contexto específicos de la capa ($x_0 W_V$) superan a los vectores compartidos ($V_1$), ya que permiten que cada capa proyecte el token en un subespacio único.

2. Arquitectura Bank of Values (BoV):
   Basándose en el hallazgo de que las capas profundas se benefician de vectores de valor libres de contexto que son deterministas e independientes de la secuencia de entrada, los autores proponen BoV.

   • Mecanismo: En el último tercio de las capas, la proyección de la matriz de valor densa estándar ($x_i W_V$) se reemplaza por una tabla de búsqueda dispersa ($E_V$).
   • Implementación: Para un token con ID $i$, el vector de valor se recupera directamente como $v_i = \gamma_v E_V[i]$, donde $E_V \in \mathbb{R}^{|V| \times d}$ es una tabla de parámetros aprendibles que almacena vectores de valor precomputados para cada token del vocabulario, y $\gamma_v$ es un coeficiente de escala aprendible.
   • Eficiencia de Inferencia: A diferencia de la atención estándar, que almacena en caché los vectores de valor para todos los tokens previos (requiriendo memoria de $O(Td)$ por capa), BoV no almacena valores en caché. En su lugar, realiza una operación de recolección (gather) de $O(Td)$ desde la tabla de parámetros estática $E_V$ durante el paso hacia adelante (forward pass). La tabla $E_V$ puede ser transferida a la memoria del host, con solo las entradas requeridas siendo precargadas.

Contribuciones Clave

• Perspicacia: El artículo demuestra que los vectores de valor en las capas profundas de los transformers se benefician principalmente de la información del token original y libre de contexto, y que el componente dependiente del contexto derivado del flujo residual es en gran medida redundante en estas capas.
• Arquitectura: Introducción de Bank of Values (BoV), un método que reemplaza la computación de los vectores de valor en las capas profundas con una búsqueda estática desde una tabla específica del token.
• Eficiencia: BoV elimina la necesidad de computar y almacenar en caché los vectores de valor para las capas profundas, reduciendo el uso de memoria para la inferencia de contexto largo y disminuyendo los FLOPs por token.

Resultados
Los autores evaluaron BoV en modelos de 135M y 780M de parámetros, entrenados bajo presupuestos controlados por cómputo ($1.50 \times 10^{18}$ y $3.91 \times 10^{19}$ FLOPs, respectivamente).

• Pérdida de Validación: BoV logró una pérdida de validación menor en comparación con las líneas base de atención estándar en una partición de control del dataset ClimbMix en ambas escalas de modelos.
• Rendimiento en Benchmarks: En la escala de 780M, BoV mejoró la puntuación promedio a través de 21 benchmarks de DCLM CORE en comparación con la línea base de atención estándar.
• Comparación: BoV igualó o superó los métodos existentes que inyectan información de tokens en los vectores de valor (por ejemplo, Value Residual Learning) utilizando menos memoria y menos FLOPs. Específicamente, superó a los métodos que añaden información de tokens sobre la computación del flujo residual, los cuales aún incurren en el costo de computar y almacenar en caché los valores dependientes del contexto.

Significancia
El artículo afirma que BoV ofrece una alternativa más eficiente a la atención estándar para las capas profundas al aprovechar la idea de que estas capas no requieren contexto para computar los vectores de valor. Al persistir los vectores de valor como parámetros dispersos del modelo en lugar de computaciones dinámicas, BoV reduce la huella de memoria para la inferencia de contexto largo y el costo computacional sin sacrificar (y en algunos casos mejorando) el rendimiento del modelo. Esto desafía el paradigma por defecto de computar todos los vectores de atención a partir del flujo residual y sugiere un camino hacia arquitecturas transformer más eficientes.