BudgetDraft: Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding

Resumen Técnico: BudgetDraft – Entrenamiento Multi-Vista Sensible a la Aceptación para la Decodificación Especulativa de KV Dispersos

1. Planteamiento del Problema

El artículo aborda un cuello de botella crítico en el despliegue de la Decodificación Especulativa (SD) para la inferencia de contexto medio a largo (4K–16K tokens). Aunque la SD acelera la generación autorregresiva (AR) utilizando un modelo "borrador" (drafter) pequeño para proponer tokens que un modelo "verificador" más grande debe validar, el despliegue estándar enfrenta un desajuste entre lo disperso y lo completo (sparse/full mismatch).

• La Restricción: Para gestionar el pico de memoria de GPU (VRAM) y la latencia, el borrador suele verse obligado a operar con un caché KV disperso (sparse) bajo un presupuesto fijo, mientras que el verificador mantiene un caché KV completo para asegurar la calidad de la salida.
• El Modo de Fallo: A medida que la longitud del contexto aumenta de 4K a 16K, la alineación entre el borrador disperso y el verificador completo se degrada rápidamente. El artículo identifica un "colapso de 16K", donde las tasas de aceptación caen a casi cero, invalidando la eficacia de la SD (aceleración $\le$ 1$\times$).
• El Problema de la Sensibilidad: El borrado disperso ingenuo es altamente sensible al presupuesto específico de KV. Un borrador ajustado para un solo presupuesto suele fallar cuando se despliega bajo diferentes restricciones de memoria, y la relación entre el tamaño del presupuesto y la tasa de aceptación es no monótona (presupuestos más pequeños a veces producen tasas de aceptación más altas que presupuestos mayores).

2. Metodología: BudgetDraft

Los autores proponen BudgetDraft, una solución en tiempo de entrenamiento que crea un único borrador robusto al presupuesto, capaz de mantener tasas de aceptación estables a través de diversos presupuestos de KV sin requerir componentes adicionales en tiempo de inferencia.

Arquitectura Central:
El método emplea una estrategia de Entrenamiento Disperso Multi-Vista donde el borrador (estudiante) es entrenado para alinearse con un verificador congelado (profesor) bajo múltiples condiciones de dispersión simultáneamente.

• Objetivos del Profesor: El verificador genera objetivos del profesor por codicia (greedy teacher targets, $x^*_t$) utilizando su caché KV completo.
• Función de Pérdida de Doble Rama: El objetivo de entrenamiento combina dos pérdidas complementarias:
  1. Pérdida Sensible a la Aceptación ($L_A$): Una rama de caché completo donde el borrador predice los tokens de codicia del profesor utilizando un prefijo de caché completo. Esto asegura la alineación directa con el criterio de aceptación del verificador (coincidencia top-1).
  2. Pérdida Dispersa Multi-Vista ($L_C$): Una rama de caché disperso donde el borrador procesa la misma secuencia pero con un presupuesto de KV muestreado aleatoriamente ($B$) de un conjunto (por ejemplo, $\{256, 512, 1024, 2048\}$). El caché disperso se construye fragmentando el prefijo y reteniendo los fragmentos superiores basados en los pesos de atención acumulados.
• Objetivo de Entrenamiento: La pérdida total es $L = L_A + \lambda L_C$. Al exponer al borrador a diversas vistas dispersas del mismo prefijo durante el entrenamiento, este aprende a generalizar a través de niveles de dispersión en lugar de sobreajustarse a un solo presupuesto.
• Inferencia: En el despliegue, el borrador entrenado utiliza un caché KV disperso con un presupuesto elegido $B$ para proponer tokens. El verificador los valida utilizando su caché completo. No se añaden cambios estructurales ni modelos intermedios al flujo de inferencia.

3. Contribuciones Clave

