Mecanismo de Atención MFA: Reducción de KV Cache en LLMs

Introducción a la Problemática del KV Cache en LLMs

El uso creciente de modelos de lenguaje grandes (LLMs) y los nuevos paradigmas en inferencia han resaltado la necesidad de una inferencia eficiente a gran escala. Un cuello de botella importante es el cache Key-Value (KV) dentro de los mecanismos de atención tradicionales. Este cache crece linealmente con el tamaño del lote y la longitud de la secuencia, convirtiéndose en un 'acaparador de memoria' que dificulta la escalabilidad de los LLMs.

Aunque han surgido variantes como MQA, GQA y MLA para abordar este problema, a menudo tienen dificultades para mantener el rendimiento bajo restricciones estrictas de memoria o introducen complejidades que plantean desafíos de ingeniería y problemas de compatibilidad.

Multi-matrix Factorization Attention (MFA): Una Nueva Solución

Un reciente estudio de investigadores de Stepes, la Universidad de Tsinghua y otras instituciones presenta una novedosa arquitectura de mecanismo de atención: Multi-matrix Factorization Attention (MFA) y su variante MFA-Key-Reuse (MFA-KR). Este mecanismo reduce significativamente el coste de la inferencia del modelo de lenguaje a la vez que mejora el rendimiento.

MFA y MFA-KR no solo superan a MLA en rendimiento, sino que también igualan el rendimiento tradicional de MHA, reduciendo el uso de KV Cache hasta en un 93.7%. MFA está diseñado para ser simple, fácil de reproducir, poco sensible a los hiperparámetros y compatible con varios métodos de Pos-embedding.

Enfoque y Análisis de MFA

El equipo de investigación analizó el diseño general y la capacidad de los mecanismos de atención, identificando dos dimensiones críticas relacionadas con la capacidad. Este análisis condujo al desarrollo de nuevos métodos analíticos y principios de diseño.

• Introdujeron el concepto de Generalized Multi-Head Attention (GMHA) como un marco unificador para entender diferentes variantes de MHA.
• El equipo también exploró la computación y el almacenamiento de key-values desde una perspectiva de inferencia y examinó la capacidad del modelo desde una perspectiva de descomposición.
• Fully Parameterized Bilinear Attention (FPBA) se estableció como el límite superior teórico de rendimiento. Descubrieron que MHA y sus variantes son descomposiciones de bajo rango de FPBA.

Comparación con MQA y MLA

El análisis se centró en dos esquemas de mejora representativos: Multi-Query Attention (MQA) y Multi-Head Latent Attention (MLA).

• MQA utiliza una estrategia de compartición de parámetros más agresiva, donde todas las cabezas de atención comparten el mismo conjunto de parámetros key-value. Esto reduce el uso de memoria, pero puede afectar la expresividad del modelo.
• MLA introduce un espacio latente compartido para la compresión de parámetros, pero la potencia expresiva real está limitada por la dimensión más pequeña. Esto significa que aumentar las dimensiones intermedias no mejora significativamente el rendimiento.

Innovaciones Clave de MFA

El desarrollo de MFA se basó en el objetivo de crear un mecanismo de atención que minimizara el consumo de recursos y se acercara a los límites teóricos de rendimiento. El diseño de MFA incorpora tres innovaciones clave:

1. Aumentar significativamente el número y la dimensión de las cabezas de atención para maximizar la capacidad del modelo.
2. Emplear una estrategia agresiva de descomposición de bajo rango para mantener la eficiencia de los parámetros mientras se amplía el número y las dimensiones de las cabezas de atención.
3. Utilizar un diseño de cabeza de key-value única para mantener el consumo de memoria al mínimo, incluso con una mayor complejidad del modelo.

Medición y Comparación de Capacidad

Para analizar más a fondo MFA y otros mecanismos de atención, el equipo introdujo dos métricas clave:

• Total Effective Rank (TER): El producto del número de cabezas de atención y el rango de factorización por cabeza (FRH).
• Shared Latent Subspace Dimension (SLSD): La dimensión del espacio oculto compartido por todas las cabezas de atención.

MFA logra un SLSD y un TER más altos en comparación con MQA. En comparación con MLA, MFA logra un tamaño de caché KV más pequeño y un TER más alto con presupuestos de parámetros similares, manteniendo un SLSD comparable. En comparación con el MHA tradicional, MFA tiene un TER más alto, aunque su SLSD es menor.

Resultados Experimentales

Se llevaron a cabo extensos experimentos para evaluar el rendimiento de la nueva arquitectura a escalas mayores, probando modelos que van desde 1B hasta 7B parámetros y datos de entrenamiento de 10B a 1T.

• MFA demostró capacidades de escalado comparables al MHA tradicional, manteniendo un excelente rendimiento incluso a escalas más grandes.
• Aunque MFA-KR mostró un rendimiento ligeramente inferior, su tendencia de escalado se alineó con MHA.
• Las ventajas de ahorro de memoria de MFA y MFA-KR continuaron expandiéndose con el tamaño del modelo, con MFA logrando un ahorro de memoria del 87.5% y MFA-KR reduciendo el uso de memoria al 6.25% en la escala más grande.

Estudios de Ablación

Los estudios de ablación validaron la eficacia de MFA y MFA-KR. Sus ventajas de rendimiento también se confirmaron en varios métodos de codificación posicional convencionales.

Perspectivas

MFA ofrece mejoras significativas con un diseño simple, abordando eficazmente el cuello de botella de memoria en la inferencia de LLM sin añadir complejidad de ingeniería adicional. Se integra a la perfección en el ecosistema Transformer existente, acelerando la aplicación de LLMs en diversos escenarios.