DepthKV: Layer-Dependent KV Cache Pruning for Long-Context LLM Inference

Aquí se presenta un resumen técnico detallado del artículo "DepthKV: Poda de Caché KV Dependiente de la Capa para Inferencia de LLM de Contexto Largo".

1. Planteamiento del Problema

Los Modelos de Lenguaje Grande (LLM) enfrentan un cuello de botella crítico de memoria durante la inferencia de contexto largo. Si bien la generación autoregresiva depende de la caché de Clave-Valor (KV) para evitar recalcular tokens pasados, la huella de memoria de esta caché crece linealmente con la longitud de la secuencia ($O(L \cdot N \cdot d)$, donde $L$ es el número de capas, $N$ es la longitud de la secuencia y $d$ es la dimensión oculta).

Las soluciones existentes post-entrenamiento utilizan principalmente la poda de caché KV, que descarta tokens con puntuaciones de atención bajas para ajustarse a un presupuesto de memoria fijo. Sin embargo, una limitación fundamental de los métodos actuales (por ejemplo, H2O, StreamingLLM) es la asunción de poda uniforme: aplican la misma proporción de poda a través de todas las capas del transformador. Los autores argumentan que esto es subóptimo porque asume implícitamente que todas las capas contribuyen por igual al rendimiento del modelo, ignorando la sensibilidad heterogénea de diferentes capas ante la eliminación de tokens.

2. Metodología: DepthKV

El marco propuesto, DepthKV, desafía la asunción de poda uniforme introduciendo una asignación de caché KV dependiente de la capa. En lugar de descartar un porcentaje fijo de tokens por capa, DepthKV redistribuye un presupuesto de caché KV global fijo entre las capas basándose en su sensibilidad específica a la poda.

A. Pre-estudio: Análisis de Sensibilidad de Capa

Antes de diseñar la estrategia de asignación, los autores realizaron un estudio de ablación por capas:

• Método: Poda de tokens en una sola capa específica a la vez, manteniendo las demás intactas, midiendo la caída resultante en el rendimiento (ROUGE-1).
• Hallazgo: La degradación del rendimiento es altamente no uniforme. Ciertas capas "críticas" causan caídas masivas cuando se podan, mientras que otras muestran un impacto insignificante o incluso redundancia.
• Amplificación de Contenido: Identificaron "capas de amplificación de contenido". Podar estas capas no solo reduce la precisión, sino que suprime significativamente la longitud y la informatividad de la salida (medido por YapScore), indicando que estas capas son cruciales para sostener la generación.
• Validación Estadística: Pruebas de permutación confirmaron que la variación en la importancia de las capas es estadísticamente significativa a través de modelos (Gemma, LLaMA, Qwen) y conjuntos de datos.

B. Métricas de Representación para la Importancia

Para automatizar la asignación sin una ablación exhaustiva, los autores analizaron propiedades de representación para encontrar un proxy de la importancia de la capa.

• Selección de Métrica: Evaluaron seis métricas (Entropía Espectral, Rango Efectivo, Curvatura, DiME, LiDAR e InfoNCE).
• Hallazgo Clave: InfoNCE (una pérdida contrastiva que mide la robustez de la representación ante perturbaciones de entrada) mostró la correlación negativa más fuerte con la degradación del rendimiento.
  • Interpretación: Las capas con valores de InfoNCE más bajos (representaciones más robustas) son más críticas. Podarlas causa caídas mayores en el rendimiento.
  • Uso: DepthKV utiliza puntuaciones de InfoNCE para clasificar las capas y determinar la asignación del presupuesto.

C. Estrategias de Asignación

DepthKV opera bajo una proporción de poda global fija ($\rho$) pero la distribuye de manera no uniforme. Se proponen tres estrategias:

1. Protección de Capas Intermedias (MLP): Un enfoque heurístico que preserva un subconjunto fijo de capas intermedias (por ejemplo, capas $\lfloor L/2 \rfloor$ y $\lfloor L/2 \rfloor + 1$) mientras poda uniformemente las demás.
2. Asignación Guiada por Métricas (MGA): El método principal propuesto. Calcula una puntuación de importancia $s^{(l)}$ basada en InfoNCE. Las capas con mayor robustez (InfoNCE más bajo) reciben un presupuesto KV mayor (proporción de poda menor), mientras que las capas menos críticas se podan de manera más agresiva.
3. Asignación de Métricas de Capas Intermedias (MLMA): Un enfoque híbrido que protege estructuralmente las capas intermedias mientras distribuye el presupuesto restante entre otras capas basándose en las puntuaciones de InfoNCE.

