Nous Research dévoile Lighthouse Attention pour accélérer l'entraînement des LLM

Nous Research a publié Lighthouse Attention, une nouvelle méthode d'optimisation de l'entraînement des grands modèles de langage sur de longs contextes. Le mécanisme agit exclusivement pendant le pré-entraînement et se désactive complètement une fois ce processus terminé : il n'affecte ni l'architecture ni le comportement du modèle fini lors de la passe avant.

Comment fonctionne Lighthouse Attention

Lighthouse Attention est un mécanisme d'attention hiérarchique et sélectif qui enveloppe l'attention standard par produit scalaire normalisé (scaled dot-product attention) pendant le pré-entraînement du modèle. Concrètement, à chaque passage dans une couche d'attention, le modèle utilise un mécanisme sélectif dédié plutôt que l'attention complète sur tous les tokens du contexte.

La différence essentielle avec les approches précédentes (comme NSA et HISA) tient au pooling symétrique de tous les composants de l'attention. Les méthodes antérieures ne réduisaient que les clés et les valeurs (K et V) et ignoraient les requêtes, alors que Lighthouse applique le pooling aux requêtes, aux clés et aux valeurs (Q, K et V) simultanément, via une pyramide de résolutions à plusieurs niveaux. La réduction du calcul est ainsi plus équilibrée et plus efficace à tous les niveaux.

Sur le plan technique, la complexité de calcul de l'opération d'attention passe de O(N·S·d) à O(S²·d), où N est la longueur totale du contexte, S la taille de la sous-séquence compacte sélectionnée et d la dimension de la couche cachée du modèle. Après la sélection, FlashAttention standard s'exécute sur la petite sous-séquence dense, ce qui économise nettement les ressources de calcul comme la mémoire GPU requise.

Des résultats impressionnants

Nous Research a testé Lighthouse Attention sur un modèle de style Llama-3 de 530 millions de paramètres, avec un contexte de 98 000 tokens, ce qui est déjà un contexte assez long pour des tests. Les résultats montrent une amélioration notable et constante des performances d'entraînement :

• Accélération de 1.40–1.69 fois de l'entraînement de bout en bout par rapport à l'implémentation de référence cuDNN SDPA sur GPU
• Perte d'entraînement (training loss) finale comparable ou inférieure, signe d'une absence de perte de qualité et de précision du modèle
• Compatibilité totale avec l'infrastructure FlashAttention existante et avec les frameworks standards comme PyTorch

Autrement dit, les organisations pourraient entraîner de grands modèles 40 à 70 pour cent plus vite sans sacrifier la qualité ni la précision. Pour les grands modèles entraînés sur d'énormes jeux de données, cela représente des semaines de temps de calcul économisées sur des clusters de GPU coûteux.

Application pratique et passage à l'échelle

Les principaux atouts de Lighthouse Attention sont sa simplicité d'adoption et l'absence d'effet sur le comportement du modèle fini. Le mécanisme n'est utilisé que pendant le pré-entraînement et se désactive automatiquement à la fin de cette étape critique. Un modèle entraîné avec Lighthouse reste donc pleinement compatible avec les applications, services et flux de travail existants, sans modification du code, de l'infrastructure ni du déploiement.

L'accélération est particulièrement précieuse pour les organisations qui entraînent de grands modèles sur des contextes de dizaines et de centaines de milliers de tokens. Parmi les usages typiques : l'analyse de longs documents et rapports, la recherche en texte intégral dans de vastes bases de connaissances, l'écriture et l'analyse de code sur des contextes de 100K+ tokens, le traitement de conversations à l'historique d'échanges profond, et le travail sur des articles scientifiques et des brevets. Chaque pour cent de calcul économisé se traduit par des économies concrètes d'électricité et par une baisse sensible de la facture de calcul dans le cloud.

Ce que cela change pour la recherche et l'industrie

L'optimisation de l'entraînement des transformers reste un domaine de recherche actif et fertile, malgré vingt ans d'investissements dans les mécanismes centraux de l'architecture. Lighthouse Attention montre que, même sur des architectures d'attention bien étudiées et polies, il reste de la place pour l'innovation, les améliorations et les optimisations inattendues.

Si la communauté de recherche adopte de telles méthodes et qu'elles sont intégrées à grande échelle dans des frameworks ouverts populaires comme PyTorch, HuggingFace Transformers et d'autres, cela pourrait abaisser nettement la barrière à l'entrée pour les organisations, les startups et les groupes de recherche qui veulent entraîner leurs propres grands modèles de langage sans disposer d'énormes ressources de calcul ni de budgets conséquents.