Taipan融合Mamba与注意力机制,实现百万级上下文高效建模。
Taipan: Efficient and Expressive State Space Language Models with Selective Attention
- 用选择性注意力层筛选关键词元,只对重要信息启用注意力计算。
- 在100万词长上下文中保持高精度,推理内存恒定且计算效率高。
- 适合需要长距离依赖的复杂任务,如文档理解、代码生成等。
高效处理长文本仍是自然语言处理的核心挑战。尽管变压器模型占据主导地位,但其训练时计算复杂度为二次方,推理时内存线性增长,难以应对长序列。近期状态空间模型(SSMs)如Mamba虽实现恒定内存使用,但在需大量上下文检索的任务中表现不佳。我们提出新架构Taipan,融合Mamba-2与选择性注意力层(SALs)。SALs识别需长程交互的词元,剔除不重要特征,并通过注意力模块增强其表征。该方法兼顾Mamba的高效性与类变压器在内存密集型任务中的性能。通过约束注意力预算,Taipan将准确预测扩展至长达100万词的上下文,同时保持计算效率。实验表明,Taipan在多尺度、多任务上均表现优异,为高效长上下文建模提供可行方案。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。