Taipan: Modelli di Linguaggio dello Spazio di Stati Efficienti ed Espressivi con Attenzione Selettiva

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

October 24, 2024

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

Abstract

Il modellamento efficiente del linguaggio a lungo contesto rimane una sfida significativa nel Processamento del Linguaggio Naturale (NLP). Sebbene i Transformers dominino i compiti linguistici, faticano con sequenze lunghe a causa della complessità computazionale quadratica durante l'addestramento e dei costi di memoria che scalano linearmente durante l'inferenza. I recenti Modelli dello Spazio di Stato (SSM) come Mamba offrono alternative con utilizzo costante della memoria, ma hanno prestazioni inferiori nei compiti che richiedono un'estesa ricerca nel contesto. Presentiamo Taipan, una nuova architettura ibrida che combina Mamba-2 con Strati di Attenzione Selettiva (SAL). Questi SAL identificano i token che richiedono interazioni a lungo raggio, rimuovono le caratteristiche meno importanti e quindi potenziano le loro rappresentazioni utilizzando il modulo di attenzione. Questo approccio bilancia l'efficienza di Mamba con le prestazioni simili a Transformer nei compiti che richiedono molta memoria. Limitando il budget di attenzione, Taipan estende le previsioni accurate a lunghezze di contesto fino a 1 milione di token preservando l'efficienza computazionale. I nostri esperimenti dimostrano le prestazioni superiori di Taipan su varie scale e compiti, offrendo una soluzione promettente per il modellamento efficiente del linguaggio a lungo contesto.

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.