提出对注意力机制的改进,实现高效且更强大的序列建模。
Log-Linear Attention
- 用对数增长的隐藏状态替代固定大小状态,提升上下文表达能力。
- 计算复杂度为对数线性,比传统注意力更高效,优于线性注意力。
- 可兼容现有线性注意力模型,适合追求速度与性能平衡的研究者。
Transformer 中的注意力机制是精确且可扩展序列建模的重要基础,但其二次计算复杂度和线性内存开销仍是显著瓶颈。线性注意力和状态空间模型虽能实现线性时间、常数内存的序列建模,并通过跨序列长度的矩阵乘法并行化高效训练,但本质上仍为 RNN,依赖固定大小的隐藏状态建模上下文,存在根本局限。本文提出对数线性注意力(log-linear attention),在保持线性注意力效率的同时增强软注意力的表达能力。该方法将固定大小的隐藏状态替换为对数增长的隐藏状态集合。通过特定增长函数设计,对数线性注意力可实现类似矩阵乘法丰富的并行形式,计算成本在序列长度上为对数线性。该方法为通用框架,可应用于现有线性注意力变体。作为案例研究,我们构建了 Mamba-2 与 Gated DeltaNet 的对数线性变体,实验表明其性能优于对应的线性时间版本。
原文摘要 · Abstract (English)
The attention mechanism in Transformers is an important primitive for accurate and scalable sequence modeling. Its quadratic-compute and linear-memory complexity however remain significant bottlenecks. Linear attention and state-space models enable linear-time, constant-memory sequence modeling and can moreover be trained efficiently through matmul-rich parallelization across sequence length. However, at their core these models are still RNNs, and thus their use of a fixed-size hidden state to model the context is a fundamental limitation. This paper develops log-linear attention, an attention mechanism that balances linear attention's efficiency and the expressiveness of softmax attention. Log-linear attention replaces the fixed-size hidden state with a logarithmically growing set of hidden states. We show that with a particular growth function, log-linear attention admits a similarly matmul-rich parallel form whose compute cost is log-linear in sequence length. Log-linear attention is a general framework and can be applied on top of existing linear attention variants. As case studies, we instantiate log-linear variants of two recent architectures -- Mamba-2 and Gated DeltaNet -- and find they perform well compared to their linear-time variants.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。