Resumen Técnico: Transferencia de Caché KV entre Modelos en Familias de LLM

1. Formulación del Problema

El servicio de modelos de lenguaje extensos (LLM) en producción depende cada vez más de la orquestación de múltiples modelos para el escalonamiento costo-calidad, el cambio en medio de una conversación y el enrutamiento. Estos flujos de trabajo frecuentemente alternan entre miembros de diferentes tamaños de una misma familia de modelos (por ejemplo, de Qwen3 14B a 32B). Aunque estos modelos comparten elecciones arquitectónicas centrales y datos de preentrenamiento, difieren en escala, número de capas y recetas de entrenamiento.

El cuello de botella principal en tales flujos de trabajo es el costo de prefill (prellenado). Cada vez que el contexto de una conversación se transfiere de un modelo de origen a un modelo de destino, este último debe procesar nuevamente todo el contexto acumulado para poblar su propio caché de Clave-Valor (KV). Este paso hacia adelante escala con el tamaño del modelo y la longitud del prompt, anulando las ganancias de eficiencia del intercambio de modelos. Las soluciones existentes como el almacenamiento en caché de prefijos solo funcionan dentro de un mismo modelo, y los enfoques previos de transferencia entre modelos (por ejemplo, fusores neuronales, adaptadores latentes) suelen requerir entrenamiento basado en gradientes o restricciones arquitectónicas estrictas (por ejemplo, conteos de capas idénticos).

Este artículo propone la Transferencia de Caché KV entre Modelos, un método para reutilizar el caché KV prellenado del modelo de origen directamente en el modelo de destino, omitiendo el prefill del modelo de destino por completo. El objetivo es encontrar un mapeo $f: C_S \to \hat{C}_T$ tal que el modelo de destino, decodificando desde el caché mapeado $\hat{C}_T$, produzca salidas equivalentes a la decodificación desde su propio caché real $C_T$.

2. Metodología

Los autores plantean la hipótesis de que las relaciones de KV entre modelos poseen una estructura lineal sustancial, lo que permite un mapeo de forma cerrada y sin gradientes. El marco propuesto opera sobre pares de "KV emparejados", donde el modelo de origen y el de destino comparten el mismo número de cabezales KV y la misma dimensión por cabezal, incluso si sus conteos totales de capas o tamaños de parámetros difieren.

El mapeador está diseñado con tres componentes principales:

A. Regresión de Ridge por Cabezal

En lugar de una red neuronal global, el método ajusta un mapeo lineal independiente para cada capa y cabezal del modelo de destino.

• Entrada: Características KV concatenadas de un subconjset seleccionado de capas de origen.
• Salida: Los vectores de Clave ($K$) y Valor ($V$) esperados del destino.
• Optimización: Los pesos se resuelven mediante regresión de ridge (solución de forma cerrada) utilizando un pequeño conjunto de calibración (500 secuencias de 1,024 tokens de FineWeb-Edu). Se utiliza la regularización de ridge ($\lambda=0.01$) para manejar la naturaleza casi singular de la matriz de características cuando las capas de origen están correlacionadas.

B. Selección de Capas de Origen Trans-capa

Dado que los modelos de origen y destino tienen profundidades distintas ($L_S \neq L_T$), el método debe determinar qué capas de origen informan a cada capa de destino.

• Estrategia de Selección: Para cada capa de destino, se seleccionan las $k$ mejores capas de origen basadas en su poder predictivo (medido por el $R^2$ promedio por cabezal en un conjunto de calibración).
• Mecanismo: Las características KV de estas $k$ mejores capas de origen se concatenan para formar la entrada para la regresión de ridge. El análisis empírico muestra que la información complementaria se distribuye a través de múltiples capas de origen; una sola capa es insuficiente, pero el rendimiento se satura rápidamente (por ejemplo, $k=6$ captura la mayor parte de la varianza de $k=todas$).

C. Mapeo de Espacio de Contenido (Factorización de RoPE)

Los LLM estándar utilizan Incrustaciones de Posición Rotatorias (RoPE), que aplican rotaciones dependientes de la posición a las claves.

• Problema: Ajustar un mapeador directamente sobre claves rotadas vincula los pesos a la distribución de posición específica vista durante la calibración, lo que limita la transferibilidad a diferentes longitudes de contexto.
• Solución: El método elimina el RoPE de las claves de origen antes del mapeo y reaplica el RoBE del destino después del mapeo.
  • $K_{stripped} = K_{source} \cdot R^{-1}_{\Theta_s}$
  • $\hat{K}_{target} = (K_{stripped} \cdot W + b) \cdot R_{\Theta_t}$
• Esto desacopla la rotación posicional del contenido semántico, haciendo que el ajuste lineal sea independiente de la posición y reutilizable en diversas longitudes de contexto (hasta 32k tokens).

