要解决的问题

全局稠密注意力在全序列计算时具有平方复杂度,标准 KV 缓存随长度线性增长。部分循环式状态空间和线性注意力架构,可用固定大小的递推状态处理新 token;这里恒定的是序列状态的容量,并非训练时全部内存。保留全局注意力层的混合架构通常仍需随长度增长的 KV 缓存,不能一概称为常数推理内存。

核心方法

  • 状态空间模型(SSM):把序列压缩进一个固定大小的隐状态,按时间递推,推理时内存恒定。Mamba 让状态更新依赖输入(选择性机制),大幅提升了语言建模能力。
  • 线性注意力与线性 RNN:把注意力改写成可以递推计算的形式,如 RWKV、RetNet、门控线性注意力(GLA)。
  • Transformer 与状态空间模型混合:大部分层用状态空间模型,少数层保留注意力,如 Jamba。

代表工作

  • RWKV(2023):兼具 RNN 推理效率和 Transformer 并行训练能力的架构。
  • RetNet(2023):提出保持(retention)机制,同时支持并行、递推和分块三种计算方式。
  • Mamba(2023):选择性状态空间模型,在论文评测中展示了有竞争力的语言建模质量和高效的线性长度扩展;不据此断言它是所有非注意力架构中的「首次」。
  • GLA(2023):门控线性注意力,并给出硬件高效的训练算法。
  • Jamba(2024):AI21 发布的 Transformer、Mamba 与混合专家的混合架构。
  • Mamba-2(2024):揭示状态空间模型与注意力之间的对偶关系,训练速度大幅提升。

现状

这类架构在长序列的吞吐和内存上有明显优势,但纯状态空间或线性架构在精确召回上下文中的具体信息时,通常仍然弱于注意力。因此,混合架构是目前务实的主线。

开放问题

  • 长程的精确召回能力。
  • 注意力层与高效层的最优混合比例。
  • 生态与计算内核的成熟度。

参考资料

相关专题