QK-Normed MLA: QK normalization without full key caching

Resumen Técnico: QK-Normed MLA

Planteamiento del Problema
Los modelos de lenguaje de gran escala modernos enfrentan dos presiones contrapuestas: la necesidad de una decodificación eficiente en memoria para escenarios de contexto largo y la necesidad de estabilidad en el entrenamiento.

• Eficiencia de Memoria: La atención multi-cabezal estándar (MHA) almacena vectores completos de clave (key) y valor (value) para cada token, lo que genera altos costos de memoria y de ancho de banda. La atención de latencia multi-cabezal (MLA) aborda esto almacenando estados latentes de baja dimensión ($c_t$) en lugar de claves completas, recuperando el efecto de la clave completa mediante una up-proyección ($W^K_g$) durante la decodificación.
• Estabilidad del Entrenamiento: Las escalas incontroladas en consultas (queries) y claves (keys) pueden provocar la saturación de la atención e inestabilidades en el entrenamiento. La normalización QK (específicamente RMSNorm) estabiliza el entrenamiento al normalizar los vectores antes del producto punto. Sin embargo, la normalización RMSNorm de QK post-proyección parece incompatible con MLA porque, aparentemente, requiere la clave totalmente proyectada ($k_{t,g} = c_t W^K_g$) para cada token almacenado para calcular la estadística de normalización. Una implementación ingenua requeriría materializar las claves completas, anulando las ventajas de memoria de MLA.

Metodología
El artículo propone QK-Normed MLA, una formulación que integra la normalización QK RMSNorm post-up-proyección en la ruta de decodificación latente de MLA sin materializar claves completas. La idea central es la descomposición algebraica de RMSNorm:
$$\text{RMSNorm}(z) = z \odot \frac{\gamma}{S(z)}$$
donde $\gamma$ es un peso afín estático y $S(z)$ es una estadística escalar dinámica.

Los autores derivan una formulación "absorbida" que separa estos componentes:

1. Absorción de Pesos Estáticos: El peso estático del lado de la clave $\gamma_k$ se absorbe en la matriz de proyección del lado de la consulta. En lugar de proyectar la consulta $q$ y luego normalizar la clave, la consulta se pre-multiplica por $\gamma_k$ (y el peso estático $\gamma_q$) antes del producto punto latente.
2. Almacenamiento de Escalares Dinámicos: La estadística dinámica $S(z)$ se reduce a un único escalar de inverso-RMS ($\alpha_{t,g} = 1/S_{t,g}$) por token y por grupo de KV.
3. Diseño del Cache: El cache de decodificación retiene el estado latente original de baja dimensión $c_t$ y añade un pequeño cache escalar $A^k_{cache}$ que contiene los valores de inverso-RMS. No almacena las claves normalizadas completas.
4. Normalización por Bloques: Para mantener la eficiencia, el método normaliza los componentes de "contenido" y "RoPE" por separado. La ruta de contenido utiliza la formulación latente absorbida, mientras que la ruta RoPE utiliza vectores materializados estándar. Esto evita puntos de sincronización que complicarían las implementaciones de kernels fusionados.

Durante la decodificación, los logits de atención se calculan como:
$$\ell_{i,t,h} = (\tilde{q}_{i,h} \cdot c_t) \cdot \alpha^q_{i,h} \cdot \alpha^k_{t,g(h)}$$
donde $\tilde{q}$ incluye los pesos absorbidos, y los factores escalares se aplican como un paso de escalado fila-columna antes del softmax.

Contribuciones Clave

1. Formulación Exacta: El artículo deriva una formulación algebraicamente exacta de la normalización QK RMSNorm post-up-proyección que es compatible con la ruta de decodificación latente de MLA.
2. Diseño de Cache Eficiente: Propone un diseño de cache que preserva el cache KV latente como el estado dominante, añadiendo solo un estadístico escalar de RMS de clave por token y grupo de KV (por ejemplo, aumentando el tamaño del cache en un ~1.56% en las configuraciones probadas).
3. Validación Empírica: Los autores demuestran que este enfoque logra una pérdida de entrenamiento menor y una mejor precisión en tareas downstream en comparación con el QK clipping, introduciendo una latencia insignificante.

Resultos
Los experimentos se realizaron en modelos MLA de 400 millones de parámetros entrenados hasta 100 mil millones de tokens en GPUs H800:

• Estabilidad del Entrenamiento: QK-Normed MLA logró consistentemente una pérdida de entrenamiento menor que QK-Clip a lo largo de la ejecución de 100 mil millones de tokens. En pruebas de estrés de alta tasa de aprendizaje ($2 \times 10^{-2}$), QK-Clip divergió con gradientes NaN, mientras que QK-Normed MLA se mantuvo estable.
• Rendimiento Downstream: En una suite de tareas downstream, QK-Normed MLA mejoró la precisión en 7 de 8 tareas, elevando la puntuación promedio de 44.75 a 46.33, y redujo la perplejidad de LAMBADA de 16.28 a 14.18.
• Eficiencia de Decodificación: Los benchmarks en contextos que van desde 4k hasta 256k tokens mostraron que el método añade menos del 2% de latencia (promedio de 1.40%). El cache escalar aumentó el tamaño total del cache en solo un ~2.8% por capa.

Significancia y Reivindicaciones
El artículo afirma que la aparente incompatibilidad entre la normalización QK y MLA es un artefacto de implementación más que una restricción arquitectónica. Al demostrar que los pesos estáticos pueden absorberse y las estadísticas dinámicas reducirse a escalares, los autores permiten la normalización QK como un mecanismo de estabilización práctico para modelos MLA sin sacrificar la eficiencia de memoria de la decodificación latente.

Los autores mantienen la modestia respecto al alcance de sus afirmaciones:

• Los resultados se basan en modelos de 400 millones de parámetros; no pretenden establecer un comportamiento a escala de frontera para modelos de miles de millones de parámetros, aunque la equivalencia algebraica se mantiene independientemente de la escala.
• El análisis de la escala de logits se presenta como un diagnóstico a nivel de mecanismo más que como una teoría de optimización completa.
• La velocidad práctica depende de la integración de kernels (fusionar los factores escalares en la ruta de atención) para evitar el tráfico de memoria innecesario.