1. Caracterización del Colapso de Aceptación: El artículo establece empíricamente que la decodificación especulativa dispersa ingenua sufre una severa caída en la tasa de aceptación a medida que la longitud del contexto se desplaza de 4K a 16K, identificando los 16K como un límite de fallo práctico. También revela la relación no monótona entre el tamaño del presupuesto de KV y la tasa de aceptación.
2. Marco de Trabajo BudgetDraft: La propuesta de un método de entrenamiento multi-vista que combina la alineación sensible a la aceptación con el entrenamiento disperso de múltiples presupuestos. Esto produce un borrador invariante al presupuesto que recupera las tasas de aceptación en todos los niveles de dispersión sin sobrecarga adicional en tiempo de inferencia.
3. Validación Empírica: Demostración de que BudgetDraft logra aceleraciones significativas sobre la decodificación AR estándar mientras mantiene la eficiencia de memoria, específicamente en el régimen desafiante de 8K–16K.

4. Resultados Experimentales

Los experimentos se realizaron en tres conjuntos de datos (PG-19, LongBench, LWM) utilizando un borrador de 68 millones de parámetros y un verificador de 7 mil millones de parámetros en una única GPU NVIDIA A100.

• Rendimiento de Aceleración:
  • Contexto de 4K: Logró una aceleración de hasta 6.55$\times$ sobre AR.
  • Contexto de 8K: Logró una aceleración de hasta 4.46$\times$ sobre AR.
  • Contexto de 16K: Logró una aceleración de hasta 2.10$\times$ sobre AR (donde la SD ingenua a menudo falla por completo).
• Robustez: A diferencia de la SD base (disperso/completo), que muestra caídas bruscas de rendimiento o aceptación cercana a cero a medida que aumenta la longitud del contexto o cambian los presupuestos, BudgetDraft mantiene tasas de aceptación estables a través de diferentes presupuestos de KV ($B \in \{256, 512, 1024, 2048\}$).
• Comparaciones: BudgetDraft superó a métodos de mitigación estructural como TriForce y métodos de borrado avanzados como EAGLE-3 en el régimen de 8K–16K. Por ejemplo, en 16K sobre LWM, BudgetDraft logró una aceleración de ~1.94$\times$, comparado con 1.19$\times$ para TriForce y 1.36$\times$ para EAGLE-3.
• Ablación: Eliminar la pérdida de dispersión multi-vista ($L_C$) resultó en un rendimiento inestable bajo cambios de presupuesto, confirmando la necesidad del muestreo de múltiples presupuestos durante el entrenamiento.

5. Significado y Reivindicaciones

El artículo afirma que BudgetDraft resuelve el problema del "desajuste entre lo disperso y lo completo" que actualmente limita el despliegue práctico de la decodificación especulativa en escenarios de contexto largo. Al trasladar la complejidad a la fase de entrenamiento, permite un flujo de inferencia simple y amigable con la memoria que es robusto ante diversas restricciones de hardware (disponibilidad de VRAM) y concurrencia de cargas de trabajo.

Los autores enfatizan que su enfoque no requiere cambios estructurales complejos (como cachés de recuperación) ni múltiples modelos en tiempo de inferencia. En su lugar, proporciona un único borrador ligero que puede adaptarse a diferentes presupuestos de despliegue, haciendo que la decodificación especulativa sea viable para aplicaciones de contexto medio a largo (4K–16K) donde antes era ineficaz.

Limitaciones Reconocidas por el Artículo:

• Extrapolación Posicional: El borrador de 68M tiene un límite nativo de incrustación posicional (2048), lo que requiere extrapolación para contextos de 8K–16K, lo cual aún puede limitar las tasas de aceptación.
• Especificidad del Verificador: El borrador es entrenado contra un verificador específico; se requiere reentrenamiento si el modelo verificador cambia.
• Decodificación por Codicia (Greedy): Los resultados se basan en la decodificación por codicia; la extensión a la decodificación basada en muestreo (ej. muestreo de temperatura) se deja para trabajos futuros.