NLP——EfficientAttention-Survey(THU-2025)

注:本文包含 AI 辅助创作

Paper Summary

  • 整体总结:
    • 本文是一篇高效 Attention 的 Survey
    • 整体涉及到了 Linear Attention 和 Sparse Attention
  • 问题提出:
    • Transformer-based 架构已成为 LLM 的主流 backbone
    • 自注意力机制存在的二次时间和内存复杂度,是实现高效长上下文建模的根本障碍
  • 当前解法:
    • 为解决这一限制,近期研究引入了两大类高效注意力机制
    • 机制一:Linear Attention 方法:
      • 通过核近似、循环公式或快速权重动态来实现线性复杂度,从而以更低计算开销实现可扩展推理
    • 机制二:稀疏注意力 (Sparse Attention) 技术:
      • 基于固定模式、分块路由或聚类策略,将注意力计算限制在选定的 Token 子集上,在保持上下文覆盖的同时提升效率
  • 本综述主要工作如下:
    • 第一:对上述机制的发展进行了系统而全面的概述,结合了算法创新和硬件
    • 第二:分析了高效注意力在大规模预训练语言模型中的整合方式
      • 包括完全基于高效注意力构建的架构,以及结合局部和全局组件的混合设计

综述引言和一些讨论

  • Transformer-based 架构 (2017) 已成为现代 LLM 事实上的 backbone 选择
    • 标准的自注意力机制仍然是一个显著的计算瓶颈,其时间和内存复杂度相对于输入序列长度呈二次方增长
    • 这一限制对 LLM 在处理日益增长的长上下文时,同时保持高性能和高效率,提出了巨大挑战
  • 为解决此问题,出现了两个主要方向来降低 softmax 注意力 (softmax Attention) 的时间和空间复杂度
    • 第一种机制是Linear Attention (2020; 2020; 2021; 2021; 2023; 2024):
      • 通过将 softmax 注意力重新参数化或近似为线性操作来降低注意力复杂度
    • 第二种候选方案是稀疏注意力 (Sparse Attention) (2019; 2021; 2021; 2022; 2024):
      • 基于固定或动态的稀疏模式,将注意力计算限制在完整键空间的某个子集上
    • 虽然这两种方法都旨在提高效率,但它们在公式、设计选择和硬件影响方面存在显著差异
  • 本综述全面回顾了高效注意力 (Efficient Attention) 机制的最新进展,同时关注算法原理和系统级实现
    • 在此基础上,论文还研究了采用这些高效注意力的预训练 LLM
  • 论文将线性注意力方法分为三大范式
    • 范式一:核化线性注意力 (kernelized linear attention):
      • 通过特征空间内的内积来近似 softmax 核,借助随机特征映射 (2020; 2021) 或固定正映射 (2020) 实现线性复杂度
    • 范式二:带有遗忘机制 (forgetting mechanism) 的循环线性注意力 (recurrent linear attention):
      • 引入了位置感知的循环,通过数据无关 (2021) 或数据相关 (2023; 2023) 的衰减机制来控制过去信息随时间衰减的方式,从而实现对长序列的建模
    • 范式三:基于快速权重 (fast-weight) 和元学习 (meta-learning) 的公式:
      • 将线性注意力重新解释为在线优化的记忆更新过程,其中如 DeltaNet (2021; 2024) 和 TTT (2024) 等模型将快速学习动态直接整合到状态演化中
  • 作者还探讨了线性注意力在硬件友好的表示形式(包括并行、循环和分块形式)重点分析了它们在计算复杂度、内存占用以及与训练或推理工作流兼容性方面的权衡
  • 论文将稀疏注意力分为固定模式稀疏度 (fixed-pattern sparsity)、分块稀疏度 (block sparsity) 和基于聚类的稀疏度 (clustering-based sparsity)
    • 固定模式稀疏度采用静态的 Token-level 掩码,如滑动窗口、扩张位置或指定的全局 Token ,提供了简单性和硬件友好性 (2020; 2019; 2023; 2023)
    • 分块稀疏度在分块粒度上选择或路由注意力,可以通过启发式评分 (2021; 2022; 2025)、可训练的选通机制 (2024; 2024) 来实现,从而实现结构化的内存访问和高效的 GPU 利用率
    • 基于聚类的稀疏度使用基于内容或位置感知的分组方法(如 k-means 或 LSH)来组织键值对,以降低内存开销并促进基于语义感知的检索 (2023; 2020; 2024)
  • 最后,本综述还讨论了将稀疏模式扩展到编码器风格模型的双向稀疏设计
    • 这些方法在稀疏粒度、选择机制以及与 FlashAttention (2023) 等硬件原语的契合度上有所不同,共同代表了现代 Transformer 中实现高效长上下文建模的基础
  • 近期已有文章将高效注意力机制整合到工业级预训练语言模型中
    • 这包括纯高效架构:如线性注意力和状态空间模型,以及结合了局部和全局注意力模式的混合设计
      • 像 EAGLE (2024)、Falcon Mamba (2024) 和 MiniCPM4 (2025) 这样的模型展示了纯线性或稀疏方法在数十亿参数规模上的可扩展性,在提供强性能的同时实现了恒定时间推理
    • 同时,混合模型 (2020; 2022; 2024; 2025; 2025; 2025) 交错使用密集、稀疏和局部注意力,以平衡计算效率和上下文建模能力
      • 反映了现代 LLM 中朝着组合化、硬件感知注意力设计发展的趋势
  • 论文的目标是为理解注意力机制在算法和硬件双重约束下的演进,以及这些设计如何被整合到可扩展的 LLM 架构中,提供一个统一的框架
    • 通过将理论洞见与实际实现相结合,作者希望本综述能为致力于高效且可部署模型设计的研究者和从业者提供有价值的参考
  • 论文将讨论安排如下:
    • 第 2 节:介绍Linear Attention,涵盖其在不同模型世代中的演变、相关的设计原理以及对硬件实现的影响
    • 第 3 节:介绍稀疏注意力 (Sparse Attention),对稀疏模式进行分类,分析部署场景,并提供实用的系统级设计建议
    • 第 4 节:回顾整合了高效注意力机制的预训练语言模型 (Pretrained Language Models),包括统一高效架构和整合了局部、稀疏、密集注意力的混合模型
    • 第 5 节:对未来的发展方向进行展望 (Outlook),讨论算法和硬件对齐研究中的开放挑战和潜在进展

