Taipan: Efficient and Expressive State Space Language Models with
  Selective Attention
Authors: name: Chien Van Nguyen; name: Huy Huu Nguyen; name: Thang M. Pham; name: Ruiyi Zhang; name: Hanieh Deilamsalehy; name: Puneet Mathur; name: Ryan A. Rossi; name: Trung Bui; name: Viet Dac Lai; name: Franck Dernoncourt; name: Thien Huu Nguyen
Published: 2024-10-24 09:25:37

Abstract: 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.