GRASPrune: Global Gating for Budgeted Structured Pruning of Large Language Models

Resumen Técnico: GRASPrune

1. El Problema

Los Modelos de Lenguaje Grande (LLMs) son costosos de implementar debido a la gran cantidad de parámetros, el alto costo computacional de la atención y, crucialmente, el uso de memoria del caché KV (Key-Value), que crece linealmente con la longitud de la secuencia.
Aunque existen métodos de compresión estructural (podado), las soluciones actuales presentan tres limitaciones principales:

1. Criterios separados: Suelen podar canales de la red Feed-Forward (FFN) y grupos de cabezas de atención (KV) con criterios independientes, ignorando que compiten por el mismo presupuesto de despliegue.
2. Asignación de presupuesto rígida: Muchos métodos asignan tasas de retención fijas por capa o usan programas dependientes de la profundidad, en lugar de aprender una asignación global óptima.
3. Desconexión en la optimización: Las pipelines existentes suelen aprender puntuaciones de importancia sin restricciones y aplicar el presupuesto de hardware solo al final (post-hoc). Esto debilita la conexión entre la arquitectura final y el objetivo de entrenamiento, a menudo requiriendo reentrenamiento costoso o reconstrucción.

2. Metodología: GRASPrune

GRASPrune es un marco de podado estructurado que se aplica después del preentrenamiento (post-training) y opera bajo un único presupuesto global de recursos.

Componentes Clave:

• Presupuesto Global Unificado:

  • El método trata el podado de canales FFN y grupos de cabezas KV como un único problema de selección discreta.
  • Define un costo unitario para cada tipo de estructura (los canales FFN tienen costo 1; los grupos KV tienen un costo $\alpha$ basado en la huella de parámetros, considerando la arquitectura GQA - Grouped Query Attention).
  • Se optimiza una sola máscara binaria global $z$ sujeta a la restricción $\sum c_i z_i \leq B$, donde $B$ es el presupuesto total.

• Entrenamiento con Estimador Straight-Through (STE) Proyectado:

  • En lugar de aprender puntuaciones libres y luego filtrarlas, GRASPrune enforza la viabilidad del presupuesto en cada paso del bucle de optimización.
  • Mecanismo:
    1. Se asignan puntuaciones continuas (gates) a cada unidad podable.
    2. En el paso de forward, las probabilidades se proyectan a una máscara dura (binaria) que satisface estrictamente el presupuesto global mediante un algoritmo de ordenación y selección (proyección).
    3. En el paso de backward, se utiliza el STE para propagar gradientes a través de las probabilidades suaves, manteniendo los pesos del modelo base (backbone) congelados.
  • Esto asegura que las puntuaciones se aprendan bajo la misma restricción que define la arquitectura final, evitando que el algoritmo favorezca unidades "baratas" pero menos importantes solo por su costo.

• Calibración de Escala (Scaling Calibration):

  • Una vez fijada la máscara, se introducen factores de escala escalares ($\gamma$) para las unidades retenidas.
  • Estos factores se optimizan en un conjunto de calibración pequeño (sin tocar los pesos del modelo) para mitigar el desajuste de escala causado por la eliminación de unidades.
  • Finalmente, estos factores se "pliegan" (fold) en los pesos podados, resultando en un checkpoint denso más pequeño sin parámetros adicionales durante la inferencia.

3. Contribuciones Clave

1. Formulación Global: Unifica la selección de canales FFN y cabezas KV bajo un solo presupuesto, reconociendo sus diferentes costos y capacidades representacionales.
2. Optimización con Restricción Activa: Propone un bucle de entrenamiento con proyección STE que mantiene la viabilidad del presupuesto en cada iteración, evitando la desconexión entre el aprendizaje de importancia y la selección final.
3. Checkpoint Densificado Eficiente: Produce un modelo final más pequeño y denso mediante la calibración de escalas, eliminando la necesidad de fine-tuning completo o reconstrucción compleja.
4. Eficiencia Extrema: El proceso de podado es extremadamente rápido (aprox. 6 minutos en una GPU A100 para LLaMA-2-7B) y requiere solo un pequeño conjunto de datos de calibración no etiquetado.

4. Resultados Experimentales

Los experimentos se realizaron principalmente en LLaMA-2-7B, comparando con métodos de vanguardia como LLM-Pruner, SliceGPT y FLAP.

• Rendimiento en Perplejidad (WikiText-2):
  • Con un 50% de retención de parámetros, GRASPrune logra una perplejidad de 12.18, superando significativamente a los baselines (LLM-Pruner: 312.9, SliceGPT: 26.1, FLAP: 24.8).
  • Incluso en un escenario de alta compresión (40% de retención), mantiene una perplejidad de 16.65, mientras que otros métodos colapsan (perplejidad > 2000 para LLM-Pruner).
• Precisión en Tareas Zero-Shot:
  • Mantiene una precisión promedio competitiva en cinco benchmarks (ARC, HellaSwag, PIQA, etc.), superando consistentemente a los métodos comparados, especialmente a medida que aumenta la tasa de podado.
• Eficiencia de Despliegue:
  • Memoria KV: La reducción de cabezas KV disminuye proporcionalmente el uso de memoria de la caché KV, permitiendo secuencias más largas o lotes (batch sizes) más grandes bajo el mismo presupuesto de memoria.
  • Throughput: Se observa una mejora en el rendimiento de generación (tokens/segundo), que aumenta con el tamaño del lote y la agresividad del podado.
• Generalización: El método se ha validado en otras arquitecturas (LLaMA-3.1-8B, Qwen3-8B/14B), mostrando consistencia.

5. Significado e Impacto

GRASPrune demuestra que la compresión de LLMs no requiere reentrenamiento masivo ni fine-tuning costoso. Al tratar el podado estructurado como un problema de optimización global con restricciones activas, logra:

• Mejor preservación de capacidades: Mantiene la capacidad de modelado de lenguaje y razonamiento básico mucho mejor que los métodos que separan la estructura o aplican el presupuesto al final.
• Viabilidad de despliegue: Genera modelos más pequeños y eficientes que reducen tanto la latencia como el uso de memoria (especialmente crítico para el caché KV en contextos largos).
• Simplicidad: Ofrece una solución práctica que puede integrarse en stacks de servicio estándar sin parámetros extra en inferencia.

En conclusión, GRASPrune establece un nuevo estándar para el podado estructurado post-entrenamiento, demostrando que una asignación global inteligente de recursos es más efectiva que las estrategias locales o secuenciales tradicionales.