Linear Attention

  • 传统的线性注意力方法旨在以与序列长度成线性关系的方式近似基于 softmax 的注意力机制
    • 核心思想是用一种基于核函数(kernel)的注意力权重近似来替代计算代价高昂的 softmax 计算
  • 在标准的自注意力中,每个输出是值 \(V\) 的加权和,权重由查询-键相似度经过 softmax 得到:
    $$
    \text{Attn}(Q,K,V)=\text{softmax}(QK^{\top})V, \tag{1}
    $$
    • \(Q,K,V \in \mathbb{R}^{L \times d}\)
    • (\(L\) 是序列长度
    • \(d\) 是每个头的模型维度)
  • softmax 为查询 \(q_i\) 和键 \(k_j\) 产生权重 \(\propto \exp(q_i^{\top}k_j)\)
  • 核化线性注意力则寻找一个特征映射 \(\phi(\cdot)\),使得 softmax 核函数可以在诱导出的特征空间中通过一个简单的点积来近似:\(\exp(q^{\top}k) \approx \phi(q)^{\top} \phi(k)\)(2019)
    • 给定这样的 \(\phi\),可以将注意力重写为:
      $$
      O = \frac{\phi(Q)(\phi(K)^{\top}V)}{\phi(Q)(\phi(K)^{\top}\mathbf{I})} \tag{2}
      $$
  • 由于 \(\exp(\cdot)\) 的值域是非负的,\(\phi(\cdot)\) 通常被选择为产生非负输出,同时应用归一化除数来模拟 softmax 概率
  • 这种重新表述将复杂度从 \(O(L^2 d)\) 降低到 \(O(L d^2)\)(或者在适当的特征降维下甚至可以达到 \(O(L d)\)),因为昂贵的 \(L \times L\) 注意力矩阵从未显式形成
  • Linear Transformer (2020)
    • 用一个固定的正特征映射替换了 softmax 核函数
    • 在实践中,他们设置 \(\phi(x) = \operatorname{ELU}(x) + 1\)
    • \(\operatorname{ELU}(\cdot)\) 在整个定义域内可微,与朴素的 \(\operatorname{ReLU}(\cdot)\) 函数相比表现出更好的性能
  • Performer (2020)
    • 引入了 FAVOR+(一种能够无偏估计 softmax 核函数的随机特征方案)
    • 它对随机特征映射 \(\phi\) 进行采样,使得 \(E[\phi(Q)\phi(K)^{\top}] = \exp(QK^{\top})\)
    • 这产生了一个仅需 \(O(N)\) 操作即可证明是完整 softmax 注意力的无偏估计量
    • 特别地,Performer 使用正交互随机特征,降低了近似中的方差
  • 随机特征注意力 (Random Feature Attention,RFA) (2021)
    • 一种通过为 softmax 核函数使用随机傅里叶特征构建的线性注意力
    • 与 Performer 类似,RFA 利用随机映射和三角激活来近似 softmax
    • RFA 在随机投影之前进一步对查询和键进行归一化以减少方差
    • RFA 还有一个变体 RFA-Gate,它增加了一个可选的选通机制以引入近因偏差
  • cosFormer (2022)
    • 使用余弦函数来近似 softmax
    • 由于 \(\cos(a+b) = \cos a \cos b - \sin a \sin b\),cosFormer 将余弦重加权注意力 \(S_{ij} = Q_i’ K_j’ \cos(\frac{\pi}{2} \times \frac{i-j}{M})\) 分解为线性注意力的形式
  • HedgeDog (2024)
    • 利用了一个尖峰核函数 \(\phi(x) = \exp(Wx + b)\),因为他们观察到 Transformer 和线性 Transformer 之间的性能差距源于缺乏尖峰和单调性属性
    • HedgeDog 展示了更好的注意力熵和单调性

Linear Attention with Forgetting Mechanism

  • 最近的一系列工作通过循环神经网络或连续状态空间模型的视角来解释注意力
  • 传统的线性注意力通常是无位置感知的,其中循环顺序对输出没有影响,但现代的线性注意力表现得更像具有状态追踪和隐藏记忆的 RNN
  • 因此,这些模型明确地结合了循环、门控或状态动力学,以线性复杂度处理长序列
  • 衰减因子是引入遗忘机制的最重要因素
Data-Independent Decay
  • Retentive Networks(RetNet) (2023)
    • RetNet 引入了一种 Retention 机制,它使用固定的衰减系数,以一种循环风格的更新来取代注意力
    • 在 RetNet 层中,每个时间步 \(t\) 维护一个状态向量 \(s_t\),该向量以指数遗忘的方式聚合过去的输入
    • 循环可以写为:
      $$
      S_t = \gamma S_{t-1} + k_t^{\top} \nu_t \tag{3}
      $$
      • \(\gamma \in (0,1)\) 是一个学习的衰减因子(每个保持头)
      • \(k_t^{\top} \nu_t\) 是来自当前 Token 的新贡献(\(\nu_t\) 是 \(x_t\) 的值投影,\(k_t\) 是键投影)
    • 然后通过一个线性“查询”投影获得输出:\(o_t = q_t s_t\);展开方程 3 得到保持的显式公式:
      $$
      o_t = q_t S_t = \sum_{n=1}^{t} \gamma^{t-n} q_t k_t^{\top} \nu_t \tag{4}
      $$
    • 这表明 Token \(n\) 的贡献在时间步 \(t\) 时以因子 \(\gamma^{t-n}\) 呈指数衰减
    • Crucially,\(\gamma\) 是一个数据无关的衰减,是层的固定参数(在多头保持中通常每个头一个),而不是输入内容的函数
      • 这使得 RetNet 能够像 RNN 一样进行 \(O(1)\) 的内存更新,同时仍然允许通过等效的矩阵公式在训练期间进行并行计算(例如,可以证明方程 3 等价于一个“保持矩阵”形式,\(\text{Retention}(X) = (Q K^{\top} \odot D) V\),其中对于 \(t \geq n\),\(D_{t,n} = \gamma^{t-n}\) 实现了衰减和因果掩码)
    • RetNet 的保持机制与其他数据无关的循环模型有共同的主题
  • Eagle (2024)
    • 通过外积记忆改进了 RWKV 设计,这等效于线性注意力
    • 在 RWKV 系列中,衰减因子参数化为 \(\gamma = \exp(-\exp(w))\),其中 \(w\) 是一个数据无关的可学习因子
    • 在实践中,RetNet 和 Eagle 都使用固定衰减来遗忘旧信息,实现了线性推理扩展和具有竞争力的性能
    • 经验上,RetNet 每个头使用一个固定标量 \(\gamma\)(通常每层有多个具有不同 \(\gamma\) 值的保持头,形成一种多尺度衰减),而 Eagle 使用可学习的标量 \(w\) 来参数化衰减因子
  • Lightning Attention (2023; 2024)
    • Lightning Attention 也提出了一种线性注意力层 ,每个头增加了一个固定标量衰减,以实现长度不变的计算速度
    • 在 Lightning Attention 中,对于某个常数 \(\lambda\)(\(\lambda\) 由模型学习或设置),隐藏状态本质上是 \(s_t = \lambda s_{t-1} + k_t^{\top} \nu_t\),这与 RetNet 的 \(\gamma\) 精神相同,但针对硬件效率进行了优化
  • H3 (2022)
    • H3 将循环状态空间模型 (2021) 引入线性注意力,通过 SSM 为键值外积隐藏状态使用学习的、数据无关的指数衰减
    • 线性注意力通过分块计算实现了高效的训练,但 H3 需要为 SSM 计算显式状态扩展,从而限制了头维度,导致表达能力有限
  • In summary,数据无关的衰减方法维持一个随时间以预定速率衰减的持久状态,实现了 \(O(1)\) 循环和每步恒定内存
    • 它们牺牲了一些适应性 ,这促使了在更近期的模型中引入数据相关机制
