From Layers to Submodules: Rethinking Granularity in Replacement-Based LLM Compression

Resumen Técnico: De Capas a Submódulos (SUBFIT)

Planteamiento del Problema
La compresión de modelos de lenguaje de gran tamaño (LLM) post-entrenamiento tiene como objetivo reducir los costes de despliegue mediante la eliminación de componentes arquitectónicos. Aunque los métodos que eliminan capas completas o bloques contiguos son efectivos, interrumpen el flujo residual, lo que suele causar una degradación brusca de la calidad. Los enfoques recientes basados en el "reemplazo" abordan esto sustituyendo los componentes eliminados por módulos ligeros ajustados. Sin embargo, los métodos existentes comparten dos restricciones de diseño restrictivas:

1. Granularidad de capa completa: Tratan a toda la capa Transformer como la unidad atómica de eliminación.
2. Selección contigua: Eliminan componentes en rangos de profundidad consecutivos.

Los autores argumentan que estas restricciones son excesivamente restrictivas. La redundancia en los transformadores preentrenados no se limita a regiones contiguas, ni se distribuye uniformemente entre las salidas de Atención y de FeedForward (FFN). En consecuencia, se necesitan diferentes estrategias para aproximar diferentes tipos de submódulos, y los componentes eliminables no necesitan agruparse en rangos de profundidad contiguos.

Metodología: SUBFIT
El artículo presenta SUBFIT (Reemplazo de residuo ajustado a nivel de submódulo), un método de compresión post-entrenamiento que opera a la granularidad de submódulos individuales de Atención y FFN. El proceso consta de dos etapas principales:

1. Selección No Contigua:

   • SUBFIT selecciona un número objetivo de submódulos de Atención y FFN de forma independiente a lo largo de la profundidad del modelo basándose en un presupuesto de dispersión $s$.
   • La selección es no contigua; los módulos se eligen independientemente por puntuación sin formar un bloque consecutivo.
   • Estrategia de Puntuación:
     • Atención: Los módulos se clasifican mediante una "puntuación de impacto" calculada sobre activaciones de calibración, que mide cuánto rota un submódulo el flujo residual y su magnitud relativa.
     • FFN: Para modelos base, se utiliza una puntuación sensible al reemplazo que combina los efectos de eliminación y el error sustituto. Para modelos con ajuste de instrucciones (instruction-tuned), se utiliza una puntuación de solo coseno para evitar el ruido de las magnitudes de flujo residual mal calibradas.
   • Ordenación: Los submódulos de Atención se seleccionan y reemplazan primero. Luego, los submódulos FFN se puntúan y seleccionan en el modelo parcialmente comprimido para tener en cuenta el desplazamiento del flujo residual inducido por el reemplazo de la Atención.

2. Reemplazo de Residuo Ajustado:

   • Cada submódulo seleccionado es reemplazado por un sustituto paramétrico ligero $\hat{F}_\ell$ que aproxima la contribución residual del submódulo original ($y = F_\ell(x)$) en lugar de su computación interna.
   • Estructura del Sustituto: El sustituto utiliza una corrección afín (puerta elemento a elemento, sesgo y centrado) más un término de interacción estructurado definido por un operador de bajo rango $U_\ell = V_\ell^\top W_\ell$.
   • Parametrización Asimétrica:
     • Atención: Cada submódulo de Atención eliminado recibe su propia base de bajo rango ($V_\ell$) y parámetros independientes, reflejando roles heterogéneos a través de las profundidades.
     • FFN: Se utiliza una base de entrada compartida $V$ en todas las capas FFN seleccionadas. Solo los factores de salida ($W_\ell$) y los términos afines siguen siendo específicos de cada capa. Este uso compartido de parámetros entre capas reduce significativamente el coste de despliegue.
   • Ajuste: Los parámetros se recuperan analíticamente mediante un ajuste de forma cerrada sobre estadísticas de calibración (mínimos cuadrados con regularización de Ridge y descomposición de autovalores), lo que no requiere optimización basada en gradientes ni ajuste fino de todo el modelo.

Contribuciones Clave

1. Identificación de Restricciones: Los autores identifican la granularidad de capa completa y la selección contigua como factores limitantes en la compresión basada en el reemplazo actual.
2. Marco de Trabajo SUBFIT: Un método novedoso que realiza la selección y el reemplazo conjuntos a nivel de submódulo (Atención y FFN) con selección no contigua.
3. Diseño Asimétrico: Una arquitectura específica donde los bypass de Atención utilizan bases independientes, mientras que los bypass de FFN comparten una base de entrada común para minimizar la sobrecarga de parámetros.
4. Evaluación Exhaustiva: Pruebas extensas en 10 LLMs (5 base, 5 con ajuste de instrucciones) y 5 niveles de dispersión (12.5% a 37.5%), demostrando mejores compromisos que cuatro líneas de base (Streamline, ReplaceMe).

Resultados Experimentales

• Perplejidad (Modelos Base): SUBFIT logra la menor degradación de la perplejidad en 9 de 10 configuraciones de modelo-dispersión. La brecha con la línea de base más fuerte se amplía significativamente bajo compresión agresiva (por ejemplo, al 25% de dispersión, SUBFIT muestra una degradación de perplejidad de 2.42× frente a 4.34× para la mejor línea de base).
• Precisión en Tareas Derivadas (Instruction-Tuned): Al 25% de dispersión, SUBFIT retiene en promedio el 84.6% de la precisión de la tarea derivada densa, superando a la línea de base más fuerte (81.6%). Es el único método que retiene más del 80% de precisión al 25% de dispersión y más del 73% al 37.5%.
• Estabilidad: SUBFIT exhibe una varianza significativamente menor en la degradación de la perplejidad en comparación con las líneas de base, que muestran inestabilidad (por ejemplo, la desviación estándar crece en órdenes de magnitud) ante una alta dispersión.
• Eficiencia: El método produce mejoras medibles en la velocidad de inferencia (1.18×–1.40× en tiempo hasta el primer token) y ahorros en el caché KV (proporcionales a la dispersión) sin requerir kernels dispersos o hardware personalizado.
• Estudios de Ablación:
  • Compensación de Residuo: Reemplazar los FFN es el factor dominante para la estabilidad; reemplazar tanto la Atención como el FFN produce los mejores resultados.
  • Base Compartida: La base compartida de FFN logra un rendimiento comparable al de las bases por capa pero con significativamente menos parámetros.
  • Sensibilidad al Rango: La calidad de FFN mejora hasta un rango de 4096, mientras que la de la Atención se satura en un rango de 256, justificando las configuraciones de rango asimétricas.

Significado y Reivindicaciones
El artículo afirma que la granularidad es un eje poco explorado en la compresión post-entrenamiento. Al pasar del reemplazo a nivel de capa al reemplazo a nivel de submódulo, SUBFIT recupera más del comportamiento del modelo denso que el reemplazo de capa completa contigua bajo el mismo presupuesto de parámetros. Los resultados sugieren que el reemplazo no contiguo y específico de cada submódulo proporciona un compromiso de calidad-eficiencia agregado más fuerte, particularmente bajo una compresión agresiva. El método está diseñado para ser barato de desplegar (ajuste de forma cerrada) y es compatible con los stacks de inferencia existentes. Los autores sugieren modestamente que este enfoque puede generalizarse a otras unidades de eliminación, como cabezales de atención o neuronas individuales, pero no reclaman una generalización inmediata a esas unidades específicas en este trabajo.