Resumen Técnico: Atención Estocástica

Enunciado del Problema

Los mecanismos de atención eficientes son críticos para escalar los modelos Transformer, sin embargo, los enfoques existentes enfrentan una compensación entre el costo computacional y la cobertura del campo receptivo. La Atención de Ventana Deslizante (SWA) reduce la complejidad a $O(nw)$ por capa al restringir los tokens a una ventana local de tamaño $w$. No obstante, esta localidad determinista limita el campo receptivo efectivo a un crecimiento lineal ($\ell w$ después de $\ell$ capas), dejando grandes porciones de la secuencia inalcanzables cuando $w \ll n$. Aunque existen remedios como tokens globales, patrones dispersos diseñados a mano o enrutamiento a nivel de bloques, estos a menudo introducen complejidad arquitectónica o fallan en proporcionar garantías de cobertura demostrables. Por el contrario, la atención completa ofrece conectividad global pero incurre en costos prohibitivos de $O(n^2)$.

El artículo se inspira en el conectoma de la mosca de la fruta (Drosophila melanogaster), una red de ~130,000 neuronas con una probabilidad de conexión de solo el 0.02%. A pesar de esta dispersión, la red logra una longitud de camino promedio corta de ~4.4 saltos. Esta topología de "mundo pequeño" surge de la interacción entre el agrupamiento local denso y las conexiones de largo alcance ampliamente distribuidas que funcionan como atajos estocásticos. Los autores postulan que un flujo de información global eficiente puede emerger de la computación local estructurada, aumentada por conexiones de largo alcance distribuidas y estocásticas, en lugar de requerir una conectividad densa de todos a todos.

Metodología: Atención Estocástica (SA)

Los autores proponen la Atención Estocástica (SA), una mejora sin parámetros para SWA que imita la estrategia de enrutamiento del conectoma.

Mecanismo Central

SA transforma una ventana local fija en una global estocástica mediante un proceso de tres etapas aplicado antes y después del cálculo estándar de SWA:

1. Permutar: Se extrae una permutación aleatoria $\sigma \sim \text{Uniforme}(S_n)$ independientemente para cada capa. Las representaciones de la secuencia de tokens ($Q, K, V$) se reorganizan según $\sigma$.
2. Atención Ventanada: Se aplica SWA estándar en el espacio permutado. Debido a que la permutación es aleatoria, la ventana fija de tamaño $w$ ahora abarca un subconjunto aleatorio de la secuencia original.
3. Deshacer Permutación: Se restaura el orden original aplicando $\sigma^{-1}$ a la salida.

En el espacio de tokens original, la posición $i$ atiende a un vecindario aleatorio $\tilde{N}^\sigma_w(i)$ de tamaño esperado $w$, distribuido uniformemente a través de toda la secuencia independientemente de la distancia original $|i-j|$. Esto crea una probabilidad de conexión par a par de aproximadamente $w/n$ para cualquier par de tokens, análogo a los atajos de largo alcance distribuidos en el conectoma.

Propiedades Teóricas

• Crecimiento del Campo Receptivo: Mientras que el campo receptivo de SWA crece linealmente ($O(n/w)$ capas para cobertura completa), las permutaciones muestreadas independientemente a través de las capas en SA generan campos receptivos de crecimiento exponencial. El artículo demuestra que la cobertura completa de la secuencia se logra en $O(\log_w n)$ capas.
• Complejidad: Las operaciones de permutación e inversa de permutación son reordenamientos de índices de $O(n)$. El cálculo de atención permanece en $O(nw)$. Por lo tanto, SA mantiene el mismo presupuesto de $O(nw)$ por capa que SWA estándar.
• Restricciones Causales: En configuraciones autorregresivas, la máscara causal se aplica en el espacio permutado, asegurando que los tokens solo atiendan a predecesores causalmente accesibles mientras habilitan conexiones de largo alcance estocásticas.

Combinación de SA con Puerta + SWA

