Nous Research propone Lighthouse Attention: una atención jerárquica basada en selección solo de entrenamiento que ofrece una aceleración de preentrenamiento de 1,4 a 1,7 veces en un contexto prolongado

Entrenar modelos de lenguaje grandes en secuencias largas tiene un problema bien conocido: la atención es costosa. La atención de producto escalado (SDPA) en el núcleo de cada transformador escala cuadráticamente Θ(N²) tanto en computación como en memoria con una longitud de secuencia N. FlashAttention abordó esto a través de mosaicos con reconocimiento de E/S que evitan materializar la matriz de atención N×N completa en memoria de gran ancho de banda, lo que reduce significativamente la huella de memoria, pero la escala de computación subyacente Θ(N²) permanece. Los investigadores de Nous Research han introducido un nuevo método llamado Lighthouse Attention que aborda este cuello de botella específicamente en el momento del preentrenamiento, logrando una aceleración del reloj de pared de un extremo a otro de 1,40× a 1,69× frente a una línea de base SDPA respaldada por cuDNN, con una pérdida de entrenamiento final coincidente o menor.

El problema central de los métodos de atención escasa existentes

Para comprender por qué Lighthouse funciona como lo hace, es útil saber qué hacen los métodos de atención escasa existentes. La mayoría de los trabajos anteriores, como NSA, HISA, DSA y MoBA, toman las mismas dos decisiones de diseño. Primero, agrupan solo el lado de la clave y el valor y dejan las consultas en resolución completa (compresión asimétrica). En segundo lugar, su lógica de selección reside dentro de un núcleo de atención personalizado, lo que significa que los equipos no pueden reutilizar los núcleos de atención densa optimizados en torno a los cuales se construyen los núcleos tensoriales de GPU modernos.

También existe una preocupación específica del entrenamiento que los métodos dispersos de sólo inferencia no enfrentan. Un método de inferencia con tiempo escaso se evalúa sólo contra su densa columna vertebral y es, como mucho, tan bueno como esa columna vertebral. Un método con poco tiempo de entrenamiento se enfrenta a una prueba más difícil: una vez finalizado el entrenamiento, ¿los pesos resultantes seguirán produciendo un modelo de atención densa competente en la inferencia? Lighthouse trata esa pregunta como su criterio central de corrección.

Lighthouse adopta un enfoque diferente en ambas decisiones de diseño. Agrupa consultas, claves y valores simétricamente en una pirámide de varios niveles y coloca la selección completamente fuera del núcleo de atención. Después de la selección, el sistema reúne las entradas elegidas en una subsecuencia densa y contigua y ejecuta FlashAttention en ella, el mismo núcleo utilizado por la línea de base densa.

Cómo funciona el proceso de cuatro etapas

Una capa de atención Lighthouse envuelve, pero no modifica, la atención del producto escalado. El oleoducto tiene cuatro etapas.

En la primera etapa, la agrupación promedio construye una pirámide de nivel L a partir de Q, K y V. Con el factor de agrupación p, el nivel ℓ de la pirámide tiene N/p^ℓ tokens, cada uno de los cuales resume las posiciones base de p^ℓ. Fundamentalmente, se aplica la misma combinación a las tres proyecciones, lo que produce tripletas coherentes (Q^(ℓ), K^(ℓ), V^(ℓ)) en todos los niveles. La construcción total de la pirámide cuesta Θ(N) tiempo y memoria.

