用状态空间模型压缩分块令牌，实现高效可调的注意力机制。
Attamba: Attending To Multi-Token States
- 用状态空间模型压缩令牌分块，再对压缩表示做注意力。
- 在相似缓存开销下，困惑度降低24%；缓存减至1/4时仅损失5%性能。
- 支持变长分块，可平滑切换从二次到线性计算复杂度。
在序列生成中预测下一个标记时，传统Transformer需对所有历史标记进行注意力计算，导致计算量随序列长度呈平方增长。状态空间模型（SSM）将整个标记序列压缩为固定维度表示以提升效率，而其他架构则通过低秩投影或稀疏注意力实现亚二次复杂度。本文提出Attamba，采用状态空间模型压缩令牌分块，并对压缩后的键值表示施加注意力。实验发现，用SSM替代Transformer中的键值投影可提升模型质量并支持灵活分块，使具有相似键值缓存与注意力开销的模型困惑度降低24%，在键值缓存和注意力浮点运算减少约4倍的情况下，仅以5%困惑度为代价。Attamba能对变长分块序列执行注意力，实现从二次到线性的平滑过渡，提供可调节的效率优势。
原文摘要 · Abstract (English)
When predicting the next token in a sequence, vanilla transformers compute attention over all previous tokens, resulting in quadratic scaling of compute with sequence length. State-space models compress the entire sequence of tokens into a fixed-dimensional representation to improve efficiency, while other architectures achieve sub-quadratic complexity via low-rank projections or sparse attention patterns over the sequence. In this paper, we introduce Attamba, a novel architecture that uses state-space models to compress chunks of tokens and applies attention on these compressed key-value representations. We find that replacing key and value projections in a transformer with SSMs can improve model quality and enable flexible token chunking, resulting in 24% improved perplexity with transformer of similar KV-Cache and attention footprint, and ~4 times smaller KV-Cache and Attention FLOPs for 5% perplexity trade-off. Attamba can perform attention on chunked-sequences of variable length, enabling a smooth transition between quadratic and linear scaling, offering adaptable efficiency gains.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。