Resumen Técnico: Más allá de la reconstrucción de KV – Reconstrucción Funcional para Modelos Draft de MLA en Decodificación Especulativa

Declaración del Problema

La atención de múltiples cabezales latentes (MLA) ha emergido como una arquitectura crítica para la inferencia de Modelos de Lenguaje de Gran Escala (LLM) de contexto largo, reemplazando el creciente caché de clave-valor (KV) con estados latentes compactos para reducir el tráfico de memoria. Aunque los métodos existentes (p. ej., TransMLA, MHA2MLA) permiten la conversión de modelos pre-entrenados de Atención de Múltiples Cabezales (MHA) o Atención de Consulta Agrupada (GQA) hacia MLA sin un reentrenamiento costoso, estas conversiones estructurales suelen fallar cuando el modelo convertido se utiliza como un modelo draft (modelo de borrador) en la decodificación especulativa.

El problema central es que la decodificación especulativa depende de una alta concordancia entre las propuestas de tokens del modelo draft y la verificación del modelo objetivo. La conversión directa de MHA/GQA a MLA introduce dos fuentes primarias de error que degradan esta concordancia:

1. Errores de factorización de bajo rango: La compresión de las proyecciones KV en un espacio latente de menor dimensión no puede reproducir exactamente los mapas K/V originales cuando el rango original supera el rango latente.
2. Errores de manejo de RoPE: La descomposición de los Embeddings de Posición Rotatorios (RoPE) durante la conversión introduce residuales en las rutas posicionales retenidas.

Si bien estos errores pueden ser tolerables para la generación independiente, causan caídas significativas en las tasas de aceptación de tokens durante la decodificación especulativa, convirtiendo al modelo draft en una carga computacional en lugar de un acelerador. Los métodos de conversión existentes se centran en la validez estructural (eficiencia del caché) en lugar de la fidelidad funcional (acuerdo de la propuesta).

Metodología

El artículo propone la Reconstrucción Funcional, un procedimiento de post-conversión de extremo a extremo (E2E) y agnóstico al convertidor, diseñado para optimizar la calidad del draft sin alterar la estructura del caché convertido ni requerir la supervisión del verificador.

Mecanismo Central

El método trata la construcción del draft de MLA como un problema de reconstrucción funcional en lugar de un problema de compresión de caché. Opera en tres etapas:

1. Inicialización: Un convertidor estándar (TransMLA o MHA2MLA) proporciona un checkpoint de MLA inicial capaz de gestionar el caché.
2. Objetivo de Entrenamiento: El método entrena únicamente las proyecciones de consulta (query) y KV introducidas por el convertidor. La proyección de salida ($W_O$) y los bloques MHA/GQA originales permanecen congelados.
   • Para un estado oculto de calibración dado $h_\ell$, el sistema calcula la salida del bloque GQA original congelado ($A^{orig}_\ell(h_\ell)$) y el bloque MLA convertido entrenable ($A^{MLA}_\ell(h_\ell)$).
   • La función de pérdida minimiza el error cuadrático medio (MSE) enmascarado entre las respuestas post-proyección de salida de ambos bloques:
     $$\mathcal{L}^{(\ell)}_{E2E} = \mathbb{E}_b \left[ \frac{\| M_b \odot (A^{MLA}_\ell(H^b_\ell) - \text{sg}[A^{orig}_\ell(H^b_\ell)]) \|_F^2}{\sum_t M_{b,t}} \right]$$
     donde $M_b$ es una máscara para tokens válidos y $\text{sg}$ denota el stop-gradient (detener el gradiente).
3. Alcance: El objetivo cubre toda la función de atención desde la entrada de la capa hasta la salida post-proyección de salida, abordando conjuntamente los errores de conversión de KV de bajo rango y los errores de manejo de RoPE.

Características Clave

• Agnóstico al Convertidor: El mismo objetivo refina los checkpoints de diferentes convertidores (TransMLA, MHA2MLA) preservando sus topologías de caché específicas.
• Sin Supervisión del Verificador: No requiere logits, activaciones o etiquetas de benchmark del modelo objetivo, basándose únicamente en los estados ocultos de calibración del modelo original.
• Libre de Inferencia: No introduce nuevos parámetros, ramas o estados de caché durante la inferencia.