En la segunda etapa, un evaluador sin parámetros asigna a cada entrada de la pirámide dos puntuaciones escalares utilizando normas ℓ₂ por cabeza: una como puntuación de consulta (∥Q^(ℓ)_i∥₂) y otra como puntuación clave (∥K^(ℓ)_i∥₂). Los niveles más gruesos heredan puntuaciones de los más finos a través de la agrupación máxima, por lo que un rango grueso adquiere la importancia de su token más fuerte. Un kernel top-K bitónico fragmentado fusionado selecciona k entradas de forma conjunta en todos los niveles de la pirámide. Un detalle de diseño que vale la pena señalar: el nivel piramidal más grueso siempre se conserva en su totalidad: es barato y garantiza al menos un contribuyente en cada posición base; el resto del presupuesto de selección se gasta en niveles más finos. Además, el diseño bitónico fragmentado produce un top-K estratificado en lugar de un top-K global estricto: el flujo de puntuación se divide en fragmentos de tamaño fijo, cada uno de los cuales mantiene un búfer top-m en el registro, por lo que si las k entradas globalmente con la puntuación más alta se agrupan en un fragmento, algunas serían reemplazadas por entradas con puntuación más baja de otros fragmentos. El resultado es una cobertura de atención más equilibrada en toda la secuencia y evita el colapso de la selección en un lapso estrecho.

El paso K superior es discreto y no diferenciable: no hay estimador directo ni Gumbel softmax. Los índices de selección no llevan gradiente. Los gradientes fluyen sólo a través de las entradas Q, K, V reunidas en WQ, WK, WV, por lo que las proyecciones aprenden a producir valores que son útiles cuando se seleccionan en lugar de puntuaciones que sean buenas para seleccionar.

En la tercera etapa, las entradas seleccionadas se reúnen en una subsecuencia contigua de longitud S = N/p^(L−1) + (L−1)·p·k y se pasan al FlashAttention estándar. En N = 1.000.000 con L = 4, p = 4, k = 4.096, S ≈ 65.000, mucho más pequeño que N. Una propiedad crítica del proceso de recopilación es que no garantiza “huecos” o espacios vacíos en la subsecuencia ensamblada. Esto es importante específicamente porque Lighthouse también comprime las consultas: una brecha en la secuencia significaría que esos tokens faltantes no tienen una ruta de gradiente durante el paso hacia atrás y podrían causar inestabilidades en el entrenamiento. Los métodos asimétricos que dejan consultas en resolución completa no enfrentan este problema, pero el diseño simétrico de Lighthouse requiere que la subsecuencia recopilada permanezca completamente densa.

En la cuarta etapa, cada entrada de salida se dispersa nuevamente a las posiciones base p^ℓ que representa a través de un núcleo de dispersión atómico entero determinista, con un desplazamiento de p^ℓ − 1 para preservar la causalidad. La distribución por posición está limitada por L independientemente de k.

Por qué la agrupación simétrica cambia la computación

La combinación de consultas junto con claves y valores cambia el carácter computacional de la llamada de atención de O (N Sd) a O (S² d) en el momento del entrenamiento. Debido a que S ≪ N en contextos largos, esto es lo que produce la ventaja de latencia. Comparado con una sola NVIDIA B200 en un contexto de 512K (bfloat16, B=1, H=8, dimensión de la cabeza 128, L=3, p=4, escasez ≈ 1:64), Lighthouse es 21 veces más rápido en el pase hacia adelante y 17,3 veces más rápido en el pase combinado hacia adelante y hacia atrás en relación con SDPA respaldado por cuDNN.

Desde un punto de vista asintótico, establecer L = logp(N/k) da un tamaño de subsecuencia recopilado de S = Θ(k log N), lo que hace que la llamada densa FlashAttention cueste Θ(k² log² N d) — polilogarítmica en N en k fijo. Combinado con las etapas de costo lineal (construcción de pirámide, puntuación, dispersión), el cálculo total por capa es Θ(T d) en k acotada (la misma clase asintótica que la atención lineal y los SSM) al tiempo que se preservan las propiedades de recuperación de la atención softmax en la subsecuencia seleccionada.

La inferencia es una restricción diferente. La decodificación autorregresiva presenta una consulta a la vez, lo que viola la suposición de que todas las consultas coexisten en un paso hacia adelante. Lighthouse es un método de solo entrenamiento y el diseño de agrupación simétrica no se puede utilizar directamente en la inferencia.

