Resumen Técnico: Deriva de Atención en la Decodificación Especulativa Autoregresiva

Declaración del Problema

La decodificación especulativa acelera la inferencia de Modelos de Lenguaje Grande (LLM) utilizando un modelo "redactor" ligero para predecir tokens futuros, los cuales son luego verificados por un modelo "objetivo" más grande. Aunque efectiva, los modelos redactores actuales (específicamente EAGLE-3 y las cabezas de Predicción de Múltiples Tokens) exhiben una fragilidad significativa bajo variaciones de implementación, como perturbaciones de plantillas y entradas de contexto largo. Los practicantes a menudo deben reentrenar redactores para motores de inferencia o prompts específicos para mantener el rendimiento, lo que sugiere una inestabilidad arquitectónica subyacente.

El artículo identifica un fenómeno previamente no reportado llamado Deriva de Atención. A medida que un redactor genera una secuencia de tokens especulativos, su mecanismo de atención se desplaza progresivamente alejándose del "sumidero de atención" (un token ancla estable, típicamente cerca del inicio del prompt) y se concentra cada vez más en sus propios tokens generados recientemente. Esta deriva se correlaciona con una degradación en las tasas de aceptación, particularmente en escenarios de contexto largo y cuando las plantillas de entrada se desvían de las distribuciones de entrenamiento.

Metodología y Análisis

Los autores investigan las causas raíz de la deriva de atención mediante una combinación de observación empírica y análisis arquitectónico a través de múltiples familias de modelos (Llama 3.1, Qwen 3/3.5 y GPT-oss).

1. Observación Fenomenológica

Utilizando mapas de calor de atención y métricas cuantitativas, los autores demuestran que:

• Migración del Sumidero: En modelos con sumideros de atención fuertes (por ejemplo, Llama, Qwen), la atención del redactor hacia el token sumidero disminuye a medida que la cadena de especulación se profundiza, mientras que la atención hacia el token redactado más reciente aumenta.
• Universalidad: Esta deriva se observa tanto en redactores EAGLE-3 (entrenados post-hoc sobre objetivos congelados) como en cabezas MTP (entrenadas conjuntamente con el objetivo), lo que sugiere que es una propiedad de diseños autoregresivos que consumen sus propias salidas.
• Ausencia de Sumidero: Incluso en modelos diseñados para suprimir sumideros (por ejemplo, GPT-oss con atención con compuertas o sesgos aprendibles), ocurre una forma de deriva donde la atención se concentra en tokens recurrentes de la plantilla en lugar del inicio del prompt.

2. Análisis de la Causa Raíz: Magnitud del Estado Oculto

El artículo rastrea la deriva de atención hacia la ruta residual no normalizada entre pasos de especulación en arquitecturas pre-norm estándar:

• Crecimiento Monótono: En transformadores pre-norm, las conexiones residuales acumulan magnitud sin normalización. Los autores muestran que la magnitud del estado oculto del redactor ($\|h^*_k\|$) crece monótonamente con la profundidad de especulación ($k$).
• Apilamiento de Capas vs. Autoregresión: Este crecimiento hace que el redactor se comporte menos como un predictor autoregresivo independiente y más como una secuencia de capas adicionales de transformador apiladas sobre el objetivo. En el paso $k$, el redactor aproxima efectivamente el modelo objetivo a una profundidad $N+k$, en lugar de predecir el siguiente token basado en una distribución estable.
• Entradas Desbalanceadas: Los estados ocultos fusionados del modelo objetivo ($h_{FC}$) a menudo están dominados por capas más profundas con magnitudes mayores, exacerbando aún más el desajuste de escala cuando se alimentan al redactor.

3. Intervenciones Propuestas

Para contrarrestar la acumulación de magnitud y estabilizar el proceso de redacción, los autores proponen dos modificaciones arquitectónicas:

