Resumen Técnico: RAP (RoPE-Aligned Pruning) para la Compresión del KV-Cache

1. Planteamiento del Problema

La inferencia de contexto largo en los Modelos de Lenguaje de Gran Escala (LLMs) se ve cada vez más limitada por los requisitos de memoria y computación del caché de Clave-Valor (KV). A medida que la longitud del contexto escala, el caché KV domina el uso de la memoria de la GPU y los costos de computación de la atención. Aunque la poda estructurada ofrece un camino directo para comprimir el caché eliminando canales menos útiles de los pesos de proyección $W_k$ y $W_v$, los métodos existentes enfrentan una incompatibilidad crítica con el Posicionamiento Rotatorio (RoPE).

Los LLMs modernos aplican RoPE después de las proyecciones de Consulta/Clave (Query/Key), rotando pares de dimensiones de características. Los métodos de poda por canal estándar (por ejemplo, KVPruner) suelen eliminar canales individuales sin considerar esta estructura de emparejamiento. Este "orfandad" de los semipares rompe la semántica rotacional de RoPE, provocando un colapso catastrófico pero silencioso en la precisión del modelo. Además, los métodos de descomposición de bajo rango existentes (por ejemplo, SVD-LLM, PaLU) que intentan comprimir los pesos a menudo no pueden fusionar sus matrices de expansión con la proyección de Consulta ($W_q$) debido a la no conmutatividad de RoPE y la multiplicación de matrices, lo que obliga a una costosa reconstrucción paso a paso de los vectores de Clave de dimensión completa durante la decodificación.

2. Metodología: RoPE-Aligned Pruning (RAP)

Los autores proponen RAP, un marco de poda estructurada diseñado específicamente para respetar la estructura de emparejamiento impuesta por RoPE.

Principio Central

En lugar de podar canales individuales, RAP restringe la granularidad de la poda a pares alineados con RoPE. Al eliminar pares de dimensiones completos, el método asegura que la operación de rotación permanezca matemáticamente válida y que la semántica posicional se preserve.

El Pipeline de RAP

1. Puntuación de Pares de RoPE:

   • Para los pesos de la proyección de Clave ($W_k$), las columnas se agrupan en pares de RoPE de acuerdo con la estrategia de emparejamiento específica del modelo (ya sea adyacente o separada a la mitad).
   • La importancia se mide utilizando puntuaciones de información de Fisher calculadas en un pequeño conjunto de datos de calibración. Para un par $p = (j, j')$, la puntuación $\sigma_p$ es la magnitud de la raíz de Fisher agregada de los dos canales.
   • Para la proyección de Valor ($W_v$), que carece de emparejamiento RoPE, se calculan las puntuaciones de Fisher de canales individuales.

2. Asignación de Presupuesto:

   • Un ratio de retención global $\rho$ se distribuye a través de capas y cabezales (heads).
   • Adaptativo por Capa: Se asignan presupuestos a las capas basadas en su sensibilidad (puntuaciones de Fisher agregadas), otorgando un presupuesto mayor a las capas más sensibles.
   • Uniforme por Cabezal: Dentro de una capa, todos los cabezales reciben el mismo número de pares retenidos para facilitar las operaciones GEMM por lotes.
   • Equilibrado K/V: Para evitar que la dimensión de la Clave sea aplastada (lo que perjudica el rendimiento en contextos largos), RAP asigna el mismo presupuesto de canales por cabezal tanto a $K$ como a $V$. Específicamente, si $K$ retiene $m$ pares, $V$ retiene $2m$ canales.

3. Construcción y Fusión:

   • Se construye una matriz de selección binaria, que preserva los pares, $E_k$ para mapear los pares retenidos de vuelta a sus índices originales.
   • Crucialmente, debido a que $E_k$ preserva la estructura de pares, conmuta con la operación RoPE: $Rope(X \tilde{W}_k E_k) = Rope(X \tilde{W}_k) E_k$.
   • Esta conmutatividad permite que la máscara de poda se fusione offline en la proyección de Consulta: $\tilde{W}_q = W_q E_k^\top$.
   • Resultado: El modelo se ejecuta en dimensiones podadas, y el caché KV almacena solo las claves latentes de baja dimensión. No se requiere reconstrucción en línea.

4. Recuperación de la Precisión:

   • Para recuperar la precisión perdida durante la poda, RAP emplea una actualización ligera de LoRA (Low-Rank Adaptation) combinada con Destilación de Conocimiento (KD).
   • Se entrena un adaptador LoRA de bajo rango ($r=8$) en el modelo podado para alinear sus salidas con el modelo profesor denso, y luego se fusiona antes del despliegue.

3. Contribuciones Clave

• Identificación de la Incompatibilidad con RoPE: El artículo destaca que la poda de canales estándar rompe los pares de RoPE, causando un colapso de la precisión, y que este fallo suele ser silencioso (sin errores de tiempo de ejecución, solo degradación del rendimiento).
• Estrategia de Poda Alineada con RoPE: RAP introduce una granularidad de poda que se adhiere estrictamente a los pares de RoPE, asegurando que la operación de rotación sea válida sin requerir un reentrenamiento complejo o reconstrucción.
• Decodificación Sin Reconstrucción: A diferencia de los métodos de descomposición de bajo rango que deben reconstruir claves de dimensión completa en cada paso de decodificación, RAP fusiona la máscara de poda en $W_q$ de forma offline, eliminando el "impuesto de reconstrucción" y reduciendo tanto la sobrecarga de memoria como de cómputo.
• Evaluación Exhaustiva: El método se evalúa en múltiples familias de modelos (Llama, Mistral, Qwen) y escalas (3B a 14B), demostrando robustez.

4. Resultados Experimentales

Las evaluaciones se realizaron bajo un entorno de iso-KV (emparejando el ratio de canales KV retenidos $\rho$) para asegurar una comparación justa de calidad y eficiencia.

• Precisión vs. Compresión:
  • En un ratio de compresión del 30% ($\rho=0.7$), RAP mantiene una precisión utilizable (por ejemplo, WikiText-2 PPL ~9.12 para Llama-3-8B frente al baseline de 8.28), superando significativamente a la poda de canales ciega a RoPE (KVPruner), que sufre un colapso masivo de precisión (PPL ~397).
  • RAP logra un rendimiento comparable o superior a los fuertes baselines de bajo rango (PaLU, SVD-LLM) operando con menores costos de atención.
• Ganancias de Eficiencia:
  • Parámetros y FLOPs: RAP reduce los parámetros de atención y los FLOPs linealmente con el ratio de retención $\rho$ (por ejemplo, ~70% en $\rho=0.7$). En contraste, los métodos de bajo rango a menudo retienen conteos de parámetros más altos debido a la necesidad de factores de expansión en línea.
  • Latencia: RAP reduce la latencia de Prefill y Decode (por ejemplo, 83%/77% del baseline denso en $\rho=0.7$) porque evita la reconstrucción paso a paso. Los métodos de bajo rango a menudo muestran un aumento en la latencia, particularmente durante la decodificación.
• Robustez:
  • RAP mantiene el rendimiento a través de diferentes tamaños de modelo y arquitecturas.
  • Es ortogonal y complementario a la cuantización de pesos de 4 bits.
  • El método maneja eficientemente el indexado de RoPE no contiguo mediante un kernel de Triton personalizado, evitando la sobrecarga de los indexados estándar de PyTorch.

5. Significado y Reivindicaciones

El artículo afirma que RAP es el único método que simultáneamente:

1. Reduce el tamaño del caché KV.
2. Reduce los parámetros del modelo y los FLOPs de proyección.
3. Mantiene una alta precisión en altos ratios de compresión.
4. Opera sin sobrecarga de reconstrucción por paso.

Los autores enfatizan que RAP es una solución de "inserción directa" (drop-in) para los LLMs basados en RoPE. Al respetar las restricciones estructurales del Posicionamiento Rotatorio, resuelve el problema del "fallo silencioso" de los métodos de poda existentes y ofrece una alternativa más eficiente para la inferencia de contexto largo. El trabajo sugiere que la alineación estructural con mecanismos específicos del modelo (como RoPE) es crítica para una compresión de modelos efectiva.