Attention Sparse Native : Une Attention Sparse Alignée sur le Matériel et Entraînable de Manière NativeNative Sparse Attention: Hardware-Aligned and Natively Trainable Sparse
Attention
La modélisation de contexte long est cruciale pour les modèles de langage de nouvelle génération, mais le coût computationnel élevé des mécanismes d'attention standard pose des défis significatifs. L'attention parcimonieuse offre une direction prometteuse pour améliorer l'efficacité tout en conservant les capacités du modèle. Nous présentons NSA, un mécanisme d'Attention Parcimonieuse Natif (Natively trainable Sparse Attention) qui intègre des innovations algorithmiques avec des optimisations alignées sur le matériel pour atteindre une modélisation de contexte long efficace. NSA utilise une stratégie parcimonieuse hiérarchique dynamique, combinant une compression grossière des tokens avec une sélection fine des tokens pour préserver à la fois la conscience du contexte global et la précision locale. Notre approche fait progresser la conception de l'attention parcimonieuse avec deux innovations clés : (1) Nous obtenons des accélérations substantielles grâce à une conception algorithmique équilibrée en intensité arithmétique, avec des optimisations d'implémentation pour le matériel moderne. (2) Nous permettons un entraînement de bout en bout, réduisant le calcul de pré-entraînement sans sacrifier les performances du modèle. Comme le montre la Figure 1, les expériences montrent que le modèle pré-entraîné avec NSA maintient ou dépasse les modèles à Attention Complète sur des benchmarks généraux, des tâches à contexte long et des raisonnements basés sur des instructions. Par ailleurs, NSA obtient des accélérations substantielles par rapport à l'Attention Complète sur des séquences de longueur 64k, que ce soit en décodage, propagation avant ou propagation arrière, validant ainsi son efficacité tout au long du cycle de vie du modèle.