When Hidden States Drift: Can KV Caches Rescue Long-Range Speculative Decoding?

1. Declaración del Problema

Decaimiento de Largo Alcance en la Decodificación Especulativa:
La decodificación especulativa acelera la inferencia de Modelos de Lenguaje Grandes (LLM) utilizando un modelo "borrador" ligero para proponer múltiples tokens, los cuales un modelo "objetivo" más grande verifica. Los borradores actuales del estado del arte (por ejemplo, EAGLE-3, MTP) dependen de la reutilización de estados ocultos, alimentando los estados internos ocultos del modelo objetivo al borrador.

• El Problema: Estos sistemas sufren de decaimiento de largo alcance: la precisión del borrado disminuye progresivamente a medida que aumenta el paso especulativo $k$.
• Explicación Existente: La comunidad atribuye esto a la discrepancia entre entrenamiento e inferencia (el borrador se condiciona sobre sus propios estados ocultos desviados durante la inferencia, a diferencia del entrenamiento). Aunque el Entrenamiento en Tiempo de Prueba (TTT) mitiga esto, el decaimiento persiste.
• La Brecha: El artículo argumenta que la discrepancia entre entrenamiento e inferencia no es la única causa. El problema fundamental radica en la preservación de la información del contexto. Los estados ocultos actúan como una "compresión de contexto sesgada" optimizada para el siguiente token inmediato, descartando información relevante para los futuros pasos especulativos.

2. Hipótesis Central y Análisis Teórico

Los autores proponen la Hipótesis de Reutilización de KV: Reutilizar la Caché KV (pares Clave-Valor) del modelo objetivo, en lugar de solo los estados ocultos, puede preservar un contexto más rico para la redacción de largo alcance.

• Reutilización de Estados Ocultos (Compresión Sesgada):
  • Los estados ocultos objetivo son el resultado de la agregación de atención: $h_t = \sum \alpha_i v_i$.
  • Los pesos $\alpha_i$ están optimizados para predecir el token actual $x_{t+1}$.
  • Los tokens débilmente relevantes para $x_{t+1}$ reciben pesos cercanos a cero y se descartan efectivamente en $h_t$.
  • Consecuencia: Para pasos especulativos posteriores ($x_{t+k}$), el borrador debe recuperar información que fue suprimida por el mecanismo de atención del objetivo, un problema de "recuperación de información" mal condicionado.
• Reutilización de KV (Contexto Explícito):
  • La caché KV retiene el conjunto completo de pares $(k_i, v_i)$ antes de la agregación.
  • El borrador puede realizar re-atención utilizando sus propias consultas futuras estimadas contra la caché KV completa del prefijo.
  • Consecuencia: No se descarta información antes del procesamiento del borrador. La dificultad se desplaza de "recuperar información perdida" a "estimar con precisión las consultas futuras" (un problema de aproximación de funciones).

3. Metodología: El Marco KVShot

Para probar esta hipótesis, los autores introducen KVShot, un marco de diagnóstico que compara tres paradigmas de reutilización bajo el mismo pipeline de Entrenamiento en Tiempo de Prueba Autoregresivo (TTT) (utilizado por EAGLE-3):

1. Reutilización Solo Oculta: Enfoque estándar de EAGLE-3/MTP (concatenando estados ocultos objetivo).
2. Reutilización Solo KV: El borrador recibe la caché KV objetivo mediante atención cruzada, reemplazando por completo los estados ocultos.
3. Reutilización Híbrida: Un mecanismo de fusión con puerta donde el estado oculto proporciona el ancla principal, y la caché KV proporciona una "corrección delta con puerta" para pasos de largo alcance.

Configuración Experimental:

• Modelo Objetivo: Qwen3-8B.
• Conjuntos de Datos: ShareGPT y UltraChat (escalados de 70k a 280k muestras).
• Métricas: Tasas de aceptación paso a paso ($\alpha_k$), Tokens Aceptados Medios (MAT) y proxies de aceleración de extremo a extremo (E2E).

4. Resultados Clave

A. Validación de Predicciones (Análisis Paso a Paso)

Los experimentos confirmaron tres predicciones teóricas:

1. Ventaja de Largo Alcance: La reutilización Solo KV se degrada con mayor gracia que la reutilización Solo Oculta a medida que aumenta $k$. En el paso $k=6$, un borrador de 4 capas Solo KV superó la línea base de EAGLE-3 ($\alpha_6 = 0.495$ frente a $0.469$).
2. Cuello de Botella en la Estimación de Consultas: El rendimiento Solo KV depende fuertemente de la profundidad del borrador. Un borrador de 1 capa tuvo un rendimiento pobre (MAT 1.84) porque no puede aproximar consultas objetivo complejas. Aumentar la profundidad a 2 capas produjo un salto masivo (MAT 2.23), pero 4 capas ofrecieron rendimientos decrecientes, quedando aún por detrás de la línea base de EAGLE-3 de 1 capa (MAT 2.37) en eficiencia general.
3. Desventaja de Corto Alcance: En $k=0$ o $1$, la reutilización de Estados Ocultos supera a la reutilización Solo KV porque los estados ocultos transportan un contenido semántico más rico que las incrustaciones crudas utilizadas para la re-atención KV.

B. Rendimiento Híbrido

El Borrador Híbrido (combinando ambas señales) logró los mejores resultados paso a paso:

• MAT: 2.54 (frente a 2.37 para la línea base de EAGLE-3).
• Retención: Mejoró la relación de retención de largo alcance del 73.5% al 77.3%.
• Mecanismo: La regla delta con puerta permitió con éxito que la ruta KV corrigiera la ruta de estados ocultos en horizontes más largos sin interrumpir la precisión de corto alcance.

C. La Verificación de Realidad de Extremo a Extremo

A pesar de las mejoras paso a paso, la aceleración de Extremo a Extremo (E2E) fue marginal:

• MAT Medido por HF: Aumentó solo de 5.01 (EAGLE-3) a 5.04 (Híbrido).
• Latencia: La ruta adicional de atención cruzada introdujo un 5–10% de latencia de redacción extra.
• Conclusión: El pipeline actual de TTT autoregresivo no puede explotar eficientemente la señal KV para superar el costo computacional añadido.

5. Cuellos de Botella Estructurales Identificados

El artículo identifica tres razones específicas por las que el pipeline actual no logra traducir las ganancias paso a paso en aceleraciones E2E:

1. Dificultad en la Estimación de Consultas: Estimar consultas objetivo futuras requiere replicar las transformaciones no lineales de múltiples capas del modelo objetivo. Los borradores superficiales (1-2 capas) carecen de la capacidad para aproximar estas consultas con precisión.
2. Señales de Gradiente Escasas: En el TTT autoregresivo, las proyecciones KV del lado del borrador ($W_K, W_V$) solo se actualizan mediante los pocos tokens generados por el borrador (por ejemplo, 7 tokens) por paso, mientras que la KV del prefijo se copia. Esto resulta en gradientes extremadamente escasos y desequilibrados para la vía KV.
3. Hambre de Gradientes Inducida por la Puerta: En el modelo híbrido, la puerta que controla la corrección KV tiende a colapsar a casi cero al inicio del entrenamiento (porque la rama KV inicializada aleatoriamente comienza como ruido). Una vez cerrada, la puerta bloquea los gradientes hacia la rama KV, impidiendo que aprenda a ser útil. Esto crea un "estado de baja inyección auto-reforzante".

6. Significado y Direcciones Futuras

• Cambio Conceptual: El artículo replantea el decaimiento de largo alcance no solo como una discrepancia entre entrenamiento e inferencia, sino como un problema de preservación de información. Demuestra que los estados ocultos son una compresión con pérdida para tareas de horizonte largo.
• Herramienta de Diagnóstico: KVShot proporciona un campo de pruebas fundamental para aislar los efectos de representación de los artefactos de entrenamiento.
• Hoja de Ruta para Arquitecturas de Próxima Generación: Los hallazgos sugieren que el TTT autoregresivo es insuficiente para la decodificación consciente de KV. Para realizar el potencial de la reutilización de KV, el campo debe avanzar hacia Paradigmas de Entrenamiento por Bloques (por ejemplo, DFlash).
  • El entrenamiento por bloques genera muchos tokens de borrador en paralelo, proporcionando señales de gradiente densas para las proyecciones KV.
  • Permite modelos de borrador más profundos capaces de estimar consultas con precisión sin penalizaciones de latencia secuencial.

Resumen: Si bien las cachés KV ofrecen teóricamente una señal superior para la decodificación especulativa de largo alcance al preservar el contexto, los pipelines actuales de entrenamiento autoregresivo fallan en aprovechar esto debido a los límites en la estimación de consultas y al hambre de gradientes. El camino a seguir reside en arquitecturas de entrenamiento paralelo y por bloques.