Resumen Técnico: LLMs Eficientes con AMP: Poda de Cabezales de Atención y MLP

Declaración del Problema

Los Modelos de Lenguaje Extensos (LLMs) han alcanzado un rendimiento de nivel humano en tareas cognitivas, pero sufren de recuentos de parámetros extensos, lo que conduce a altos costos computacionales y una inferencia lenta. Estos factores dificultan su despliegue en entornos con recursos limitados. Si bien la poda es una estrategia prometedora para reducir el tamaño del modelo, los métodos existentes enfrentan compromisos significativos:

• Poda Estructurada: Elimina componentes enteros (por ejemplo, capas, cabezales), pero a menudo requiere un ajuste fino computacionalmente costoso o provoca la degradación del rendimiento si se eliminan componentes críticos.
• Poda Semi-estructurada: Requiere hardware especializado (por ejemplo, patrones de dispersión específicos de GPU como 2:4) para lograr aceleraciones, lo que limita su despliegue en hardware de consumo.
• Poda No Estructurada: Crea matrices dispersas irregulares que no logran acelerar la inferencia en hardware de propósito general sin optimizaciones de software especializadas.
• Limitaciones Generales: Muchas técnicas existentes dependen de métricas basadas en la pérdida, expansiones de Taylor o heurísticas complejas que son computacionalmente costosas, inflexibles o requieren hardware especializado. Existe la necesidad de un método agnóstico al hardware, eficiente, que comprima los LLMs con una pérdida mínima de rendimiento y aceleraciones de inferencia tangibles.

Metodología: AMP (Poda de Cabezales de Atención y MLP)

Los autores proponen AMP, un método de poda estructurada que elimina cabezales de atención menos críticos dentro de la Atención de Multicabezal (MHA) y neuronas dentro de las Perceptrones Multicapa (MLP). A diferencia de los métodos que dependen únicamente de las magnitudes de los pesos o de los cambios en la pérdida, AMP evalúa la importancia estructural proyectando los datos de entrada sobre los pesos para evaluar la contribución de cada componente al flujo residual.

Mecanismo Central

El método opera bajo el principio de que la información fluye a través del flujo residual. AMP calcula una puntuación de importancia para cada componente basada en la magnitud de su contribución de activación a la salida final de la capa.

1. Componente MHA:

   • La salida de MHA se descompone en una suma de contribuciones de cabezales individuales ($O = \sum h_n W_n$).
   • La importancia del $n$-ésimo cabezal ($I_n$) se define como la norma $\ell_1$ del producto de la salida del cabezal ($h_n$) y su matriz de proyección de salida correspondiente ($W_n$):
     $$I_n = \|h_n W_n\|_1$$
   • Este enfoque tiene en cuenta las interacciones no lineales entre los pesos y las entradas, a diferencia de la poda simple por norma de peso.

2. Componente MLP:

   • El método se dirige al mecanismo SwiGLU utilizado en modelos como LLaMA, específicamente a las proyecciones Up y Gate.
   • Dado que SwiGLU utiliza multiplicación elemento a elemento, las neuronas deben podarse en pares de ambas proyecciones para mantener la integridad arquitectónica.
   • La importancia del $m$-ésimo par de neuronas ($I_m$) se calcula como el promedio de la norma $\ell_1$ del producto elemento a elemento de las proyecciones Gate y Up a través de una muestra de tokens de entrada:
     $$I_m = \frac{1}{S} \sum_{s=1}^{S} | \text{SiLU}((x_s W_{\text{Gate}})_m) \cdot (x_s W_{\text{Up}})_m |$$

3. Proceso de Poda:

   • Se utiliza un pequeño conjunto de calibración (por ejemplo, 50 muestras aleatorias de Alpaca) para extraer valores de activación promedio.
   • Los componentes (cabezales y neuronas) se clasifican por importancia en todas las capas.
   • Para lograr una tasa de compresión objetivo $c$, el $c\%$ de los cabezales de atención y el $c\%$ de las neuronas MLP inferiores se eliminan uniformemente en todas las capas.
   • Post-entrenamiento: El modelo podado se somete a un proceso de ajuste fino ligero utilizando Adaptación de Bajo Rango (LoRA) para recuperar el rendimiento.