D. Extensión No Lineal (MLP)

Para los pares donde el mapeador de ridge lineal no logra retener suficiente precisión, los autores proponen un reemplazo directo utilizando un pequeño Perceptrón Multicapa (MLP). Este MLP redistribuye el error residual lejos de los subespacios sensibles a la atención, recuperando el rendimiento en pares difíciles.

3. Contribuciones Clave

1. Marco de Mapeo de Forma Cerrada: Un marco de regresión de ridge por cabezal, sin gradientes, que permite la transferencia de caché KV entre modelos sin necesidad de entrenamiento. Se basa en la selección de las $k$ mejores capas de origen y el mapeo de contenido con RoPE eliminado.
2. Validación de Multifamilia: Validación a través de tres familias de modelos (Qwen3, Llama 3.1, Ministral 3) y seis pares de KV emparejados.
3. Recuperación No Lineal: Demostración de que un MLP ligero puede recuperar el rendimiento en pares "difíciles" donde el mapeo lineal falla, específicamente al redistribuir el error lejos de las dimensiones críticas de la atención.
4. Métrica Diagnóstica: Identificación de que la similitud de coseno de la salida de atención es un mejor predictor de la retención de la transferencia descendente que las métricas de reconstrucción estándar ($R^2$). Los autores muestran una correlación de Pearson de $r=+0.57$ entre la similitud de coseno de la salida de atención y la retención de HellaSwag, comparado con un $r=-0.20$ para $R^2$.

4. Resultados Experimentales

El marco fue evaluado en seis pares de KV emparejados a través de tres familias utilizando cinco benchmarks de precisión (ARC-C, HellaSwag, WinoGrande, MMLU, GSM8K) y la perplejidad de WikiText-2.

• Retención de Precisión:
  • Nivel 1 (Éxito): Cuatro pares (Qwen3 14B→32B, 8B→32B; Llama 3.1 8B→70B; Ministral 3B→8B) retuvieron entre el 73% y el 98% de la precisión autónoma del modelo de destino.
  • Nivel 2 (Fallo): Dos pares (Ministral 3B→14B, 8B→14B) degradaron significativamente (retención del 42–44%). Sin embargo, reemplazar el mapeador lineal con un MLP recuperó la retención de HellaSwag en un +24.3 a +36.8 puntos porcentuales, llevando a estos pares por encima del 90%.
• Latencia:
  • El mapeador se ejecuta de 2.7 a 25 veces más rápido que el re-prellenado del modelo de destino.
  • Las aceleraciones son más pronunciadas para transferencias de pequeño a grande y longitudes de contexto largas (por ejemplo, una aceleración de 25× a 32k tokens para Qwen3 14B→3B).
• Estabilidad de Multi-turno:
  • En escenarios de transferencia de múltiples turnos (CoQA), la deriva permanece pequeña. La deriva de pequeño a grande aumenta solo 1.7 pp en 10 turnos; la deriva de grande a pequeño crece linealmente a 0.33 pp/turno, lo que sugiere estabilidad para sesiones típicas.
• Eficiencia de Calibración:
  • El mapeador requiere solo 500 secuencias para la calibración.
  • El ajuste toma aproximadamente 47–87 minutos por par en un solo nodo de 8×H100, sin entrenamiento basado en gradientes.

5. Significado y Reivindicaciones

El artículo afirma que la transferencia de caché KV entre modelos dentro de una familia de LLM admite un ajuste de forma cerrada y sin entrenamiento. Los autores argumentan que la relación de KV entre modelos es mayormente lineal, desafiando la suposición de que los adaptadores neuronales complejos son necesarios para tales transferencias.

La importancia de este trabajo radica en:

1. Practicidad: Permite el escalonamiento costo-calidad y el cambio en medio de una conversación sin el costo prohibitivo del re-prellenado, haciendo que la orquestación dinámica de modelos sea factible en producción.
2. Cambio en la Evaluación: Los hallazgos sugieren que los criterios de evaluación para los mapeadores entre modelos deberían cambiar de las métricas de reconstrucción bruta ($R^2$) a diagnósticos conscientes del subespacio (como la similitud de coseno de la salida de atención), ya que estos últimos predicen mejor el rendimiento en tareas descendentes.
3. Redistribución del Error: El trabajo destaca que el fallo de la transferencia se debe a menudo a dónde cae el error (en subespacios sensibles a la atención) más que a la magnitud del error, lo que motiva futuros objetivos alineados con la atención.

Los autores mantienen la modestia respecto a las limitaciones, señalando que el enfoque está actualmente restringido a modelos de atención completa densa dentro de la misma familia con configuraciones de KV emparejadas, y que la calibración se realiza en un único dominio (FineWeb-Edu). Dejan la transferencia entre familias y las condiciones de KV no emparejados como preguntas abiertas para trabajos futuros.