La receta del entrenamiento en dos etapas y la recuperabilidad

La configuración experimental utilizó un decodificador estilo Llama-3 con parámetros 530M (dmodel=1024, 30 capas, 8 cabezas, dimensión de la cabeza 128, ancho FFN 1536, tokenizador a nivel de bytes), entrenado en C4 en un contexto de 98.304 tokens con AdamW a una tasa de aprendizaje de 2×10⁻³, β1=0,9, β2=0,95, caída de peso. 0.1, calentamiento lineal en pasos de 2k, clip 1 con norma de gradiente, bfloat16 y FSDP. Un detalle de implementación que es importante para los profesionales: de las 30 capas, las capas {0, 1, 28, 29} retienen SDPA denso en todas partes; solo las otras 26 capas usan Lighthouse. La llamada de atención interna dentro de esas 26 capas de Lighthouse utiliza el mismo núcleo SDPA respaldado por cuDNN como línea de base densa.

El enfoque de formación consta de dos etapas. Trenes de la etapa 1 con selección Lighthouse habilitada para la mayor parte del presupuesto de pasos. La etapa 2 reanuda el punto de control de la etapa 1 bajo SDPA denso (mismo estado de optimizador, mismo cargador de datos) durante una cola corta. Si la Etapa 1 hubiera vaciado la capacidad de atención densa del modelo, la recuperación de la Etapa 2 fallaría.

No falla. Al realizar pruebas con un presupuesto total de 16 000 pasos (~50,3 mil millones de tokens), se evaluaron tres puntos de división (10k+6k, 11k+5k, 12k+4k) frente a una línea de base SDPA densa desde cero. En cada punto de reanudación, la pérdida de entrenamiento aumenta transitoriamente entre 1,12 y 1,57 nats a medida que el modelo se ejecuta por primera vez a través de la atención contra la que no fue entrenado, luego se recupera dentro de aproximadamente 1000 a 1500 pasos SDPA y cruza por debajo de la línea de base densa. En el paso 16.000, las tres ejecuciones de Lighthouse reanudadas alcanzan pérdidas finales de 0,6980–0,7102, frente a 0,7237 de la línea de base densa, mientras que pasan de 22,5 h a 27,0 h en reloj de pared en comparación con 37,9 h para SDPA denso desde cero con el mismo presupuesto simbólico.

Ablaciones y rendimiento

La cuadrícula de ablación completa cubre el tipo de anotador, el factor de agrupación p, el número de niveles de la pirámide L y el presupuesto k superior. Hallazgos clave: el puntaje de la norma de proyección está dentro de ~0.01 del puntaje de atención softmax dilatado en cualquier dirección (sin ganador uniforme), pero es aproximadamente un 9% más barato en B200 horas, ya que omite por completo el paso de la atención sobre la pirámide. Las pirámides menos profundas (L=3) superan consistentemente a las más profundas (L=4, L=5) con presupuestos equivalentes. Los valores k más pequeños producen una pérdida posterior a la reanudación más baja dentro del rango probado: la configuración de pérdida más baja en toda la cuadrícula es L=3, p=2, k=1536 con el puntaje dilatado, alcanzando una pérdida final de 0,6825, un resultado contrario a la intuición que los equipos de investigación atribuyen a la selección jerárquica que actúa como un regularizador en esta escala de presupuesto simbólico.

El rendimiento de la etapa 1 en toda la red de ablación oscila entre 84 000 y 126 000 tokens/s/GPU frente a aproximadamente 46 000 para SDPA denso. El anotador de norma de proyección en L=3, p=4, k=1536 encabeza el rango con 126.000 tokens/s/GPU al saltarse por completo el pase de atención dilatada.

Recuperación de contexto largo

