Resumen Técnico: SEKV – Caché KV de Resolución Adaptativa con Memoria Semántica Jerárquica

Declaración del Problema

Los Modelos de Lenguaje de Gran Escala (LLMs) se despliegan cada vez más en escenarios que requieren razonamiento de contexto largo (por ejemplo, análisis de documentos, diálogo multiturno, planificación de agentes). Sin embargo, el caché de Clave-Valor (KV), que almacena estados intermedios para acelerar la decodificación, crece linealmente con la longitud de la secuencia. Retener un caché KV completo para contextos que superan los 128K tokens resulta prohibitivamente costoso en términos de memoria de GPU, lo que a menudo hace que la inferencia de contexto largo sea inviable sin compresión.

Las estrategias de compresión existentes enfrentan un compromiso fundamental entre eficiencia y fidelidad:

1. Evicción de Tokens: Los métodos que descartan permanentemente tokens basados en estadísticas de atención (por ejemplo, H2O, SnapKV) corren el riesgo de fragmentar evidencia coherente o perder información de forma irreversible que puede volverse crítica solo más tarde en el proceso de generación.
2. Agrupación Semántica Estática: Los métodos que comprimen tokens en bloques o frases semánticas en el momento del prefill (por ejemplo, ChunkKV, SentenceKV) fijan la resolución de la compresión de forma temprana. Carecen de la capacidad de recuperar detalles finos a nivel de token de un intervalo comprimido una vez que dicho intervalo se vuelve relevante durante la decodificación.

El desafío central es diseñar una organización de memoria que preserve el acceso a la evidencia relevante bajo estrictas restricciones de memoria sin pérdida irreversible de información, permitiendo al modelo ajustar dinámicamente la resolución del contexto almacenado según la relevancia de la consulta.

Metodología: SEKV

Los autores proponen SEKV (Semantic KV Cache), un caché KV semántico de resolución adaptativa que organiza el contexto en una memoria semántica jerárquica que abarca GPU y CPU. SEKV evita la evicción irreversible de tokens almacenando el contexto a resoluciones variables y realizando un "zoom" dinámico hacia el nivel de detalle apropiado durante la decodificación.

1. Segmentación de Intervalos Guiada por Entropía

En lugar de una fragmentación de tamaño fijo, SEKV segmenta el contexto de entrada en intervalos semánticamente coherentes utilizando la sorpresa del token (token surprisal) como una señal de frontera de costo cero.

• Mecanismo: Durante la pasada de prefill, se calcula la sorpresa $H_t = -\log p(x_t | x_{<t})$ para cada token. Los tokens de alta sorpresa, que a menudo marcan transiciones semánticas (cambios de tema, introducciones de entidades), se identifican como fronteras de intervalo.
• Tokens Ancla: Estos tokens de frontera de alta sorpresa se retienen a resolución completa en la GPU como "tokens ancla". Estos portan información semántica concentrada esencial para el enrutamiento y la reconstrucción posteriores.
• Intervalos (Spans): Los tokens entre dos anclas vecinas forman un intervalo semántico $S$.

2. Representación de Intervalo de Doble Resolución y Jerarquía de Memoria

Cada intervalo semántico se representa simultáneamente a dos resoluciones a través de una jerarquía GPU–CPU:

• GPU (Enrutamiento Grueso): Cada intervalo se representa mediante un vector de resumen ligero (un promedio ponderado por la sorpresa y normalizado de los vectores de clave proyectados). Este resumen reside en la GPU y se utiliza para un enrutamiento eficiente y de grano grueso para determinar qué intervalos son relevantes para la consulta actual.
• CPU (Reconstrucción de Grano Fino): La estructura completa a nivel de token del intervalo se codifica de forma compacta en la CPU utilizando una base SVD de bajo rango. Un predictor ligero asigna compuertas suaves (soft gates) por componente a los valores singulares, lo que permite al sistema retener de forma adaptativa un presupuesto de rango $r_S$ para cada intervalo basado en su densidad de información. Esto permite una reconstrucción fiel y bajo demanda de los pares KV a nivel de token sin almacenar la matriz completa en la GPU.

3. Mecanismo de Zoom-In Entrenado

Durante la decodificación, SEKV emplea un mecanismo aprendido para seleccionar dinámicamente qué intervalos expandir:

