Resumen Técnico: RAP (Poda Adaptativa en Tiempo de Ejecución para Inferencia de LLM)

Enunciado del Problema

Los Modelos de Lenguaje Grande (LLM) enfrentan barreras significativas para su despliegue debido a sus enormes huellas computacionales y de memoria. Aunque las técnicas de compresión de modelos (poda, cuantización, destilación) ofrecen soluciones, los métodos existentes adolecen de limitaciones críticas:

1. Heurísticas Estáticas: La mayoría de las estrategias de poda actuales se basan en heurísticas fijas y de un solo paso que asumen cargas de trabajo estáticas. No logran adaptarse a las variaciones de memoria en tiempo de ejecución causadas por solicitudes de usuarios fluctuantes (por ejemplo, tamaños de lote y longitudes de secuencia variables).
2. Negligencia de la Caché KV: La poda tradicional se centra principalmente en reducir los pesos estáticos del modelo (parámetros). Sin embargo, en la inferencia autoregresiva, la memoria de la caché de claves y valores (KV) crece linealmente con el tamaño del lote y la longitud de la secuencia, a menudo superando con creces la memoria de los parámetros estáticos. Las ratios de poda fijas no pueden acomodar el cambio dinámico de cuellos de botella dominados por parámetros a cuellos de botella dominados por la caché KV.
3. Varianza a Nivel de Sistema: Los entornos reales de borde y nube experimentan disponibilidad de memoria estocástica debido a aplicaciones concurrentes e interferencias, lo que invalida las suposiciones de presupuesto fijo.

El desafío central abordado es cómo seleccionar una política de poda óptima para LLM que se adapte dinámicamente a cargas de trabajo de solicitudes heterogéneas y variables en el tiempo, satisfaciendo al mismo tiempo presupuestos de memoria fluctuantes sin degradar la calidad de la salida.

Metodología: El Marco RAP

Los autores proponen RAP (Poda Adaptativa en Tiempo de Ejecución), un marco de poda elástico impulsado por Aprendizaje por Refuerzo (RL) que ajusta dinámicamente las estrategias de compresión basándose en estados del sistema en tiempo real.

1. Importancia Secuencial Codiciosa (GSI)

Para abordar las limitaciones de la poda de "un solo paso" (que ignora las dependencias intercapas y conduce a una degradación del rendimiento), RAP emplea la Importancia Secuencial Codiciosa (GSI).

• Mecanismo: En lugar de puntuar todos los bloques una sola vez, GSI elimina iterativamente el bloque cuya exclusión resulta en el deterioro mínimo del rendimiento del modelo (medido por la perplejidad).
• Recalibración: Después de cada eliminación, las puntuaciones de importancia de los bloques restantes se reevalúan dentro de la nueva arquitectura contraída. Este proceso paso a paso controla la acumulación de errores y tiene en cuenta la sensibilidad dinámica de las capas a medida que cambia la estructura del modelo.

2. Formulación de Aprendizaje por Refuerzo

RAP plantea la selección de la política de poda como un problema de toma de decisiones secuencial (Proceso de Decisión de Markov) donde un agente de RL aprende a equilibrar la eficiencia de la memoria y la fidelidad del modelo.

• Estado ($S_t$): El agente observa un estado compuesto en cada paso de inferencia $t$:
  • Características de la Solicitud: Tamaño del lote ($R_{bs}$) y longitud de la secuencia ($R_{sql}$).
  • Configuración del Modelo: Puntuaciones de importancia de todos los bloques de Atención Multi-Cabeza (MHA) y Red Neuronal de Alimentación Directa (FFN) derivadas de GSI.
  • Estado del Sistema: Memoria del sistema disponible actual ($Sys_{avail}$) y sobrecarga de memoria estimada ($Sys_{req}$) para la acción propuesta.
• Acción ($A_t$): El agente selecciona una política de poda decidiendo si retener o eliminar bloques específicos de MHA y FFN. Para gestionar la explosión combinatoria del espacio de acciones ($2^{2N}$), la decisión se descompone en selecciones secuenciales de un solo bloque.
• Recompensa ($R_t$): Una función de recompensa unificada equilibra dos objetivos:
  • Perplejidad ($R_{ppl}$): Fomenta la preservación del rendimiento de modelado del lenguaje.
  • Penalización de Memoria ($R_{mem}$): Penaliza el consumo máximo de memoria durante la inferencia.
  • La recompensa es una suma ponderada: $R_t = \sum (At)_i (\alpha R_{ppl, i} - \beta R_{mem, i})$, donde los hiperparámetros $\alpha$ y $\beta$ ajustan la compensación entre precisión y memoria.

3. Flujo de Ejecución

1. El agente de RL observa el entorno en tiempo real (longitud de entrada, lote, disponibilidad de memoria).
2. Selecciona una acción de poda (qué bloques eliminar) basándose en el estado actual.
3. El modelo base ejecuta la inferencia en la arquitectura podada.
4. El agente recibe una recompensa basada en el uso de memoria resultante y la perplejidad, actualizando su política para optimizar decisiones futuras.

Contribuciones Clave

1. Marco Adaptativo en Tiempo de Ejecución: RAP es el primer marco que considera conjuntamente los pesos del modelo y la memoria de la caché KV "sobre la marcha", ajustando dinámicamente los niveles de dispersión basándose en la carga de trabajo instantánea y el estado del dispositivo.
2. Política Impulsada por RL: Los autores formulan la poda como un problema de RL, entrenando a un agente para aprender una política óptima que equilibra el ahorro de memoria con la fidelidad de generación, superando las líneas base de heurísticas estáticas.
3. GSI para Poda Estructural: La introducción de la Importancia Secuencial Codiciosa aborda el problema de la dependencia intercapa en la poda de transformadores, proporcionando un método más robusto para determinar la importancia de los bloques bajo compresión agresiva.
4. Validación Empírica: Experimentos extensivos demuestran que RAP supera consistentemente a las líneas base de poda estática más avanzadas (por ejemplo, LLM-Pruner, SliceGPT, ShortGPT) en diversos modelos (Llama2, Llama3, Qwen) y restricciones de memoria.

Resultados Experimentales

Los experimentos se realizaron en Llama2-7B, Llama3-8B, Qwen1.5-7B y Qwen2.5-7B bajo presupuestos de memoria unificada del 80% y 60% (que cubren tanto parámetros como caché KV).

• Capacidad de Generación: RAP mostró la menor deriva de perplejidad. En Llama2-7B con un presupuesto del 80%, RAP aumentó la perplejidad solo en +6.33 (WikiText2) y +21.47 (PTB), superando significativamente al siguiente mejor método por márgenes de 22.95 y 146.93, respectivamente.
• Tareas de Descenso: En el razonamiento de sentido común en cero disparos (zero-shot), RAP preservó el 86% de la precisión del modelo denso bajo un presupuesto del 80%, superando a las líneas base líderes en +8.2% (Llama2-7B) y +11.5% (Llama3-8B).
• Resiliencia: Bajo restricciones de memoria severas del 60%, RAP fue el único método que retuvo más del 60% del rendimiento del modelo denso, mientras que otros métodos sufrieron una degradación catastrófica.
• Estudios de Ablación:
  • Eliminar GSI (usando puntuación estática de un solo paso) causó una caída del 13.17% en la precisión promedio de sentido común.
  • Eliminar el agente de RL (usando eliminación aleatoria de bloques) resultó en un aumento masivo de la perplejidad (+301.71 en WikiText2) en comparación con la política aprendida, destacando la necesidad de una toma de decisiones adaptativa.
• Sobrecarga: El controlador de RL agrega una sobrecarga de tiempo de ejecución insignificante (<1% de aumento en la latencia) debido a su arquitectura compacta (18K parámetros), mientras que el cálculo de GSI fuera de línea es un costo único amortizado durante el despliegue.

Significado y Afirmaciones

El artículo afirma que RAP cierra la brecha entre las técnicas de compresión estáticas y los escenarios de despliegue dinámicos del mundo real. Al alejarse de las heurísticas de "talla única", RAP permite:

• Despliegue Escalable en el Borde: Permite que los LLM se ajusten en el dispositivo bajo restricciones de memoria ajustadas y fluctuantes, podando pesos sustanciales cuando es necesario, mientras retiene la capacidad completa cuando los recursos lo permiten.
• Eficiencia sin Pérdida de Calidad: Logra un ahorro de memoria superior y una inferencia más rápida con un impacto mínimo en la calidad de la salida, navegando efectivamente por las compensaciones entre cuellos de botella pesados en parámetros y cuellos de botella pesados en caché KV.
• Generalización: Se demuestra que el marco es agnóstico a la arquitectura, funcionando bien en diferentes familias de modelos (Llama, Qwen) y generaciones.

Los autores posicionan a RAP como un mecanismo de compresión flexible y consciente del contexto que hace viable la inferencia de LLM en entornos heterogéneos donde las estrategias fijas fallan.