Nous Research presenta Lighthouse Attention para acelerar el entrenamiento de LLM

Nous Research ha publicado Lighthouse Attention, un nuevo método para optimizar el entrenamiento de grandes modelos de lenguaje con contextos largos. El mecanismo actúa exclusivamente durante el preentrenamiento y se desactiva por completo cuando este termina, por lo que no afecta a la arquitectura ni al comportamiento del modelo terminado en el paso hacia adelante.

Cómo funciona Lighthouse Attention

Lighthouse Attention es un mecanismo de atención jerárquica y selectiva que envuelve la atención estándar de producto escalar escalado (scaled dot-product attention) durante el preentrenamiento del modelo. En la práctica, en cada paso por una capa de atención el modelo usa un mecanismo selectivo específico en lugar de la atención completa sobre todos los tokens del contexto.

La diferencia clave con enfoques anteriores (como NSA y HISA) es el pooling simétrico de todos los componentes de la atención. Los métodos previos reducían solo las claves y los valores (K y V) e ignoraban las consultas, mientras que Lighthouse aplica pooling a consultas, claves y valores (Q, K y V) a la vez mediante una pirámide de resoluciones de varios niveles. Así se logra una reducción del cómputo más equilibrada y eficiente en todos los niveles.

Técnicamente, esto reduce la complejidad computacional de la operación de atención de O(N·S·d) a O(S²·d), donde N es la longitud total del contexto, S es el tamaño de la subsecuencia compacta seleccionada y d es la dimensión de la capa oculta del modelo. Tras la selección, el FlashAttention estándar se ejecuta sobre la pequeña subsecuencia densa, lo que ahorra de forma considerable tanto recursos de cómputo como la memoria de GPU necesaria.

Resultados llamativos

Nous Research probó Lighthouse Attention en un modelo estilo Llama-3 de 530 millones de parámetros con un contexto de 98 000 tokens, que ya es un contexto bastante largo para pruebas. Los resultados mostraron una mejora significativa y constante del rendimiento del entrenamiento:

• Aceleración de 1.40–1.69 veces en el entrenamiento de extremo a extremo frente a la implementación base cuDNN SDPA en GPU
• Pérdida de entrenamiento (training loss) final comparable o inferior, lo que indica que no hay pérdida de calidad ni de precisión del modelo
• Plena compatibilidad con la infraestructura existente de FlashAttention y con frameworks estándar como PyTorch

Eso significa que las organizaciones podrían entrenar modelos grandes entre un 40 y un 70 por ciento más rápido sin sacrificar calidad ni precisión. En modelos grandes entrenados con enormes conjuntos de datos, esto se traduce en semanas de tiempo de cómputo ahorradas en costosos clústeres de GPU.

Aplicación práctica y escalabilidad

Las principales ventajas de Lighthouse Attention son su facilidad de adopción y que no influye en el comportamiento del modelo terminado. El mecanismo se usa solo durante el preentrenamiento y se desactiva automáticamente al concluir esa etapa crítica. Por ello, un modelo entrenado con Lighthouse es totalmente compatible con las aplicaciones, servicios y flujos de trabajo existentes, sin cambios en el código, la infraestructura ni el despliegue.

La aceleración resulta especialmente valiosa para organizaciones que entrenan modelos grandes con contextos de decenas y cientos de miles de tokens. Entre los usos típicos están el análisis de documentos e informes largos, la búsqueda de texto completo en grandes bases de conocimiento, la escritura y el análisis de código con contextos de 100K+ tokens, el procesamiento de conversaciones con un historial extenso y el trabajo con artículos científicos y patentes. Cada punto porcentual de cómputo ahorrado supone un ahorro concreto de electricidad y un recorte considerable en la factura de la computación en la nube.

Qué significa para la investigación y la industria

La optimización del entrenamiento de transformers sigue siendo un campo de investigación activo y fértil, pese a veinte años de inversión en los mecanismos centrales de la arquitectura. Lighthouse Attention demuestra que incluso en arquitecturas de atención muy estudiadas y pulidas queda margen para la innovación, las mejoras y las optimizaciones inesperadas.

Si la comunidad investigadora adopta métodos como este y se integran de forma masiva en frameworks abiertos populares como PyTorch, HuggingFace Transformers y otros, podrían reducir de forma notable la barrera de entrada para organizaciones, startups y grupos de investigación que quieran entrenar sus propios grandes modelos de lenguaje sin necesitar enormes recursos de cómputo ni presupuestos.