3. Contribuciones Clave

1. Insight Teórico: Demostró que las capas del transformador exhiben una sensibilidad estadísticamente significativa y no uniforme a la poda de caché KV, invalidando la asunción de poda uniforme.
2. Descubrimiento de Amplificación de Contenido: Identificó capas específicas cuya poda suprime la generación de contenido, vinculando la robustez de la representación directamente con la calidad y longitud de la salida.
3. Propuesta de Marco: Introdujo DepthKV, un marco de asignación dependiente de la capa que redistribuye un presupuesto de memoria fijo basado en la importancia de la capa (utilizando específicamente InfoNCE).
4. Validación Empírica: Mostró que DepthKV supera consistentemente a la poda uniforme en diversas tareas y modelos sin requerir entrenamiento adicional.

4. Resultados Experimentales

Los autores evaluaron DepthKV en tres familias de modelos (Gemma-7B, LLaMA-3.1-8B, Qwen2.5-7B) a través de siete conjuntos de datos:

• Tareas: Resumen de documentos largos (arXiv, PubMed, GovReport, LegalCase), QA fundamentado en documentos (HotpotQA, Qasper) y Razonamiento Matemático (GSM-$\infty$).
• Configuración: Reducción global fija de KV del 60% (retención del 40%).
• Rendimiento:
  • Resumen: DepthKV (específicamente MGA) mejoró significativamente ROUGE-1 y la similitud SBERT en comparación con la poda uniforme (por ejemplo, +3.0 ROUGE-1 en arXiv para Gemma).
  • QA y Razonamiento: Las variantes de DepthKV lograron puntuaciones Exact Match (EM) y F1 más altas. Por ejemplo, en Qasper, MLMA-6L superó a todos los baselines.
  • Robustez: Los beneficios se extendieron a tareas de razonamiento (GSM-$\infty$), donde todas las variantes de DepthKV superaron a la poda uniforme.
  • LLM como Juez: Las evaluaciones utilizando Prometheus confirmaron que las salidas de DepthKV eran más correctas, completas y concisas que los baselines de poda uniforme.
• Ablación: El estudio confirmó que cómo se distribuye el presupuesto (asignación por capas) importa más que el estimador específico de importancia de tokens (por ejemplo, H2O con V vs. sin V).

5. Significancia y Limitaciones

Significancia:

• Eficiencia: DepthKV ofrece una solución simple, post-entrenamiento, al cuello de botella de memoria de contexto largo, permitiendo ventanas de contexto más largas en hardware existente sin reentrenamiento.
• Generalizabilidad: El método es agnóstico a la arquitectura y efectivo a través de diferentes familias de modelos y tipos de tareas.
• Cambio de Paradigma: Mueve el campo de la poda a "nivel de token" a la asignación de recursos a "nivel de capa", reconociendo la heterogeneidad funcional de las capas del transformador.

Limitaciones:

• No Consciente de la Consulta: DepthKV estima la importancia de los tokens durante la etapa de prellenado basándose en la atención histórica. No se ajusta dinámicamente a la consulta de decodificación específica, perdiendo potencialmente tokens que se vuelven críticos solo para consultas futuras específicas.
• Granularidad a Nivel de Cabeza: El marco actual asigna presupuestos a nivel de capa pero agrega la atención a través de cabezas. No tiene en cuenta la variabilidad intra-capas donde cabezas de atención específicas podrían ser más críticas que otras.

Conclusión

DepthKV establece que no todas las capas del transformador son iguales en la inferencia de contexto largo. Al redistribuir presupuestos de memoria para preservar las capas más sensibles y robustas (identificadas mediante InfoNCE), el método logra compensaciones rendimiento-por-memoria superiores en comparación con la poda uniforme, ofreciendo un camino práctico hacia adelante para escalar las ventanas de contexto de los LLM.