Taipan: Effiziente und ausdrucksstarke Sprachmodelle für den Zustandsraum mit selektiver Aufmerksamkeit

Taipan: Efficient and Expressive State Space Language Models with Selective Attention

October 24, 2024

Autoren: Chien Van Nguyen, Huy Huu Nguyen, Thang M. Pham, Ruiyi Zhang, Hanieh Deilamsalehy, Puneet Mathur, Ryan A. Rossi, Trung Bui, Viet Dac Lai, Franck Dernoncourt, Thien Huu Nguyen

cs.AI

Zusammenfassung

Effizientes Sprachmodellieren über lange Kontexte bleibt eine bedeutende Herausforderung im Bereich der natürlichen Sprachverarbeitung (NLP). Obwohl Transformer dominierend in Sprachaufgaben sind, haben sie Schwierigkeiten mit langen Sequenzen aufgrund quadratischer Rechenkomplexität beim Training und linear skalierenden Speicherkosten während der Inferenz. Neueste Zustandsraummodelle (SSMs) wie Mamba bieten Alternativen mit konstanter Speichernutzung, jedoch erbringen sie schlechtere Leistungen bei Aufgaben, die umfangreiche Kontextabrufungen erfordern. Wir stellen Taipan vor, eine neuartige Hybridarchitektur, die Mamba-2 mit Selektiven Aufmerksamkeitsschichten (SALs) kombiniert. Diese SALs identifizieren Tokens, die langreichweitige Interaktionen erfordern, entfernen weniger wichtige Merkmale und erweitern dann ihre Darstellungen mithilfe des Aufmerksamkeitsmoduls. Dieser Ansatz balanciert die Effizienz von Mamba mit Transformer-ähnlicher Leistung bei speicherintensiven Aufgaben. Durch die Begrenzung des Aufmerksamkeitsbudgets ermöglicht Taipan präzise Vorhersagen für Kontextlängen von bis zu 1 Million Tokens bei gleichzeitiger Beibehaltung der Recheneffizienz. Unsere Experimente zeigen die überlegene Leistung von Taipan über verschiedene Maßstäbe und Aufgaben hinweg und bieten eine vielversprechende Lösung für effizientes Sprachmodellieren über lange Kontexte.

English

Efficient long-context language modeling remains a significant challenge in Natural Language Processing (NLP). While Transformers dominate language tasks, they struggle with long sequences due to quadratic computational complexity in training and linearly scaling memory costs during inference. Recent State Space Models (SSMs) such as Mamba offer alternatives with constant memory usage, but they underperform in tasks requiring extensive in-context retrieval. We introduce Taipan, a novel hybrid architecture that combines Mamba-2 with Selective Attention Layers (SALs). These SALs identify tokens requiring long-range interactions, remove less important features, and then augment their representations using the attention module. This approach balances Mamba's efficiency with Transformer-like performance in memory-intensive tasks. By constraining the attention budget, Taipan extends accurate predictions to context lengths of up to 1 million tokens while preserving computational efficiency. Our experiments demonstrate Taipan's superior performance across various scales and tasks, offering a promising solution for efficient long-context language modeling.