Data-Dependent Decay
  • 虽然固定衰减提供了简单性和速度,但它们可能未能充分利用输入流中的信息
  • 门控或数据相关的方法使遗忘因子本身成为当前输入的学习函数
  • 这种循环更新的一般形式是:
    $$
    S_t = G_t S_{t-1} + k_t^{\top} \nu_t \tag{5}
    $$
    • \(S_{t-1}\) 是前一个状态
    • \(G_t\) 是由 Token \(x_t\) 确定的门控张量
      • 如果 \(G_t\) 在某个分量上接近 0,则该分量中的过去状态在时间 \(t\) 被大量遗忘;
      • 如果 \(G_t \approx 1\),则过去状态被保留
    • 与 RetNet 中的常数 \(\gamma\) 不同,这里的 \(G_t\) 通过 \(x_t\) 随 \(t\) 变化
  • 在大语言模型设计中,这种策略的两个显著例子是 Mamba (2023) 和门控线性注意力 (Gated Linear Attention, GLA) (2023)
  • Mamba
    • 一种循环状态空间模型,它赋予状态衰减率以输入依赖性
    • 在每个 Mamba 层中,基本的状态演化类似于 S4 (2021),但状态矩阵实际上变得动态
    • \(G_t\) 是一个范围在 0 到 1 之间的分组向量,作为动态遗忘门
      • 这弥合了注意力和纯 SSM 之间的差距
    • 实证结果表明,Mamba2 在语言建模任务上可以超越相似甚至更大规模的 Transformer,凸显了数据相关衰减在长序列建模中的能力
  • GLA
    • 直接在线性注意力中引入了门控机制,将一个门控函数嵌入到线性化的注意力层中以提高其表达能力
    • GLA 通过一个可学习的 Element-wise 遗忘门 \(G_t\) 来修改保持循环
  • 除此之外,其他几个模型同样赋予了其循环以内容相关的门控
  • xLSTM (2024)
    • 用线性门控信号(带归一化)的指数变换取代了标准的 sigmoid 遗忘门,对其单元状态产生了平滑的、输入条件化的衰减
  • GateLoop (2023)
    • 在保持机制上应用了头级门控,实现了一个简单但有效的数据相关衰减,同时保持了高效的硬件实现
  • HGRN (2023)
    • 在线性 RNN 中引入了门控循环
    • HGRN2 (2024) 进一步在 HGRN 框架中增加了状态扩展
    • 状态扩展等效于线性注意力中的键值外积
  • Finch (2024)
    • 在 Eagle 上使用了数据相关的门控
    • 由于 Eagle 与其他正交修改的保持机制相似,Finch 也显示出与上述模型的深厚联系
  • In summary,数据相关衰减模型通过基于内容的门控来增强线性注意力或 RNN 风格的架构,这些门控控制着信息的流动
    • 论文中的结果表明,这些模型在语言任务上通常能够匹配或超越 Transformer 的性能 ,同时能够扩展到非常长的输入

Linear Attention as In-Context Learners

  • 除了线性注意力机制带来的效率提升外,一项重大进展在于它们应用于增强上下文学习能力
    • 这指的是模型能够从给定的 Prompt 中快速适应或学习,而无需对其预训练权重进行显式的梯度更新
    • 理解:这里是说注意力机制本身是上下文学习期,即增加一些 Prompt Token,就可以借助修改注意力内容来实现对模型的快速学习
  • 大型 Transformer 模型通过将 Prompt 解释为一种训练数据的形式,固有地表现出上下文学习
    • 但最近的创新已将快速学习规则直接整合到注意力机制中,有效地将序列处理视为一个在线训练过程
    • 问题:这里的方法具体是什么?
  • FWP (2021) 建立了现有线性注意力机制与快速权重编程器之间的形式等价性
    • 在 FWP 范式中,一个慢速神经网络学习去编程另一个神经网络的“快速权重”,通常通过自发明的键和值模式的外积加法来实现
  • 表 1:
    • 不同线性注意力变体的更新规则
    • 每个模型都是关于矩阵记忆 \(S_t\) 的循环
Learning Objective
  • 从元学习的角度来看,这些模型定义了一个在推理过程中优化的隐式学习目标
  • 用 \(q_t, k_t, \nu_t\) 表示时间步 t 的查询、键和值,上下文记忆 \(S_t\) 通过以下目标进行优化:
    $$
    \mathcal{L}_t(S) = \frac{1}{2} || f_S(k_t) - \nu_t ||^2 \tag{6}
    $$
  • DeltaNet
    • DeltaNet 融合了经典的 delta 规则 (2021),其中 \(f_S(k_t) = S k_t\)
    • DeltaNet 更新规则为
      $$ S_t = S_{t-1} + \eta_t (\nu_t - S_{t-1} k_t) k_t^{\top} $$
      • 可以通过最小化当前记忆检索 \(S_{t-1} k_t\) 与新值 \(\nu_t\) 之间的误差推导出来
      • 这标志着朝着在线学习键值映射、基于即时上下文改进记忆迈出了一步
  • TTT (2024)
    • TTT 用不同的建模架构概括了元学习目标:
      $$
      f_S(k_t) = \begin{cases}
      \text{LM}(S k_t) + k_t, & \text{TTT-Linear} \\
      \text{LM}(\text{MLP}_S(k_t)) + k_t, & \text{TTT-MLP}
      \end{cases} \tag{7}
      $$
    • 上下文网络 \(f_S\) 增强了上下文元学习的能力
    • 但由于 \(f_S\) 的梯度比简单的线性投影复杂得多,在线更新不能写成一个简单的规则
  • 批量更新 (Batch Update)
    • 批量更新试图解决当 \(f_S\) 作为神经网络工作时训练并行性的困难
    • 通常,上下文记忆是以批大小为 1 进行元学习的,这对于一般的 TTT 模型来说不可行
    • 相反,类似于分块并行,TTT 将整个块视为一个批次
    • 在批次内没有状态更新(即 \(S\) 保持不变)
    • 处理完批次后,\(S\) 使用来自批次中所有样本的聚合梯度或更新信号进行一次更新
    • 这种策略在保持并行效率的同时,适应了更复杂架构的训练要求
  • Momentum Titans (2025)
    • Titans 中引入了优化中常用的动量,以加强记忆更新机制的能力:
      $$
      \mathcal{M}_t = (1 - \alpha_t) \mathcal{M}_{t-1} + S_t \\
      o_t = q_t \mathcal{M}_t \tag{8}
      $$
    • 动量项允许记忆通过对状态 \(S\) 进行指数移动平均来逐渐积累信息
    • 这可以看作是元学习的一种形式,其中更新规则本身学会了在长序列上更加稳定和鲁棒
  • 权重衰减 (Weight Decay)
    • 权重衰减是训练中的另一种正则化技术,对应于线性注意力模型中的遗忘机制
    • Gated DeltaNet (2024) 和 Titans 在其记忆更新中使用了权重衰减,作为学习的遗忘门来限制非常旧或嘈杂数据的影响
    • 它对应于在 RetNet (2023) 和 Mamba (2023) 等架构中发现的选择性状态保持机制,其中衰减机制被证明对语言建模性能至关重要:
      $$
      S_n = \gamma_n S_{t-1} + \eta_t (v_t - S_{t-1} k_t) k_t^{\top} \tag{9}
      $$
  • In summary,线性注意力机制中的这些进步通过明确地将元学习原则整合到其架构中,正在推动上下文学习的边界
    • 通过快速权重更新、复杂的内存管理技术和在线学习规则,这些模型正朝着一种范式发展,在这种范式中,训练和推理之间的区别变得越来越模糊,从而产生了能够直接从上下文中学习和利用知识的更高效、适应性更强的大语言模型

