Resumen Técnico: S4R - Muestreo Selectivo, Subespacios y Reconstrucción Dispersa para el Caché KV de Contexto Largo Comprimido

1. Planteamiento del Problema

La expansión de las ventanas de contexto en los Modelos de Lenguaje de Gran Escala (LLMs) a cientos de miles o millones de tokens ha creado un cuello de botella crítico en la memoria. El caché de Clave-Valor (KV), requerido para almacenar computaciones intermedias durante la generación autorregresiva, crece linealmente con la longitud del contexto, superando a menudo la huella de memoria de los propios parámetros del modelo.

Las soluciones existentes para la compresión del caché KV enfrentan un compromiso fundamental entre eficiencia y adaptabilidad:

• Métodos Offline/Fijos: Enfoques como Palu, LoRC y Eigen Attention realizan una descomposición de bajo rango (p. ej., SVD) sobre datos de calibración o pesos antes de la inferencia. Aunque son eficientes durante la decodificación, son sensibles a los cambios de distribución y dependen de diseños de compresión fijos que pueden no alinearse con prompts específicos.
• Métodos Online/Dependientes del Prompt: Enfoques como ShadowKV y xKV estiman representaciones de bajo rango a partir del prompt actual durante la inferencia. Aunque son más robustos ante cambios de contexto, incurren en costos computacionales prohibitivos, particularmente durante la etapa de prellenado (prefilling), donde descomponer todo el prompt largo es costoso, y a menudo requieren la reconstrucción de todo el caché durante la decodificación, afectando el rendimiento (throughput).

El artículo identifica una brecha: un método que combine la baja sobrecarga de prellenado de la compresión fija con la adaptabilidad al contexto de los métodos dependientes del prompt, sin incurrir en el costo de la descomposición de todo el prompt o la reconstrucción de todo el caché.

2. Metodología: S4R

S4R (Selective Sampling, Subspaces, and Sparse Reconstruction) aborda este dilema mediante un proceso de dos etapas: Prellenado (Prefilling) y Decodificación (Decoding).

2.1 Etapa de Prellenado: Construcción de Subespacios Selectivos

En lugar de aplicar la Descomposición de Valores Singulares (SVD) a toda la secuencia de contexto largo, S4R construye subespacios de bajo rango a partir de un subconjunto de tokens muestreados selectivamente.

• Preservación de Tokens de Sumidero (Sink Tokens): Los primeros $s$ tokens (identificados como "tokens de sumidero" que actúan como pivotes globales en la atención) se excluyen de la compresión y se almacenan en precisión completa para mantener dependencias de largo alcance estables.
• Muestreo Híbrido: Para los tokens restantes no pertenecientes al sumidero, S4R muestrea un subconjunto representativo $X[I]$ utilizando una estrategia híbrida:
  1. Muestreo uniforme de posiciones tempranas.
  2. Un bloque contiguo de los tokens más recientes.
• Normalización y SVD: Los keys y values muestreados se normalizan a lo largo de la dimensión de características para evitar que los vectores de gran magnitud dominen la descomposición. Luego, se realiza una SVD truncada sobre este subconjunto para derivar las bases de bajo rango ($B_K, B_V$).
• Almacenamiento de Coeficientes: Los estados KV no pertenecientes al sumidero se proyectan sobre estas bases aprendidas, almacenando solo los coeficientes de bajo rango ($C_K, C_V$) en lugar de los vectores de dimensión completa.
• Proyección de Query: Las queries locales también se proyectan en el mismo subespacio de bajo rango para permitir una estimación de relevancia eficiente posteriormente.

2.2 Etapa de Decodificación: Reconstrucción Dispersa

Durante la generación, S4R evita reconstruir todo el caché. En su lugar, reconstruye solo un conjunto disperso de tokens considerados relevantes.

