Resumen Técnico: Proxies Singulares para Caché Adaptativa en Modelos de Lenguaje de Difusión

Declaración del Problema

Los Modelos de Lenguaje de Difusión (DLM) ofrecen una alternativa flexible a los modelos autoregresivos (AR) al permitir la decodificación de orden arbitrario y la generación paralela. Sin embargo, su naturaleza no causal y su dependencia de la atención bidireccional hacen que la caché estándar de claves y valores (KV), piedra angular de la eficiencia de los AR, sea incompatible. En los DLM, los estados ocultos no pueden reutilizarse simplemente sin recálculo para capturar dependencias globales en evolución, lo que conduce a una latencia de inferencia prohibitiva (complejidad $O(T \times N^2)$).

Los intentos existentes de mitigar esto mediante caché enfrentan dos limitaciones principales:

1. Identificación Costosa: Los métodos actuales dependen de cálculos de similitud de alta dimensión (por ejemplo, monitorear la deriva del estado de Valor) para identificar qué tokens requieren actualizaciones. Esta sobrecarga de identificación a menudo compensa las ganancias de las actualizaciones dispersas a medida que aumentan las longitudes de secuencia.
2. Asignación de Presupuesto Rígida: Los enfoques existentes suelen aplicar un presupuesto de actualización uniforme en todas las capas del modelo. Esto no tiene en cuenta la dinámica heterogénea de los estados ocultos, donde algunas capas son estables (requieren pocas actualizaciones) y otras son volátiles (requieren actualizaciones frecuentes), lo que conduce a redundancia computacional o degradación del rendimiento.

Metodología: SPA-Cache

Los autores proponen SPA-Cache (Proxies Singulares para Caché Adaptativa), un marco que optimiza conjuntamente la identificación de actualizaciones y la asignación de presupuesto. El flujo de trabajo para una sola capa Transformer se divide en tres fases:

1. Identificación de Actualizaciones mediante Proxies Singulares

En lugar de calcular la similitud en el espacio completo de alta dimensión ($d$), SPA-Cache introduce un proxy singular para identificar tokens críticos para la actualización en un subespacio de baja dimensión ($r \ll d$).

• Base Teórica: Los autores demuestran que la estabilidad de los estados de Valor es una condición suficiente para la estabilidad de las salidas de atención y las salidas subsiguientes de la red neuronal feed-forward (FFN). Por lo tanto, monitorear los estados de Valor proporciona una garantía conservadora para la estabilidad de la salida.
• Reducción de Dimensionalidad: Para evitar el cuello de botella computacional de la proyección de alta dimensión, el método aplica la Descomposición en Valores Singulares (SVD) a la matriz de proyección de Valor $W$. Conserva solo los $r$ vectores y valores singulares superiores para crear una proyección truncada $W_r$.
• Mecanismo: Los estados de entrada se proyectan en esta variedad de baja dimensión. Se calcula la similitud del coseno entre el proxy singular actual y el proxy almacenado en caché. Se seleccionan los tokens con la menor similitud (mayor deriva) para su recálculo.
• Eficiencia: Esto reduce la complejidad de identificación de $O(d)$ a $O(r)$ y la sobrecarga de proyección de $O(d^3)$ a $O(rd^2)$.

2. Asignación de Presupuesto Adaptativa

Reconociendo que la deriva del estado oculto varía significativamente entre las capas del modelo, SPA-Cache reemplaza las proporciones de actualización estáticas y uniformes con una estrategia dinámica y adaptativa por capas.

• Observación: La deriva es baja en las primeras capas (transformaciones de incrustación), alcanza su punto máximo en las capas intermedias (transformaciones semánticas complejas) y se estabiliza en las últimas capas.
• Implementación: La proporción de actualización $\rho(l)$ para la capa $l$ se parametriza utilizando una función gaussiana por partes. Esto asigna un presupuesto de actualización más alto a las capas intermedias volátiles y almacena en caché de manera agresiva las capas iniciales y finales estables, maximizando el rendimiento sin sacrificar la calidad.

3. Ejecución Dispersa

Los tokens seleccionados (índices $I$) se utilizan para generar Consultas, Claves y Valores dispersos. El mecanismo de Atención calcula las salidas utilizando estas consultas dispersas frente a una caché KV parcialmente actualizada. La Red Neuronal Feed-Forward (FFN) procesa solo las salidas de atención dispersas, y la salida de la capa final se actualiza solo en los índices seleccionados, dejando que otras posiciones reutilicen las características almacenadas en caché.

Contribuciones Clave

1. Justificación Teórica para Proxies de Valor: El artículo proporciona pruebas formales (Teoremas 3.1 y 3.2) que demuestran que la similitud del estado de Valor limita la divergencia de las salidas de atención y FFN, validando el uso de los estados de Valor como un identificador principiado para la selección de actualizaciones.
2. Proxy Singular para Identificación de Baja Sobrecarga: Al aprovechar la SVD truncada, los autores introducen un método para identificar tokens críticos en un subespacio de baja dimensión, reduciendo significativamente el costo computacional del propio proceso de identificación.
3. Asignación de Presupuesto Adaptativa: La introducción de una estrategia de actualización dinámica que alinea los recursos computacionales con la heterogeneidad por capas, evitando las ineficiencias de la asignación uniforme.
4. Ortogonalidad a la Decodificación Paralela: SPA-Cache está diseñado para ser ortogonal a las optimizaciones a nivel de secuencia y puede integrarse con marcos de decodificación paralela para amplificar las ganancias de rendimiento.

Resultados Experimentales

Los autores evaluaron SPA-Cache en dos DLM (LLaDA-8B y Dream-v0-Instruct-7B) a través de siete puntos de referencia (incluyendo GSM8K, MMLU-PRO, MBPP y HumanEval).

• Mejoras en el Rendimiento: SPA-Cache logra hasta una mejora de rendimiento de 8× sobre la decodificación estándar y una aceleración de 2–4× sobre las líneas base de caché existentes (dLLM-Cache, Fast-dLLM).
• Puntos de Referencia Específicos:
  • En MBPP con LLaDA-8B, SPA-Cache obtuvo una aceleración de 8.0×.
  • En MMLU-PRO, logró aceleraciones de 6.0× (LLaDA-8B) y 6.2× (Dream-7B).
• Integración con Decodificación Paralela: Cuando se combina con métodos de decodificación paralela, SPA-Cache produjo una aceleración total de 28× sobre la decodificación estándar en el punto de referencia BBH, manteniendo una precisión comparable.
• Estudios de Ablación:
  • Reemplazar los identificadores de Valor de dimensión completa con el proxy singular ($r=128$) mejoró los Tokens por Segundo (TPS) de ~165 a ~179 sin pérdida de precisión.
  • La asignación de presupuesto adaptativa aumentó aún más los TPS a ~189 al reducir la proporción de actualización promedio del 25% al 16%, mientras que una proporción uniforme del 16% causó degradación de la precisión.

Significado y Afirmaciones

El artículo afirma que SPA-Cache establece una nueva línea base de eficiencia para los Modelos de Lenguaje de Difusión, posicionándolos como una alternativa computacionalmente viable para aplicaciones de alto rendimiento en el mundo real. Al abordar fundamentalmente el cuello de botella de la inferencia mediante un análisis teórico principiado y una asignación adaptativa de recursos, el método permite que los DLM escalen eficazmente sin comprometer la fidelidad de la generación. Los autores enfatizan que su enfoque no requiere entrenamiento y es compatible con las estrategias existentes de decodificación paralela, allanando el camino para el despliegue práctico de modelos de difusión discretos.