Discussion on Other Designs

Element-wise Linear Attention
  • 无需注意力的 Transformer (Attention-Free Transformer) (2021)
    • 利用一个简单的权重 \(\exp(K_{\nu’} + w_{t,\nu’})\) 来代替 \(\exp(QK^{\top})\):
      $$
      O_t = \sigma_q(Q_t) \odot \frac{\sum_{\nu’=1}^{t} \exp(K_{\nu’} + w_{t,\nu’}) \odot V_{\nu’} }{\sum_{\nu’=1}^{t} \exp(K_{\nu’} + w_{t,\nu’})} \tag{10}
      $$
    • 其中 \(w_{t,\nu’}\) 是学习的成对位置偏置
    • 在 AFT 的变体中,AFT-Simple 移除了 \(w_{t,t’}\),实现了线性化的推理模式
    • 由于 \(K\) 和 \(V\) 的乘积是 Element-wise 的,循环状态大小是 \(\mathbb{R}^d\) 而不是外积状态 \(\mathbb{R}^{d \times d}\)
  • RWKV (2023)
    • 在 AFT-Simple 上利用了衰减机制。具体来说,RWKV 通过指数衰减 \(w_{t,i} = -(t-i)w\) 改进了 AFT 的位置偏置
    • 指数形式在引入位置偏置的同时保留了循环属性
  • Element-wise 线性注意力带来了强大的推理优势,但它受到状态大小瓶颈的影响,性能低于基于矩阵的状态大小
    • Besides,尽管 Element-wise 内存比外积内存快得多,但由于其他组件在拥有外积内存的情况下占据了超过 95% 的延迟 (2023),端到端的优势仍然有限
Multi-Pass Linear Attention
  • 带有限内存控制的注意力 (Attention with Bounded-memory Control)
    • 将线性注意力视为一个有界内存模块:
      $$
      \begin{split}
      \tilde{K}_n &= \sum_{i=1}^{n} K_i \otimes \phi_i, \quad \tilde{V}_n = \sum_{i=1}^{n} V_i \otimes \phi_i \\
      O_n &= \text{softmax}(Q_n \tilde{K}_n^{\top}) \tilde{V}_n
      \end{split} \tag{11}
      $$
    • 其中 \(\tilde{K}_n, \tilde{V}_n\) 是在线更新的、大小受限的键和值。在实现中,ABC 可以简化为两遍线性注意力
  • 门控 Slot 注意力 (Gated Slot Attention) (2024)
    • 进一步将 GLA 引入到 ABC 框架 (2021) 中
    • 由于 \(\tilde{K}_n, \tilde{V}_n\) 作为一个隐式的线性注意力工作,GSA 将更新改进为门控形式:
      $$
      \tilde{K}_n = \text{Diag}(\alpha_n) \tilde{K}_{n-1} + (1 - \alpha_n) \otimes K_n, \quad \tilde{V}_n = \text{Diag}(\alpha_n) \tilde{V}_{n-1} + (1 - \alpha_n) \otimes V_n \tag{12}
      $$
  • Multi-Pass 是增强线性注意力表达能力的一种有效方式,但它也带来了额外的计算开销,这使得架构设计需要在训练效率和性能之间进行权衡
Bidirectional Linear Attention
  • 双向注意力在编码器风格的架构(如 BERT (2019))中扮演着重要角色
  • 单向和双向注意力在线性表述中的关键区别在于推理瓶颈和计算模式
  • Decoder-only 模型通常表现出 \(O(N^2)\) 的复杂度,且 Decoder-only 模型中的每个 Token 都可以访问全局信息
    • 因此双向线性注意力通常维护一个恒定长度的全局 Token 池以降低复杂度,同时保留 softmax 函数的使用
  • 例如
    • Linformer (2020) 通过一个额外的矩阵投影将键和值的数量减少到一个恒定长度
    • Luna (2021) 通过跨模型层编码全局 Token 池进一步扩展了 Linformer 设计
  • 虽然双向线性注意力对 Decoder-only 架构有效,但这些设计在应用于因果设置时面临重大挑战,因为基于全局池的方法往往计算成本高昂
    • 因此,此类架构不太适合大语言模型