1. Post-Norm en Estados Ocultos del Redactor: Aplicar RMSNorm a los estados ocultos del redactor después de la conexión residual (entre pasos de especulación) previene el crecimiento monótono de la magnitud. Esto obliga al redactor a aprender una regla autoregresiva invariante a la profundidad en lugar de refinamientos específicos de la profundidad.
2. RMSNorm por Estado Oculto: Añadir capas RMSNorm antes de fusionar los estados ocultos del objetivo ($h_{low}, h_{mid}, h_{high}$) en la entrada del redactor para equilibrar la magnitud de la señal entrante.

Los autores también evalúan la Atención con Compuertas (suprimiendo la dependencia de un sumidero fijo) como una intervención separada. Descubren que, aunque la compuerta elimina el sumidero, no previene el crecimiento de la magnitud, y combinar la compuerta con post-norm puede llevar a una sobre-regularización (colapso de entropía).

Resultados Clave

La arquitectura post-norm propuesta fue evaluada contra el actual estado del arte pre-norm EAGLE-3 en diversas pruebas y condiciones:

• Perturbación de Plantillas: Bajo perturbaciones de plantillas (por ejemplo, eliminar plantillas de chat o tokens de inicio de secuencia), el redactor post-norm mejoró la longitud de aceptación hasta en 2× en comparación con pre-norm. Los modelos pre-norm sufrieron caídas catastróficas (hasta un 52% de pérdida), mientras que los modelos post-norm permanecieron resilientes (máximo 5% de pérdida).
• Tareas de Contexto Largo: En tareas de contexto largo (LongBench, chat multi-turno con atención de ventana deslizante), post-norm mejoró la longitud de aceptación en 1.18×. En configuraciones de contexto largo con atención completa, los modelos pre-norm colapsaron a una aceptación cercana a cero, mientras que post-norm mantuvo un rendimiento utilizable.
• Pruebas Estándar: A través de siete pruebas estándar (chat multi-turno, matemáticas, codificación), post-norm mostró una mejora consistente de 1.10× en la longitud de aceptación.
• Eficiencia de Entrenamiento: Los redactores post-norm se generalizan a profundidades de inferencia más allá de su horizonte de entrenamiento. Esto permite entrenar con profundidades "tiempo-entrenamiento-tiempo-prueba" (TTT) más cortas (por ejemplo, reduciendo de 8 a 4 pasos), recortando el tiempo de entrenamiento en aproximadamente un tercio sin sacrificar el rendimiento.
• Robustez al Ruido: Las arquitecturas post-norm toleran significativamente más ruido en los estados ocultos (retención de aceptación del 58% en ruido $\alpha=0.5$ frente al 5% para pre-norm), lo que sugiere una mejor robustez ante la cuantización y cambios de distribución.

Significado y Afirmaciones

El artículo afirma que la Deriva de Atención es un modo de fallo crítico, previamente no reportado, en la decodificación especulativa que surge de la interacción entre rutas residuales no normalizadas y la retroalimentación autoregresiva.

• Diagnóstico: Los autores argumentan que la deriva de atención es un síntoma de un problema más profundo: la acumulación de la magnitud del estado oculto, lo que hace que el redactor aprenda transformaciones dependientes de la profundidad en lugar de una predicción de token estable.
• Solución: Afirman que simplemente fijar el sumidero de atención (por ejemplo, mediante compuertas) es insuficiente; deben abordarse las dinámicas subyacentes de la magnitud. La arquitectura post-norm propuesta se presenta como una solución simple y efectiva que estabiliza el proceso de redacción.
• Generalización: Los hallazgos sugieren que post-norm debería considerarse una opción arquitectónica predeterminada para redactores de decodificación especulativa, particularmente para sistemas de producción que requieren robustez frente a prompts variables y contextos largos.

El artículo concluye que nombrar y caracterizar la "deriva de atención" proporciona un diagnóstico útil para la robustez del redactor, y que las intervenciones arquitectónicas dirigidas a la acumulación de magnitud son esenciales para una decodificación especulativa confiable.