Para reproducir el régimen de "mundo pequeño" (alto agrupamiento local + caminos cortos), los autores proponen una arquitectura de doble camino que combina SA y SWA:
$$Y = g_{sa} \odot Y_{sa} + g_{swa} \odot Y_{swa}$$
donde $Y_{sa}$ y $Y_{swa}$ son las salidas de los caminos estocástico y de ventana deslizante, respectivamente. Los pesos de puerta ($g_{sa}, g_{swa}$) se calculan mediante proyecciones de sigmoide independientes, permitiendo que el modelo aprenda a equilibrar la coherencia local (de SWA) y la cobertura global (de SA). Esta combinación añade parámetros insignificantes y mantiene una complejidad de $O(nw)$.

Contribuciones Clave

1. Atención Estocástica (SA): Una mejora lista para usar, sin parámetros, para SWA que aplica permutaciones aleatorias para lograr un crecimiento exponencial del campo receptivo ($O(\log_w n)$) dentro del mismo presupuesto computacional de $O(nw)$.
2. SA con Puerta + SWA: Una arquitectura de doble camino que reproduce teórica y empíricamente el régimen de mundo pequeño observado en conectomas biológicos, equilibrando el agrupamiento local y los atajos estocásticos.
3. Validación Empírica:
   • Pre-entrenamiento: En modelos de 360M parámetros entrenados desde cero, la combinación de SA con puerta + SWA logra la mejor precisión promedio en cero disparos, superando a SWA puro y igualando la atención completa en perplejidad mientras la supera en tareas posteriores.
   • Inferencia sin Entrenamiento: Aplicado a modelos preentrenados Qwen3-8B y Qwen3-30B-A3B, SA supera consistentemente a SWA y iguala o excede a Mezcla de Atención de Bloques (MoBA) con presupuestos computacionales comparables. Recupera la calidad de la atención completa significativamente más rápido a medida que aumenta el tamaño de la ventana.

Resultados

• Pre-entrenamiento: El modelo SA + SWA logró una precisión promedio en cero disparos de 35.9 (vs. 35.1 para SWA y 34.9 para Atención Completa). Mejoró notablemente el rendimiento en LAMBADA (que requiere contexto discursivo amplio) mientras mantenía una baja perplejidad en WikiText. SA puro sufrió una mayor perplejidad, confirmando la necesidad de coherencia local para el modelado del lenguaje.
• Inferencia en Qwen3:
  • Con un tamaño de ventana efectivo de 128, SA en Qwen3-8B alcanzó una precisión promedio del 70.9% (dentro de 1 punto del baseline de atención completa), mientras que SWA se quedó atrás en 62.2%.
  • En Qwen3-30B-A3B, SA alcanzó una precisión del 73.2% con un tamaño de ventana de 64, superando significativamente a SWA (47.0%) y MoBA (66.3%).
  • SA demostró robustez en tareas intensivas en conocimiento (MMLU, BoolQ) incluso con tamaños de ventana muy pequeños (ej. $w=32$), donde el rendimiento de SWA colapsó.
• Eficiencia: SA mantiene una escalabilidad lineal $O(nw)$. En una GPU A100, SA mostró aceleraciones de 1.5× a una longitud de secuencia de 2,048, aumentando a 28× a 32,768 en comparación con la atención completa.

Significado y Afirmaciones

El artículo afirma que el enrutamiento estocástico inspirado en el conectoma es un primitivo práctico para mejorar la expresividad de la atención eficiente. Los autores argumentan que el flujo de información global no requiere una conectividad densa de todos a todos, sino que puede emerger de la interacción entre la computación local estructurada y los atajos de largo alcance dispersos y estocásticos acumulados a través de la profundidad.

El significado de este trabajo radica en:

1. Simplicidad: SA no requiere parámetros aprendibles en el mecanismo de atención en sí y solo añade una sobrecarga de índices de $O(n)$, lo que lo convierte en una actualización "lista para usar" para los kernels SWA existentes (ej. en Mistral, Gemma 2, gpt-oss).
2. Garantía Teórica: Proporciona una profundidad de cobertura de $O(\log_w n)$ demostrable, una mejora significativa sobre la cobertura lineal de SWA.
3. Complementariedad: Se presenta como complementario a los enfoques lineales y dispersos existentes, ofreciendo un nuevo principio de diseño que cierra la brecha entre la eficiencia local y la expresividad global sin la complejidad de patrones dispersos diseñados a mano o tokens globales.