Natürliche Sparse Attention: Hardware-ausgerichtet und nativ trainierbare Sparse AttentionNative Sparse Attention: Hardware-Aligned and Natively Trainable Sparse
Attention
Die Modellierung von langen Kontexten ist entscheidend für Sprachmodelle der nächsten Generation, jedoch stellt der hohe Rechenaufwand herkömmlicher Aufmerksamkeitsmechanismen bedeutende Rechenherausforderungen dar. Sparse Attention bietet eine vielversprechende Möglichkeit, die Effizienz zu verbessern, während die Modellfähigkeiten erhalten bleiben. Wir präsentieren NSA, einen nativ trainierbaren Sparse Attention Mechanismus, der algorithmische Innovationen mit hardwareorientierten Optimierungen integriert, um eine effiziente Modellierung langer Kontexte zu erreichen. NSA verwendet eine dynamische hierarchische Sparse-Strategie, die eine grobkörnige Token-Kompression mit einer feinkörnigen Token-Auswahl kombiniert, um sowohl das Bewusstsein für den globalen Kontext als auch die lokale Präzision zu erhalten. Unser Ansatz verbessert das Design von Sparse Attention durch zwei Schlüsselinnovationen: (1) Wir erzielen erhebliche Beschleunigungen durch eine algorithmische Gestaltung mit ausgewogener arithmetischer Intensität und Implementierungsoptimierungen für moderne Hardware. (2) Wir ermöglichen ein End-to-End-Training, das die Vorab-Berechnung reduziert, ohne die Modellleistung zu beeinträchtigen. Wie in Abbildung 1 gezeigt, zeigen Experimente, dass das mit NSA vorab trainierte Modell bei allgemeinen Benchmarks, langen Kontextaufgaben und instruktionsbasiertem Denken die Full Attention-Modelle beibehält oder übertrifft. Gleichzeitig erzielt NSA erhebliche Beschleunigungen gegenüber Full Attention bei Sequenzen mit einer Länge von 64k über Dekodierung, Vorwärts- und Rückwärtspropagation und bestätigt so seine Effizienz im gesamten Modelllebenszyklus.