Taipan: Efficiënte en Expressieve Taalmodellen van de Toestandsruimte met Selectieve Aandacht

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

October 24, 2024

Auteurs: 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

Samenvatting

Efficiënte taalmodellering op basis van lange context blijft een aanzienlijke uitdaging in de verwerking van natuurlijke taal (NLP). Hoewel Transformers domineren bij taaltaken, hebben ze moeite met lange sequenties vanwege de kwadratische rekencomplexiteit tijdens training en lineair schalende geheugenkosten tijdens inferentie. Recente State Space Modellen (SSM's) zoals Mamba bieden alternatieven met constant geheugengebruik, maar presteren minder goed bij taken die uitgebreide in-context ophaling vereisen. We introduceren Taipan, een nieuw hybride architectuur die Mamba-2 combineert met Selectieve Aandachtslagen (SAL's). Deze SAL's identificeren tokens die langeafstandsinteracties vereisen, verwijderen minder belangrijke kenmerken en versterken vervolgens hun representaties met behulp van de aandachtsmodule. Deze benadering balanceert de efficiëntie van Mamba met Transformer-achtige prestaties bij geheugenintensieve taken. Door het aandachtsbudget te beperken, breidt Taipan nauwkeurige voorspellingen uit naar contextlengtes tot 1 miljoen tokens, terwijl de rekenkundige efficiëntie behouden blijft. Onze experimenten tonen de superieure prestaties van Taipan aan over verschillende schalen en taken, en bieden een veelbelovende oplossing voor efficiënte taalmodellering op basis van lange context.

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.