Contribuciones Clave

• Métrica de Importancia Novedosa: AMP introduce una métrica basada en datos que proyecta las activaciones de entrada sobre los pesos para evaluar la importancia del componente, diferenciándose de las heurísticas tradicionales de solo peso o basadas en la pérdida.
• Aceleración Agnóstica al Hardware: El método logra aceleraciones prácticas de inferencia (hasta 1.25$\times$) sin requerir hardware especializado o núcleos de tensores dispersos, lo que lo hace adecuado para GPUs de consumo.
• Eficiencia: El proceso de decisión de la poda toma solo minutos en una GPU de consumo, y el ajuste fino posterior requiere recursos insignificantes en comparación con el reentrenamiento completo.
• Validación de Coherencia: Los autores proponen y ejecutan un experimento de "verificación de coherencia", demostrando que eliminar los componentes más importantes identificados por AMP causa caídas catastróficas en el rendimiento, validando la capacidad de la métrica para distinguir estructuras críticas de las no críticas.
• Flexibilidad: El método es aplicable a diferentes familias de LLM (LLaMA, Phi) y permite tasas de compresión flexibles sin estar ligado a patrones de dispersión fijos (por ejemplo, 2:4).

Resultados Experimentales

Los autores evaluaron AMP en los modelos LLaMA 7B, LLaMA-2 7B, Phi-1.5 y Phi-2 utilizando benchmarks de cero disparos (WinoGrande, HellaSwag, ARC-e/c, PIQA) y perplejidad (WikiText2).

• Rendimiento:
  • LLaMA 7B: Con una tasa de poda del 20%, AMP logró una precisión promedio del 64.52%, superando a métodos de vanguardia como LLM-Pruner (59.71%) y Shortened LLaMA (63.46%).
  • LLaMA-2 7B: Con una tasa de poda del 30%, AMP logró un 61.02% de precisión promedio, superando a SliceGPT (51.50%) y LLM Surgeon (59.03%).
  • Mejora: AMP superó a las técnicas de poda estructurada existentes hasta en 1.49 puntos porcentuales en la precisión promedio de las tareas.
• Velocidad de Inferencia:
  • AMP logró una aceleración de 1.25$\times$ para LLaMA 7B y de 1.19$\times$ para LLaMA-2 7B.
  • Esto supera a Shortened LLaMA (1.13$\times$) e iguala a métodos semi-estructurados como Wanda (1.24$\times$) sin requerir hardware especializado.
• Perplejidad: Los modelos podados y ajustados mostraron una perplejidad significativamente menor en WikiText2 en comparación con las versiones no ajustadas, con la brecha ampliándose a tasas de poda más altas.
• Estudios de Ablación:
  • Verificación de Coherencia: Revertir el orden de la poda (eliminando los componentes más importantes) resultó en caídas de precisión de >25 puntos porcentuales, confirmando la validez de la métrica.
  • Aislamiento de Componentes: Podar solo cabezales MHA o solo neuronas MLP resultó en un rendimiento significativamente peor que podar ambos simultáneamente, indicando la necesidad de un enfoque equilibrado.

Significado y Reivindicaciones

El artículo afirma que AMP aborda el compromiso crítico de la compresión de modelos ofreciendo un método que es:

1. Efectivo: Supera a los métodos actuales de poda estructurada de vanguardia en precisión manteniendo altas tasas de compresión.
2. Eficiente: Requiere mínimos recursos computacionales tanto para la decisión de poda (minutos) como para la recuperación del rendimiento (ajuste fino ligero de LoRA).
3. Práctico: Permite aceleraciones de inferencia tangibles en hardware estándar, facilitando el despliegue en entornos con recursos limitados.
4. Sostenible: Alineado con los principios de IA Verde, el método reduce los costos operativos en aproximadamente un 17.39% y las emisiones de carbono en un 19.97% debido a la reducción de la latencia de inferencia y el recuento de parámetros.

Los autores enfatizan que su método identifica y elimina con éxito las estructuras no esenciales, permitiendo el despliegue de LLMs altamente optimizados sin la necesidad de infraestructura de hardware especializada. También destacan el papel crítico del post-entrenamiento (ajuste fino) para restaurar el rendimiento del modelo después de la poda.