Contribuciones Clave

1. Distinción de Objetivos: El artículo establece que la conversión estructural (asegurar una parametrización de MLA válida) y el drafting (asegurar el acuerdo de la propuesta para la decodificación especulativa) son objetivos distintos. Un modelo puede ser estructuralmente válido pero funcionalmente pobre como draft.
2. Optimización en Tiempo de Entrenamiento: Demuestra que la calidad del draft puede optimizarse enteramente en el tiempo de entrenamiento utilizando un objetivo funcional que iguala el comportamiento de atención del modelo original, sin necesidad de supervisión del verificador.
3. Aplicabilidad General: Se demuestra que el método es agnóstico al convertidor, mejorando o manteniendo el rendimiento a través de diferentes convertidores (TransMLA, MHA2MLA), familias de modelos (Llama, Qwen) y backends (Hugging Face, vLLM).

Resultados Experimentales

Los autores evaluaron 192 configuraciones que abarcan cuatro pares draft-target (Llama-1B/3B $\to$ Llama-8B; Qwen-1.5B/3B $\to$ Qwen-7B), dos convertidores, dos backends y cuatro tareas (HumanEval, Alpaca, Natural Questions, CNN/DailyMail).

• Tasas de Aceptación: En comparación con una línea base de "Partial RoPE" (un método de reconstrucción local del convertidor), la Reconstrucción Funcional produjo mejoras materiales en la aceptación en 37 de 64 celdas de tareas emparejadas (usando una tolerancia de 0.5 puntos porcentuales).
  • 26 celdas permanecieron prácticamente sin cambios.
  • 1 celda mostró una disminución material.
• Throughput (Rendimiento): Entre las 37 celdas con mejora en la aceptación, 12 también mostraron aumentos materiales en el throughput. El método no introduce sobrecarga de inferencia, por lo que las ganancias de throughput son impulsadas por mayores tasas de aceptación.
• Efectos de Interacción: Los resultados revelaron un patrón de "cruce":
  • Llama + TransMLA: Consistentemente mejoró.
  • Qwen + MHA2MLA: Consistentemente mejoró.
  • Llama + MHA2MLA: Mayormente sin cambios (sugiriendo que la conversión inicial ya era funcionalmente cercana).
  • Qwen + TransMLA: Resultados mixtos, con cierta sensibilidad al backend.
• Límites de Recuperación: Aunque el método recupera una parte significativa de la aceptación perdida durante la conversión, no siempre restaura completamente el rendimiento del draft GQA original no convertido, particularmente cuando la conversión inicial o la implementación del backend introducen desajustes severos (p. ej., semántica de paralelismo de tensores).

Significancia y Reivindicaciones

El artículo argumenta que la conversión de MLA y la construcción de drafts de MLA deben tratarse como tareas relacionadas pero separadas.

• Éxito Jerárquico: Los autores definen tres niveles de éxito: validez estructural (parámetros ejecutables), fidelidad funcional (igualar mapas de atención) y utilidad de draft (aceptación especulativa). Los convertidores actuales a menudo satisfacen el primero pero fallan en los otros dos.
• Interfaz Funcional: La significancia de la Reconstrucción Funcional radica en situar el objetivo de optimización en una interfaz funcional compartida (post-$W_O$ salida) en lugar de asumir que cada convertidor deja el mismo error de distribución. Actúa como un "control deslizante de entrenamiento" para refinar el modelo draft dentro de las restricciones fijas del rango latente y el presupuesto de caché.
• Implicación Práctica: El trabajo sugiere un flujo de trabajo de tres niveles para futuros pipelines de MLA: (1) establecer parámetros eficientes de caché, (2) aplicar la reconstrucción funcional para optimizar el mapa de atención, y (3) validar mediante la aceptación emparejada. Enfatiza que, si bien la reconstrucción funcional es un refinamiento general, no puede compensar los desajustes fundamentales en el rango latente o la implementación del backend (p. ej., semántica de paralelismo de tensores).

El artículo concluye que, si bien la Reconstrucción Funcional mejora materialmente la aceptación en muchos escenarios, lograr una aceleración especulativa completa aún requiere convertidores iniciales de alta calidad y kernels de MLA optimizados.