• Enrutamiento: La consulta actual se proyecta en el mismo espacio de enrutamiento que los resúmenes de los intervalos. Una compuerta de relevancia sigmoidea calcula la probabilidad de que un intervalo sea relevante.
• Expansión: Un umbral aprendido, por cabeza y por capa, determina si un intervalo debe expandirse. Si se selecciona, el sistema recupera de forma asíncrona la base SVD de bajo rango correspondiente desde la CPU, reconstruye los pares de clave-valor a nivel de token e integra los mismos en el cómputo de atención.
• Atención de Resolución Mixta: El mecanismo de atención opera sobre una memoria de resolución mixta: los tokens ancla contribuyen con pares KV completos, los intervalos no expandidos contribuyen con sus vectores de resumen gruesos (con una corrección de tamaño) y los intervalos expandidos contribuyen con pares de nivel de token reconstruidos.
• Entrenamiento: El LLM base permanece congelado. Solo las proyecciones de enrutamiento ligeras, los umbrales de zoom-in y el predictor de compuerta de rango se entrenan utilizando una combinación de pérdida de destilación (alineando con un profesor de KV completo), supervisión de zoom, pérdida de reconstrucción y una penalización de presupuesto.

Contribuciones Clave

1. Arquitectura SEKV: Un caché KV semántico de resolución adaptativa que almacena el contexto a múltiples resoluciones y ajusta dinámicamente el detalle del caché según la relevancia de la consulta, evitando la evicción irreversible de tokens.
2. Segmentación Guiada por Entropía: Un método que utiliza la sorpresa del token para identificar intervalos semánticos coherentes y preservar tokens ancla de alta sorpresa a resolución completa para una recuperación fiable.
3. Estructura de Doble Representación: Una novedosa estructura KV que combina resúmenes de intervalos residentes en GPU para enrutamiento grueso, bases SVD de bajo rango residentes en CPU para la aproximación de la reconstrucción a nivel de token, y un mecanismo de zoom-in entrenado que expande selectivamente solo los intervalos relevantes para la consulta.

Resultados Experimentales

SEKV fue evaluado en cuatro benchmarks (LongBench, RULER, InfiniteBench, Needle-in-a-Haystack) y cinco modelos de base (incluyendo variantes de Llama-3, Mistral y Qwen).

• Rendimiento: SEKV supera al referente más fuerte de compresión semántica (SentenceKV) en un promedio del 5.9% en todos los benchmarks y modelos. Logra los mejores resultados en 20 de 20 comparaciones de benchmark-modelo.
• Fidelidad de Recuperación: En el benchmark Needle-in-a-Haystack (NIAH), SEKV muestra una estabilidad de recuperación superior a través de diversas longitudes de contexto y profundidades de la aguja en comparación con los métodos de evicción de tokens y métodos semánticos estáticos.
• Eficiencia de Memoria: A un contexto de 128K, SEKV reduce el consumo de memoria de la GPU en un 53.3% en comparación con el caché KV completo, manteniendo un rendimiento cercano a la línea base sin compresión.
• Eficiencia: SEKV logra una latencia y un rendimiento competitivos, superando a menudo al caché KV completo en latencia debido a la reducción de la presión de ancho de banda de memoria, mientras supera significativamente en precisión a los métodos de evicción agresiva.
• Eficiencia de Parámetros: El método añade menos del 0.05% de parámetros entrenables (aproximadamente 4.3 millones de parámetros para Llama-3-8B) manteniendo el modelo base congelado.

Significancia y Reivindicaciones

El artículo argumenta que la inferencia eficiente de contexto largo requiere un cambio de la compresión estática a la recuperación de resolución dinámica. Los autores afirman que, mientras la compresión estática preserva el contexto grueso, falla en recuperar evidencia de grano fino cuando esta se vuelve relevante tardíamente en el proceso de generación. SEKV aborda esto organizando el contexto en intervalos semánticos guiados por la entropía y utilizando un sistema de memoria jerárquica que permite al modelo hacer "zoom" en la información relevante bajo demanda.

La significancia de este trabajo radica en demostrar que:

1. La evicción irreversible no es necesaria: Se puede lograr una inferencia de contexto largo de alta fidelidad mediante la transferencia de representaciones de bajo rango recuperables a la CPU en lugar de descartar los tokens.
2. La resolución adaptativa es crítica: Asignar resolución de grano fino solo donde es necesario (mediante el mecanismo de zoom-in) permite a los modelos equilibrar las restricciones de memoria con la fidelidad de recuperación de manera más efectiva que las políticas estáticas.
3. La estructura semántica ayuda a la compresión: El uso de la segmentación guiada por la entropía crea unidades más compresibles que los bloques de tokens arbitrarios, permitiendo una aproximación de bajo rango más eficiente.

Los autores reconocen limitaciones, señalando que la efectividad depende de la calidad de la construcción de los intervalos semánticos (que puede tener dificultades con discursos altamente fragmentados) y que el enfoque introduce una dependencia del ancho de banda CPU-GPU, lo que podría impactar la latencia en escenarios de consulta adversos. Sin embargo, concluyen que SEKV proporciona un marco robusto para la inferencia de LLM de contexto largo escalable y de alta fidelidad.