Hardware Implementation

  • 并行表示 (The Parallel Representation)
    • 可以将带门控衰减的因果线性注意力定义为:
      $$
      \begin{split}
      Q &= \phi(XW_Q), \quad K = \phi(XW_K), \quad V = XW_V, \quad \gamma = f_\gamma(X) \\
      D_{nm} &= \begin{cases}
      \prod_{i=m+1}^{n} \gamma_i, & n \geq m \\
      0, & n < m
      \end{cases}, \quad O(X) = \text{LN}((Q K^{\top} \odot D) V)
      \end{split} \tag{13}
      $$
    • 其中
      • \(W_Q, W_K, W_V \in \mathbb{R}^{d \times d}\)
      • \(f_\gamma\) 控制衰减的锐度
      • 矩阵 \(D \in \mathbb{R}^{N \times N}\) 编码了具有衰减模式的因果掩码,确保信息的单向流动
    • 当衰减是数据无关时,\(f_\gamma(\cdot) = \text{const} \in (0,1]\)
      • 注意,线性注意力后的分组归一化 (GroupNorm) (2018) 已经是一个强制性的组件 (2023),方程 2 中核化线性注意力的显式除数变得不必要
    • 并行表示简单易懂,但有两个缺点
      • 第一:并行形式仍然保持了与 softmax 注意力相同的 \(O(N^2)\) 复杂度
      • 第二:当表示第 2.3 节中的 ICL 风格线性注意力时,其复杂性会增加
  • 循环表示 (The Recurrent Representation)
    • 上述并行公式可以等效地表达为逐步解码的循环形式,如图 2(b) 所示,在每个时间步 \(n\),输出计算为:
      $$
      S_n = f_{\text{update} }(S_{n-1}, K_n, V_n), \quad O_n = Q_n S_n \\
      f_{\text{update} }(S_{n-1}, K_n, V_n) = \begin{cases}
      \gamma_n S_{n-1} + K_n^{\top} V_n, & \text{Linear Attention with Decay} \\
      \gamma_n S_{n-1} + \eta_n (V_n - S_{n-1} K_n) K_n^{\top}, & \text{ICL-style Linear Attention}
      \end{cases} \tag{14}
      $$
    • 这种循环表示通过维护单个状态向量 \(S_n\),实现了具有恒定内存的高效自回归生成
    • 虽然循环表示将计算复杂度从 \(O(N^2)\) 降低到 \(O(N)\),但它在训练期间会产生大量的内存开销
      • 因为 \(S_n\) 涉及存储 \(K_n\) 和 \(V_n\) 的外积,这对于长序列来说是极其昂贵的
      • 因此,循环形式通常仅限于解码阶段
  • 分块循环表示 (The Chunkwise Recurrent Representation)
    • 分块表示结合了线性复杂度和硬件友好的并行性的优势 (2022; 2023)
    • 如图 2(a) 所示,以衰减风格线性注意力为例,给定一个块大小 \(B\),令 \(x_{[i]}\) 表示第 \(i\) 块。定义块内的累积衰减为:
      $$
      \beta_{(i-1)B+j} = \prod_{k=(i-1)B+1}^{(i-1)B+j} \gamma_k, \quad D_{[i]}(j,k) = \begin{cases}
      \frac{\beta_{(i-1)B+k} }{\beta_{(i-1)B+j} }, & j \leq k \\
      0, & \text{Other wise}
      \end{cases} \tag{15}
      $$
    • 块级记忆状态 \(R_i\) 计算为:
      $$
      R_i = K_{[i]}^{\top} (V_{[i]} \odot \frac{\beta_{iB} }{\beta_{[i]} }) + \beta_{iB} R_{i-1} \tag{16}
      $$
    • 第 \(i\) 块的输出为:
      $$
      O_{[i]} = (Q_{[i]} K_{[i]}^{\top} \odot D_{[i]}) V_{[i]} + (Q_{[i]} R_{i-1}) \odot \beta_{[i]} \tag{17}
      $$
    • 这种表述提供了循环和并行的统一视图:
      • 第一项捕获块内依赖关系
      • 第二项通过单个矩阵-向量乘积传播块间记忆
      • 由于其效率和可并行性,分块表示通常在训练和预填充阶段被采用
  • 对于 ICL 风格的线性注意力,已经使用 Householder 变换 (2024) 开发了硬件友好的分块表示
    • 但对于更复杂的变体,如 TTT 和 Titans,构建显式的分块形式仍然具有挑战性
    • 相反,这些架构通常依赖于大的批次大小进行记忆更新,通过固定超参数有效地模拟分块计算
  • 内核级优化对于实现高性能至关重要
    • 广泛采用的 FLA (2024) 为许多常见的线性注意力模块提供了基于 Triton 的实现
    • 或者,开发者还提供了用 CUDA 或 TileLang (2024) 编写的自定义实现,可以用于进一步加速
  • 图 2:线性注意力的双重形式

Sparse Attention

  • 稀疏注意力(Sparse Attention)方法利用注意力计算中固有的稀疏特性,通过以下公式近似完整注意力计算:
    $$
    \text{Attn}(Q,K,V) = \text{softmax}(QK_{[\mathcal{S}]}^T)V_{[\mathcal{S}]} \tag{18}
    $$
    • 其中 \(\mathcal{S}(t)\) 是查询向量 \(Q(t)\) 所关注索引的子集
    • 不同的方法基于选择准确性和硬件效率的考量,设计了不同的选择标准 \(\mathcal{S}(t)\)
    • 目标是在预填充(prefilling)阶段实现亚线性或线性复杂度,或在解码(decoding)阶段控制固定的计算预算

Fixed-pattern Sparse Attention (固定模式稀疏注意力)

  • 一些研究工作利用 Token-level 别稀疏性的结构化模式,为注意力计算构建固定的稀疏掩码
  • Local Window Attention (局部窗口注意力)
    • 局部窗口注意力将每个查询限制在仅与固定滑动窗口 \(w\) 内的相邻 Token 交互,从而在保留局部上下文的同时降低内存和计算需求
    • Sparse Transformer (2019) 首先应用局部窗口(行)注意力,其中 \(w\) 接近 \(\sqrt{N}\),然后通过额外的列注意力来总结先前位置并在全局传播信息
    • GPT-3 (2020) 也采用了与 Sparse Transformer 类似的稀疏注意力模式
    • StreamingLLM (2023) 发现大量的注意力得分被分配给了输入序列的初始 Token ,他们称之为“注意力汇集点(attention sink)”。他们提出了一种简单的固定模式注意力,只保留汇集点 Token 和滑动窗口内的 Token 。例如,给定一个长度为 \(n\) 的输入序列,StreamingLLM 中查询 Token \(q_t\) 的选定 Token 子集 \(\mathcal{S}(t)\) 被定义为:
      $$
      \mathcal{S}(t) = \left\{ j \mid 0 \leq j \leq s \ \lor \ t-w \leq j \leq t \right\}, \ \forall t \in [1, n] \tag{19}
      $$
      • 其中 \(s\) 是汇集点 Token 的大小,\(w\) 是滑动窗口的大小
      • 为了获得更好的硬件效率,采用块粒度的 StreamingLLM (2024) 以块方式保留汇集点 Token 和局部 Token ,从而实现高效的内存加载和计算
  • Dilated Attention (膨胀注意力)
    • LongNet (2023)
      • 引入了膨胀注意力作为长上下文训练和推理的固定稀疏模式。膨胀注意力随着距离的增长呈指数级扩展注意力范围,从而将注意力的复杂度从 \(O(n^2)\) 降低到 \(O(n)\)。具体来说,在沿序列维度将输入分割成长度为 \(w\) 的片段后,从每个片段中以间隔 \(r\) 选择膨胀稀疏索引。第 \(i\) 个片段的选择索引为:
        $$
        \hat{I}_i = \left[iw, iw+r, iw+2r, …, (i+1)w-1 \right] \tag{20}
        $$
      • 稀疏化的片段 \(Q_{\hat{I}_i}, K_{\hat{I}_i}, V_{\hat{I}_i},\ i \in \{0, 1, …, \frac{n}{w}\}\) 被并行输入到注意力计算中,得到注意力输出 \(O\)。结合不同片段大小和膨胀率 \(\{r_i, w_i\}^{k}\) 的注意力输出,最终注意力计算如下:
        $$
        O = \sum_{i=1}^{K} \alpha_i O|_{r_i, w_i}, \quad \alpha_i = \frac{s_i}{\sum_{j} s_j} \tag{21}
        $$
      • 其中 \(s_i\) 表示 \(O|_{r_i, w_i}\) 的注意力 softmax 的分母
  • LogSparse (2019)
    • 采用指数稀疏注意力方案,其中每个位置仅关注 \(\log N\) 个 Token ,这可以看作是指数膨胀注意力的一个实例

Block Sparse Attention (块稀疏注意力)

  • 给定一个长度为 \(n\)、块大小为 \(b\) 的输入序列,可以将 \(Q, K, V \in \mathcal{R}^{n \times d}\) 各自划分为 \(\frac{n}{b}\) 个块,每个块大小为 \(b \times d\)
  • 目标是近似一个块级掩码 \(M \in \{0, 1\}^{n/b \times n/b}\),用于选择关键块进行计算,如图 3 所示
    $$
    \text{Attn}(Q, K, V)_i = \sum_{j=1}^{n/b} M_{ij} \cdot \text{softmax}(Q_i K_j^T) V_j \tag{22}
    $$
  • 块级选择对于在现代 GPU 上实现高效计算至关重要
  • 图 3: 块稀疏注意力:长序列被分成若干块,每个 Token 仅关注其局部窗口和 top-k 相关块
Block-Sparse Attention for Prefill (用于预填充的块稀疏注意力)
  • 使用块稀疏注意力进行预填充的方法,近似选择覆盖大部分注意力得分且具有高召回率的 Top-K 块,从而将注意力计算复杂度从 \(O(n^2)\) 降低到 \(O(K)\)
    $$
    S = \text{softmax}(QK^T - c(1 - M)) \\
    \min \ |S(M) - S_{\text{dense} }|
    $$
    • 其中 \(M\) 是如上定义的块级稀疏掩码,\(c\) 是一个大常数(例如 1e5),确保重要性较低的注意力权重在 softmax 计算后趋近于零
    • 块稀疏注意力的目标是在最小开销下实现更大的加速,同时尽可能保留更多的注意力权重
  • MInference (2024)
    • 观察到注意力权重中存在三种模式:
      • 流式(A 型)模式(Streaming (A Shape) Pattern)
      • 垂直斜线模式(Vertical-Slash Pattern)
      • 块稀疏模式(Block-Sparse Pattern)
    • 它离线确定每个注意力头的最佳模式,并在推理过程中基于指定的模式动态构建稀疏索引
  • FlexPrefill (2025)
    & 提出了一种上下文感知的稀疏注意力机制,能够实时动态调整注意力模式和计算预算
  • XAttention (2025)
    • 提出了一个块稀疏注意力框架,利用反对角线(antidiagonal)评分来预测注意力块的重要性,从而能够高效识别并剪除非必要块,实现高稀疏性和显著的计算增益
  • SpargeAttn (2025)
    • 同样在预填充阶段采用块级稀疏注意力,通过一个双阶段在线过滤过程完成:
      • 第一阶段快速预测注意力图以跳过某些矩阵乘法
      • 第二阶段应用 softmax 感知的过滤器以进一步消除不必要的计算
Block-Sparse Attention for Decode (用于解码的块稀疏注意力)
  • 用于解码的块稀疏注意力方法动态选择包含每个解码步骤最关键的 Token 的子集 \(S\) 的 \(K, V\) 向量,从而减少内存加载并提高效率
  • Quest (2024) 通过计算注意力权重的上界来近似每个块的关键性。对于块 \(K_i\),论文通过以下公式维护 Element-wise 的 Min 和 Max Key \(m_i\) 和 \(M_i\):
    $$
    m_{i,d} = \min(K_{i,d}), \quad M_{i,d} = \max(K_{i,d})
    $$
    • 其中 \(\min(\cdot)\) 和 \(\max(\cdot)\) 在每个维度 \(d\) 上 Element-wise 应用
  • 给定查询 \(q\),块 \(i\) 的近似注意力得分由下式给出:
    $$
    \text{score}_i = \sum_{j=1}^{d} \max(q_j \times M_{i,j}, q_j \times m_{i,j})
    $$
  • 然后选择得分最高的 Top-K 块作为注意力计算的稀疏子集 \(S\):
    $$
    S = \text{argtopk}(\text{score}, k)
    $$
  • DoubleSparsity (2024) 通过降低计算 \(QK^T\) 乘积的矩阵乘法维度来高效近似关键 Token
    • 它首先离线计算 \(QK^T\) 中的离群通道,记为 \(C\)
    • 然后选择具有最高近似注意力得分 \(\hat{s}\) 的 Top-K Token 作为稀疏子集 \(S\):
      $$
      Q_{\text{label} } = Q_{\{C\} }, \quad \hat{s} = Q_{\text{label} } K_{\text{label} }^T, \quad S = \text{argtopk}(\hat{s}, k)
      $$
  • ReSA (2025)
    • 结合了无需训练的块稀疏估计和 GQA 共享,有助于提高效率
    • ReSA 还提出了一个修正阶段来控制 KV 缓存累积误差
    • ReSA 在长序列生成任务上显示出优势
Routing-based Block-Sparse Attention (基于路由的块稀疏注意力)
  • 基于路由的块稀疏注意力通过可训练的 MLP 层学习每个 Token 块的重要性,该层在推理期间充当门控网络以选择关键块
  • Learnable Sparsity on Pretrained Models (预训练模型上的可学习稀疏性)
  • SeerAttention (2024)
    • 通过自蒸馏(self-distillation)的方式在预训练的 LLM 上训练门控网络
    • 为了获得每个块的重要性分数,它首先沿序列维度对 \(Q\) 和 \(K\) 进行池化,记为 \(P_q\) 和 \(P_k\)。下采样后的 \(Q, K\) 然后通过一个可学习的线性层 \(W_q\) 和 \(W_k\)
    • 投影后的 \(W_q P_q(Q)\) 和 \(W_k P_k(K)\) 的矩阵相乘结果通过 softmax 运算符作为门控过程:
      $$
      \text{score} = \text{softmax}((W_q P_q(Q)) \cdot (W_k P_k(K)))
      $$
    • 可学习的线性层通过自蒸馏的方式训练,以与原始 LLM 的 2D 最大池化结果对齐。蒸馏损失计算如下:
      $$
      gt = \text{MaxPool2D}(\text{softmax}(QK^T)), \quad loss = D_{KL}(gt \ || \ \text{score})
      $$
    • 在推理过程中,门控分数通过 Top-K 或阈值化来预测块级稀疏性,用于稀疏计算和效率提升
  • Training-aware Sparse Attention (训练感知的稀疏注意力)
    • Landmark (2023)
      • 提出使用特殊的“地标”(landmark) Token 来表示每个块,并训练注意力机制通过这些地标 Token 直接检索 Top-K 块
      • 然而,它没有在大型预训练模型上进行实验
    • MoBA (2025)
      • 将可训练的稀疏注意力集成到预训练阶段
      • 它提出了 Mixture of Block Attention,应用来自 MoE(专家混合)的 Top-K 机制作为门控机制,为每个查询 Token 决定关键块
      • 每个块的重要性分数通过查询 Token \(q\) 和块 \(K_i\) 沿 Token 维度的平均池化结果的内积计算:
        $$
        s_i = \langle q, P_{\text{mean} }(K_i) \rangle
        $$
      • 然后选择具有最高 \(s\) 分数的 Top-K 块用于计算 \(q\) 的注意力
      • Notably,MoBA 使用的 Top-K 块选择是不可微分的
        • 因此,在预训练阶段,稀疏模式仍然以无需训练的模式进行估计,从而实现高效的推理和加速的训练
    • NSA (2025)
      • 引入了一种训练感知的多粒度稀疏注意力机制,包含三个分支 \(C \in \{\text{cmp}, \text{slc}, \text{win}\}\),分别对应压缩(compression)选择(selection)滑动窗口(sliding window)策略
      • NSA 利用一个可微分的压缩分支来学习块选择分数。结合三个分支,NSA 的注意力输出由下式给出:
        $$
        o = \sum_{c \in \mathcal{C} } g^c \cdot \text{Attn}(q, K^c, V^c), \quad g_c \in [0, 1]
        $$
      • 对于压缩分支 \(c = \text{cmp}\),块 \(i\) 的键 \(K_i \in \mathcal{R}^{d_i \times b}\) 通过一个可学习的 MLP 层 \(\varphi\) 被压缩为单个键 \(K_i^{\text{cmp} } \in \mathcal{R}^{d_i \times 1}\)
      • 对于选择分支 \(c = \text{slc}\),基于块重要性分数 \(p\) 选择 Top-K 块,该分数可以直接从压缩分支获得
    • InfLLM-v2 (2025)
      • 采用了与 MoBA 类似的训练感知 Top-K 块稀疏注意力机制
      • 为了提高 Top-K 块选择的准确性,它将块划分为具有重叠的小粒度内核,并在每个块内对内核重要性分数进行聚合
System-level Design Choices (系统级设计选择)
  • 训练感知的稀疏注意力 (2025; 2025; 2025) 开始考虑内核实现和高效执行
  • 为了实现块稀疏注意力的高效实现,FlashAttention (2023) 被用于在高效的平铺机制中进行注意力计算,这为更好地利用硬件资源带来了要求和机会,包括:
    • 为避免内存访问不一致,在 SeerAttention (2024) 和 MInference (2024) 中,块大小 \(b\) 通常设置为至少 64 的相对较大值
    • 为了与 GPU 张量核心上分组矩阵乘法指令(Grouped Matrix Multiplication)的最低要求对齐,在 NSA (2025) 和 InfLLM-v2 (2025) 中,一个查询组内的 K、V 头数被设置为至少 16
    • 为了减少内存访问,NSA (2025) 和 InfLLM-v2 (2025) 强制查询组之间共享选定的块,这是通过在查询组内对块级重要性分数进行池化来完成的

Clustering Attention (聚类注意力)

  • 与块稀疏注意力类似,聚类注意力旨在为解码选择最关键地 Token ,但将 Token 组织在数据结构中以获得更好的语义属性或采样效率
  • RetrievalAttention (2024)
    • 采用近似最近邻搜索(Approximate Nearest Neighbor Search, ANNS)来选择关键的 K 个聚类
    • 为了解决注意力机制中查询向量和键向量之间的分布外(out-of-distribution)性质带来的挑战,它引入了一种适应查询向量分布的注意力感知向量搜索算法
  • ClusterKV (2024)
    • 在语义聚类的粒度上选择 Token ,克服了诸如 Quest 等页面级检索方法内部碎片化的问题
    • 在预填充阶段之后, Token 通过 K-means 算法进行聚类
    • Token \(i\) 和 \(j\) 之间的语义相似性通过键向量的余弦相似性度量:\(\mathcal{D}(i, j) = 1 - \frac{\langle k_i, k_j \rangle}{|k_i| \cdot |k_j|}\)
    • 语义聚类由其质心 \(\mu_1, \mu_2, …, \mu_C \in \mathcal{R}^d\) 表示
    • 在每个解码步骤,基于查询 Token \(q\) 和质心 \(\mu_i\) 的注意力权重(即 \(q \mu_i^T\))的排名来选择聚类
  • MagicPIG (2024)
    • 利用局部敏感哈希(Locality Sensitive Hashing, LSH)采样来高效近似注意力计算
    • 它使用 LSH 将相似的查询向量和键向量映射到相同的哈希桶中,并将存储和部分计算卸载到 CPU,以解决 KV 缓存的瓶颈
    • 它还引入了 Oracle Top-K 采样作为比暴力 Top-K 更好的策略

Bidirectional Sparse Attention (双向稀疏注意力)

  • 双向稀疏注意力建立在编码器风格的架构上,使用静态模式或块级稀疏性来加速注意力计算
  • 块稀疏性在双向稀疏注意力中被广泛使用
    • BigBird (2020)
      • 使用块级随机注意力,作为缩短 Token 之间间接路径的桥梁
    • Longformer (2020)
      • 使用静态的全局-局部混合注意力
      • Longformer 也依赖于块级稀疏性,并带有额外的全局和随机链接,以促进结构化计算和内存高效的并行性
  • 基于聚类的方法也用于双向稀疏注意力
    • Reformer (2020)
      • 使用局部敏感哈希(LSH)将相似的 Token 分配到同一个桶中
    • Routing Transformer (2020)
      • 在每一层执行在线 \(k\)-means 聚类
    • ClusterFormer (2022)
      • 引入了一个与下游目标共同训练的可微分聚类模块
      • 这些方法通过分组相关 Token 来减少计算,同时通过学习的适应性来保持性能

采用高效注意力机制的预训练大语言模型