Para complementar los resultados de recuperabilidad basados ​​en pérdidas, el equipo de investigación realizó una evaluación simplificada de Needle-in-a-Haystack (NIAH): un único dígito de clave de acceso oculto en un relleno alfanumérico aleatorio a profundidades de 0 a 100% en longitudes de contexto de 4K a 96K tokens, con recuperación calificada como un argmax de un token sobre los tokens de diez dígitos (posibilidad aleatoria: 10%). Se probaron cuatro configuraciones de Lighthouse (k ∈ {1536, 2048} variable y goleador ∈ {dilatado, norma} en L=3, p=4) frente a la línea de base de SDPA denso desde cero. Tres de los cuatro ensayos de Lighthouse igualan o superan la tasa de recuperación media de la línea de base densa de 0,72: k=2048 dilatado alcanza 0,76, k=1536 dilatado alcanza 0,73 y k=2048 norma coincide con la línea de base en 0,72. Sólo la norma k=1536 cae a 0,65. Un patrón emerge a lo largo de la cuadrícula: una k mayor es el eje dominante para el rendimiento de recuperación, y el anotador norma perjudica la recuperación más que la pérdida de entrenamiento al mismo k. La implicación práctica es que la configuración óptima depende de si la tarea posterior está impulsada por pérdidas o por recuperación.

Escalado de paralelismo contextual

Para secuencias de más de ~100 000 tokens, Lighthouse se ejecuta bajo paralelismo de contexto (CP). La agrupación piramidal, la puntuación y el top-K se ejecutan localmente en cada rango sin comunicación entre rangos, ya que la ventana del grupo más gruesa (por ejemplo, 64 tokens) es órdenes de magnitud más pequeña que el tamaño del fragmento. La subsecuencia recopilada es densa, por lo que participa en la atención de anillo estándar sin colectivos con conocimiento disperso, algo que los métodos basados ​​en índices dispersos no pueden hacer sin ingeniería específica para el diseño disperso. El paralelismo de contexto introduce aproximadamente un 10 % de sobrecarga de rendimiento por rango debido a la rotación del anillo, pero se conserva la relación de aceleración de Lighthouse frente a SDPA. El método se escala a un entrenamiento de 1 millón de tokens en 32 GPU Blackwell (4 nodos, grado de CP 8) sin cambios en el núcleo de atención interno.

Explicador visual de Marktechpost

Conclusiones clave

Lighthouse Attention de Nous Research agrupa Q, K y V simétricamente a través de una pirámide de varios niveles, a diferencia de NSA y HISA que solo agrupan K y V, cortando la llamada de atención de O(NS d) a O(S² d) y haciendo que el costoso paso almacene FlashAttention en una subsecuencia pequeña y densa. Es un método de solo entrenamiento: una breve reanudación de SDPA denso al final convierte el punto de control en un modelo normal de atención total que iguala o supera el denso desde cero con el mismo presupuesto de token (pérdida final 0,6980–0,7102 frente a 0,7237 de referencia, 16 000 pasos, ~50,3 mil millones de tokens). En un contexto de 512K en un solo B200, Lighthouse es 21 veces más rápido en el pase hacia adelante y 17,3 veces más rápido en avance + retroceso en comparación con cuDNN SDPA, lo que se traduce en una aceleración del reloj de pared de preentrenamiento de extremo a extremo de 1,40 × –1,69 ×. El paso de selección top-K es deliberadamente no diferenciable (no es un estimador directo, no es Gumbel softmax), por lo que las matrices de proyección aprenden a producir valores que son útiles cuando se seleccionan, no a engañar a un anotador que se puede aprender. Se escala a un entrenamiento de 1 millón de tokens en 32 GPU Blackwell (4 nodos, grado de CP 8) bajo paralelismo de contexto sin cambios en el núcleo de atención interno, porque la subsecuencia recopilada es densa y participa en la atención de anillo estándar.

Consulte el documento, el repositorio de GitHub y los detalles técnicos. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 150.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.

¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros