Log-Linear Attention

2025年06月05日
  • 简介
    Transformer 中的注意力机制是实现精确且可扩展序列建模的重要基础组件。然而,其二次计算复杂度和线性内存复杂度仍然是显著的瓶颈。线性注意力机制和状态空间模型能够实现线性时间复杂度和常数内存复杂度的序列建模,并且可以通过沿序列长度的矩阵乘法密集型并行化高效训练。然而,这些模型本质上仍然是循环神经网络(RNN),因此使用固定大小的隐藏状态来建模上下文是其根本限制。本文提出了对数线性注意力(log-linear attention),这是一种在保持线性注意力效率的同时兼具 softmax 注意力表达能力的注意力机制。对数线性注意力用一个以对数速率增长的隐藏状态集合替代了固定大小的隐藏状态。我们证明,通过采用特定的增长函数,对数线性注意力可以转化为一种类似矩阵乘法密集型并行的形式,其计算成本与序列长度呈对数线性关系。对数线性注意力是一个通用框架,可以应用于现有的线性注意力变体之上。作为案例研究,我们在两个最近的架构——Mamba-2 和 Gated DeltaNet——上实例化了对数线性变体,并发现它们相较于线性时间复杂度的变体表现良好。
  • 图表
  • 解决问题
    该论文试图解决注意力机制在Transformer模型中的计算复杂度问题,特别是其二次计算复杂度和线性内存复杂度带来的瓶颈。虽然已有线性注意力和状态空间模型可以实现线性时间复杂度和常数内存复杂度,但它们本质上仍是RNN,受限于固定大小的隐藏状态。
  • 关键思路
    论文提出了一种名为log-linear attention的新机制,通过用对数增长的隐藏状态集替代固定大小的隐藏状态,在效率上接近线性注意力的同时保留了softmax注意力的表达能力。此外,log-linear attention可以通过特定的增长函数转化为矩阵乘法密集型并行形式,从而降低计算成本至对数线性复杂度。
  • 其它亮点
    1. 提出了log-linear attention框架,并验证了它在两种现有架构(Mamba-2和Gated DeltaNet)上的有效性。 2. 实验表明,log-linear attention变体相较于线性时间变体具有更好的性能。 3. 该方法支持高效并行化,适用于长序列建模。 4. 尽管未提及具体数据集或开源代码,但研究方向为未来工作提供了明确路径,例如优化增长函数或探索更多应用场景。
  • 相关研究
    近期相关研究包括: 1. "Linformer: Self-Attention with Linear Complexity" - 提出了一种线性复杂度的自注意力机制。 2. "Performer: Generalized Attention with RFF Kernels" - 使用随机特征核函数实现高效的注意力计算。 3. "Longformer: The Long-Document Transformer" - 针对长文档设计的局部与全局结合的注意力机制。 4. "SSM: State Space Models for Sequence Modeling" - 利用状态空间模型进行序列建模的研究。 这些研究均致力于提高注意力机制的效率,而本论文的独特之处在于引入对数增长的隐藏状态集以平衡效率与表达能力。
PDF
原文
点赞 收藏 评论 分享到Link

沙发等你来抢

去评论