采用统一高效注意力机制的预训练模型

  • 尽管早期的线性注意力探索通常局限于小规模模型,但近期的进展已证明其能够成功扩展到数十亿参数的规模,使其成为标准 Transformer 的一种可行且高效的替代方案
  • 这些模型完全基于线性注意力或其架构等价物,如状态空间模型和循环神经网络,即使在大规模下也能保持其标志性的推理效率
  • 基于 RWKV 的模型
    • RWKV 项目代表了一项持续且有影响力的努力,旨在创建一个可扩展的循环神经网络架构,它结合了 Transformer 的可并行化训练与传统 RNN 的高效推理 (2023)
    • 例如,EAGLE 系列引入了矩阵值状态以增加容量,而后续迭代如 Finch (RWKV-6) (2024) 和 Goose (RWKV-7) (2025) 则结合了动态循环和更具表达力的状态演化机制,以支持更复杂、数据依赖的状态转换
  • 基于 Mamba 的模型
    • Mamba 架构的成功及其数据依赖的选择机制,已引发了主要研究实验室的一波采用和扩展浪潮
    • Falcon Mamba (2024)
      • 基于纯 Mamba 架构,在一系列通用语言基准测试中展现出与领先 Transformer 模型相竞争的性能,验证了该架构在此类任务上的可行性,同时保持了其标志性的恒定时间推理
    • Codestral Mamba (2024)
      • 基于 Mamba-2 架构,进一步证明了该范式的潜力
      • 虽然专门用于代码生成,但它在相关基准测试中取得了最先进的结果,并支持 256K 个 Token 的上下文长度,展示了 SSM 方法在复杂结构化领域内的可扩展性和有效性
  • 基于稀疏注意力的模型
    • MiniCPM-4 (2025) 引入了一种两阶段稀疏注意力机制,根据语义相似性为每个查询 Token 动态选择相关的键值块
    • MiniCPM-4 利用 InfLLM-v2(一种块稀疏注意力变体)来替代标准注意力机制
    • 此外,一种轻量级的 LogSumExp 近似实现了高效的 top-k 选择,使得该方法能够扩展到极长序列
    • 这些技术共同使 MiniCPM-4 能够在细粒度上下文感知能力与可控的内存和计算需求之间取得平衡,使其成为长上下文建模的有力候选者

采用混合高效注意力机制的预训练模型

  • 随着对高效长上下文建模和多样化计算范式需求的增长,最近的研究广泛探索了混合注意力机制
  • 此类策略结合了全局和局部注意力组件,通常交错使用专门设计的层,以平衡计算成本和性能
  • 候选模型架构示意图如图 4 所示
  • 稀疏混合模型
    • GPT-3 (2020) 通过交错使用稠密注意力和局部带状稀疏注意力层,集成了一种混合注意力机制,其灵感来自 Sparse Transformer (2019)
    • 稠密注意力提供全上下文建模,而稀疏层则采用固定或跨步模式来减少被关注的 Token 数量
    • 这种设计使得 GPT-3 能够在 2048 个 Token 的固定上下文窗口内高效扩展到大规模模型,平衡了建模能力和计算效率
  • 线性-全注意力混合模型
    • Jamba (2024) 和 MiniMax-01 (2025) 结合了线性注意力和全注意力层,以在吞吐量和表达力之间实现高效的权衡
    • MiniMax-01 在大多数层中使用 Lightning Attention,并每八层插入一次基于 Softmax 的全注意力
    • Jamba 采用了相似的比例,在每八层的 Mamba 块中插入一个 Transformer 层
    • 两者都通过限制计算密集的全注意力的使用,实现了更快的解码和改善的长序列性能
  • 局部-全注意力混合模型
    • Gemma 3 (2025)、Command A (2025) 和 LLaMA-4-Maverick (2025) 在局部和全局注意力层之间交替使用,其共享的设计理念是稀疏地使用全局层(例如,每 4-6 层一次)以提升效率
    • 虽然局部层采用滑动窗口模式,但关键区别在于位置编码策略
    • Gemma 3 调制 RoPE 的基础频率——为局部层分配 10K,为全局层分配 1M(以更好地捕获长距离依赖关系)
    • Command A 和 LLaMA-4-Maverick 混合了基于 RoPE 的局部层与完全省略位置嵌入的全注意力层,从而实现了更强的长序列性能
  • 先进混合模型
    • Character.AI (2024) 将滑动窗口的局部注意力与每六层应用一次的稀疏全局注意力层交错使用
    • 特别是,它们在多个非相邻层之间复用全局注意力层的键值表示
    • 这种 KV 共享机制能以减少的内存和延迟开销实现高效的长上下文处理
  • YOCO (2024) 和 Phi-4-mini-flash (2025) 采用了一种双解码器架构,将预填充阶段和生成阶段分离开来
    • 自解码器在预填充和生成中都使用 RetNet 和滑动窗口注意力等线性注意力机制,而交叉解码器仅在生成期间激活
    • 整个过程中使用单层全局 KV 缓存,实现了线性时间的预填充和高效的解码,同时 GPU 内存消耗最小
  • In summary,这些最新进展强调了混合注意力机制以在不同计算约束和序列长度下实现平衡性能的趋势。每种架构都独特地贡献了关于如何有效结合局部细节管理与全局上下文整合的见解,从而为未来注意力机制的发展提供了有价值的框架

Outlook

  • 本综述全面概述了高效注意力机制,重点关注其算法基础、实际实现以及在大规模预训练语言模型中的集成
  • 通过将线性和稀疏注意力分类为明确定义的范式,论文识别了实现可扩展性、计算效率和长上下文能力的关键设计原则
  • 论文还分析了这些机制在最先进模型中如何部署,无论是作为独立架构还是作为平衡局部和全局计算的混合设计的一部分
  • 展望未来,论文强调几个预计将塑造该领域未来研究的关键方向:
  • 对混合模型的架构理解
    • 虽然先前关于线性注意力的工作主要集中于独立的线性架构,但混合模型通常通过结合现成的线性注意力模块与稠密或局部组件来构建
      • 但更强的线性主干是否直接转化为改进的混合性能尚不清楚
    • 未来的工作应将混合模型作为一个独特的架构类别来研究,试图理解它们的组成、相互作用效应和优化动态
  • 无损稀疏注意力与扩展上下文
    • 稀疏注意力仍然受到精度和计算增益之间权衡的挑战。完全训练的稀疏模型通常性能不如稠密模型,而训练后稀疏近似则由于缺乏端到端训练而面临限制
    • 一个主要的研究前沿在于开发能够保持稠密注意力的表达力和精度,同时扩展到更长上下文的稀疏注意力机制
    • Besides,稀疏预算与上下文长度之间的关系尚不明确,固定的 top-k 方案在序列更长时可能会退化,这需要更具适应性的策略
  • 对稀疏和混合注意力的机制性洞见
    • 实证研究反复证明,混合注意力模型可以用更少的注意力计算匹配甚至超越稠密模型,但其有效性的根本原因仍未得到充分探索
    • 此外,研究在合成基准测试中表现良好的稀疏模式是否适用于现实世界任务,以及描述基于稀疏性的泛化极限,尤为重要
  • 随着基于注意力的模型不断发展,论文预计架构创新、理论洞见和硬件感知设计之间将进一步融合
  • 作者希望本综述能为未来高效、高性能语言建模系统的研究奠定坚实的基础