Attamba：注意力于多 token 状态
机器学习
2024-11-27 v1 计算与语言
摘要
在预测序列中下一个 token 时，vanilla transformer 对所有前置 token 计算注意力，导致计算量随序列长度呈二次增长。状态空间模型（state-space models）将整个 token 序列压缩为固定维度表示以提高效率，而其他架构则通过低秩投影或稀疏注意力模式实现亚二次复杂度。本文引入 Attamba 架构，使用状态空间模型压缩 token 块，并对这些压缩后的 key-value 表示应用注意力机制。我们发现，用 SSM 替代 transformer 中的 key 和 value 投影，可提升模型质量并实现灵活的 token 分块，实现与类似 KV-Cache 和注意力 footprint 的 transformer 相比的 24% 意义提升，以及在 5% 意义牺牲换取约 4 倍更小的 KV-Cache 和注意力 FLOPs。Attamba 可对变长块序列执行注意力，实现从二次到线性比例缩放的平滑过渡，提供可调的效率收益。
引用
@article{arxiv.2411.17685,
title = {Attamba: Attending To Multi-Token States},
author = {Yash Akhauri and Safeen Huda and Mohamed S. Abdelfattah},
journal= {arXiv preprint arXiv:2411.17685},
year = {2024}
}