• Estimación de Relevancia en el Espacio Latente: La relevancia de los tokens se estima directamente en el espacio de coeficientes latentes. Los coeficientes de la query proyectada se multiplican con los coeficientes de los keys almacenados en caché para calcular los logits de atención sin la reconstrucción completa.
• Selección Dispersa: Se selecciona un subconjunto de tokens para la reconstrucción basado en:
  1. Ventana Local: Los últimos $W$ tokens siempre se retienen.
  2. Importancia Global: Los $k$ mejores tokens del resto de la historia se seleccionan basándose en los puntajes de relevancia promedio derivados del espacio latente.
• Reconstrucción Dispersa: Solo los tokens locales y globales seleccionados se reconstruyen a partir de sus coeficientes utilizando las bases aprendidas. Estos se concatenan con los tokens de sumidero de precisión completa para realizar la computación de atención final.

3. Contribuciones Clave

1. Compresión de Bajo Rango Consciente del Prompt: S4R construye las bases de key/value a partir de tokens de prompt muestreados selectivamente en lugar de datos de calibración externos o el prompt completo, intercambiando una pequeña cantidad de adaptación de prellenado por un mejor alineamiento con el contexto.
2. Estrategia de Reconstrucción Dispersa: Al retener solo las posiciones informativas (ventana local + tokens globalmente importantes) durante la decodificación, S4R evita el cuello de botella computacional de la reconstrucción de todo el caché.
3. Preservación de Tokens de Sumidero: El método preserva explícitamente los primeros tokens "sumidero" en precisión completa, abordando la inestabilidad que suele causar la compresión de pivotes críticos de atención global.
4. Compromiso Eficiencia-Precisión: El diseño combina con éxito la baja sobrecarga de prellenado de los métodos fijos con la adaptabilidad de los métodos dependientes del prompt.

4. Resultados Experimentales

Los autores evaluaron S4R en los benchmarks LongBench y RULER utilizando familias de modelos Llama (1B, 8B) y Qwen (4B, 14B).

• Precisión:
  • LongBench: S4R logra hasta 5× de compresión del caché KV manteniendo una precisión cercana a la línea base de caché completo. En Qwen3-4B, alcanza un puntaje promedio de 47.79 (frente a 47.79 para el caché completo), superando a las líneas base fijas como Palu y a métodos competitivos dependientes del prompt como ShadowKV.
  • RULER: Bajo una relación de compresión de 2.5× con un contexto de 128k, S4R alcanza un puntaje promedio de 93.61, igualando casi al caché KV completo (93.97) y superando ligeramente a xKV (93.47). Muestra un desempeño particularmente fuerte en tareas de recuperación de múltiples llaves y múltiples consultas.
• Eficiencia:
  • S4R reduce sustancialmente la sobrecarga de los métodos dependientes del prompt. Comparado con xKV, S4R reduce el Tiempo al Primer Token (TTFT) de 79.90s a 27.65s y la latencia de generación de 512 tokens de 260.76s a 61.54s.
  • Aunque es ligeramente más lento que los métodos fijos como Palu en implementaciones no optimizadas, S4R ofrece un compromiso de precisión-eficiencia superior bajo relaciones de compresión agresivas.

5. Significado y Reivindicaciones

El artículo afirma que S4R proporciona un camino práctico hacia la compresión adaptativa del caché KV para la inferencia de LLMs de contexto largo. Su principal significancia radica en resolver el "dilema" de los métodos de bajo rango existentes:

• Evita la dependencia de datos de calibración de los métodos offline.
• Evita los costos prohibitivos de prellenado y reconstrucción de los métodos online.

Al demostrar que el subespacio dominante de KV puede aproximarse a partir de un subconjunto cuidadosamente seleccionado de tokens y que la reconstrucción completa no es necesaria en cada paso de decodificación, S4R permite una inferencia de contexto largo de alta fidelidad con una memoria persistente significativamente reducida y una sobrecarga computacional manejable. Los autores señalan que su implementación no está optimizada para sistemas, lo que sugiere que las métricas de eficiencia reportadas reflejan la sobrecarga algorítmica en lugar de un sistema totalmente desplegado, y que es posible obtener más ganancias con optimizaciones a nivel de kernel.