挑战Transformer架构的前沿模型:Mamba、Hyena、RWKV?

超越O(n²)复杂度:挑战Transformer架构的前沿模型

一、 卫冕冠军的“阿喀琉斯之踵”

在介绍挑战者之前,我们必须首先理解Transformer架构为何会暴露出弱点。其强大的全局信息交互能力源于自注意力机制,但这同样也是其致命弱点:二次方复杂度 (Quadratic Complexity)

  • 计算量:对于一个长度为 n 的序列,自注意力需要计算每个元素与所有其他元素的关系,计算量与 成正比。
  • 内存占用:存储巨大的注意力分数矩阵也需要 级别的空间。

这意味着序列长度翻倍,计算和内存需求就会变成四倍。处理一篇几千词的文档已然非常吃力,而处理一整本书、一个代码库或DNA序列等超长序列,对标准Transformer而言几乎是不可能完成的任务。

因此,整个AI社区都在寻找一个核心问题的答案:是否存在一种架构,既能拥有Transformer强大的长程依赖建模能力,又能摆脱 的诅咒?以下四种技术路径为此提供了极具创造力的答案。

二、挑战 (1) :Mamba —— 具备选择性记忆的状态空间模型

核心原理:通过一个“选择性机制”,将RNN的线性复杂度与内容感知能力相结合,让模型能动态地记住关键信息、遗忘次要信息。

Mamba可以看作一种现代化的循环网络(RNN)。传统RNN的致命弱因在于其“记忆衰减率”是固定的,导致重要和不重要的信息以同等速度被遗忘。Mamba革命性地解决了此问题,实现了 O(n) 的线性复杂度。

其计算过程是顺序的,一次处理一个词元(token),并维护一个固定大小的“记忆状态”(state)h

  1. 当第 t 个词元 x_t 输入时,模型会更新记忆:h_t = A * h_{t-1} + B * x_t
  2. 并基于新记忆产生输出:y_t = C * h_t

此过程的计算量仅与状态 h 的大小相关,与序列长度 n 无关,因此总复杂度是线性的。

Mamba的创新之处:选择性机制

  • 在传统的状态空间模型(SSM)中,矩阵 A(决定遗忘多少旧记忆)和 B(决定吸收多少新信息)是固定不变的。
  • 在Mamba中,AB 不再是固定的,而是由当前输入 x_t 动态生成

这种设计带来了一种智能的状态管理机制:

  • 当一个关键信息(如主角名字)输入时,模型可以学会生成一个接近1的 A 值(“锁住记忆,不要遗忘”)和一个较大的 B 值(“完全吸收新信息”)。
  • 当一个无关紧要的词(如“的”、“了”)输入时,模型则可以生成一个较小的 A 值(“让旧记忆快速衰减”)和一个很小的 B 值(“忽略这个词”)。

通过这种方式,Mamba能动态管理其有限的记忆空间,确保最重要的信息在长序列中得到保持,从而高效地建模长距离依赖。

三、挑战 (1) :Hyena —— 基于长卷积的高效全局信息混合

为何突出Hyena?

在众多挑战者中,Hyena架构之所以备受关注,不仅因为它在理论上通过快速傅里叶变换(FFT)将卷积的复杂度从 O(n²) 优化至 O(n log n),更重要的是,它在实际应用中取得了巨大成功。从HyenaDNA在百万级长度的基因组学数据上取得突破,到StripedHyena混合模型在通用语言任务上比肩顶级的Transformer,Hyena用一系列成果证明了长卷积路线的强大潜力与实用价值,成为了该方向当之无愧的标杆。

核心原理:通过数学方法动态生成一个与输入序列等长的结构化卷积核,然后利用FFT进行一次高效的全局卷积操作,以模拟全局信息关联。

这等效于使用一个巨大的卷积掩码(mask),但它并不直接存储这个巨大的核,从而避免了 O(n) 的参数量和 O(n²) 的计算量。

1. Hyena的机制:参数化生成与FFT加速

Hyena的精妙之处在于,它不直接学习卷积核 h 的每一个值,而是学习一个生成 h 的函数

  • 第一步:长卷积核的参数化
    Hyena规定,长卷积核 h 必须遵循一个特定的数学形式,例如一个由位置 t 和几个可学习参数 a, b, c... 决定的函数:h(t) = f(t, a, b, c, ...)。这个函数形式(如阻尼振荡)天然适合模拟信息随距离衰减和传播的特性。如此一来,学习一个任意长度的卷积核,从学习 n 个值简化为学习少数几个参数。

  • 第二步:动态化生成
    更进一步,Hyena认为“回声”模式应根据内容动态调整。它引入了几个小型神经网络(MLP),这些网络接收当前词元 x_t 的向量作为输入,并输出用于生成长卷积核 h 的那一组最优参数 a, b, c...

  • 第三步:高效卷积运算
    有了输入信号 x 和生成的长卷积核 h,就需要执行卷积 y = x * h。直接计算的复杂度为 O(n²)。Hyena利用了数学中的卷积定理

    时域中的卷积运算,等价于频域中的乘法运算。

    这提供了一条高效的计算路径:

    1. 进入频域 (FFT):对输入序列 x 和卷积核 h 分别执行快速傅里叶变换,得到它们的频谱 FFT(x)FFT(h)。此步复杂度为 O(n log n)
    2. 频域混合 (Element-wise Product):将两个频谱逐点相乘 FFT(y) = FFT(x) ⊙ FFT(h)。此步极其高效,复杂度仅为 O(n)
    3. 返回时域 (Inverse FFT):对结果 FFT(y) 执行逆傅里叶变换,得到最终输出 y。此步复杂度为 O(n log n)

通过这条路径,总复杂度从 O(n²) 成功降至 O(n log n),使得Hyena能够高效处理超长序列。

2. Hyena的应用与挑战

  • 开源情况与发展:Hyena的核心论文于2023年初发布,其代码库(如HazyResearch/safari, togethercomputer/stripedhyena)也随之开源。在不到两年的时间里,它经历了从理论奠基、到在基因组学领域的“杀手级应用”(HyenaDNA)、再到与注意力机制融合的通用语言模型(StripedHyena-7B),并开始探索边缘计算部署(Hyena Edge),发展速度惊人。

  • StripedHyena-7B的性能表现:作为一个混合架构(Hyena + Attention)模型,StripedHyena-7B在标准(短上下文)基准测试上,其性能与LLaMA 2 7B和Mistral 7B等顶级Transformer模型不相上下。而在其“主场优势”——长上下文任务中,随着序列长度增加到32k甚至128k,它在困惑度、训练/推理速度和内存占用方面的优势愈发明显,全面超越了依赖特定优化技术的Transformer模型。

  • 应用瓶颈

    1. 归纳偏置的局限性:卷积天然适合处理连续、结构化的数据。对于需要“跳跃式”逻辑和精确信息抽取的任务,其表现可能不如注意力机制。混合架构是目前最好的解决方案。
    2. 生态系统不成熟:与Transformer相比,Hyena在硬件优化、软件算子、社区资源和教程方面仍有差距,这增加了开发者的学习和迁移成本。
    3. 现实困境:随着FlashAttention等技术的普及,Transformer在处理中等长度上下文(8k-32k)时效率已大幅提升。Hyena需要证明其在更常见任务中的效率优势足以抵消迁移成本。

四、挑战 (3) :RWKV —— 将Transformer翻译成RNN形式

核心原理:通过巧妙的数学重写,将自注意力的核心计算过程,变换为一个可以顺序迭代的RNN形式,从而兼具Transformer的性能与RNN的 O(n) 效率。

RWKV的洞察在于,标准注意力 Output = Softmax(Q * K^T) * V 中的 Q * K^T 是计算瓶颈。如果用一个线性函数近似Softmax,那么注意力计算就可以被拆解。对于第 i 个词,其输出可以看作是与所有过去词 j (j < i)K_jV_j 的某种加权求和。

这个“加权求和”可以被整理成一个递归的形式:

当前状态 = 时间衰减因子 * 上一时刻状态 + 当前词的KV信息

这个“状态”向量包含了到目前为止所有历史信息的精华。每次有新词元输入,模型只需执行一次固定大小的更新操作,而无需重新计算与所有历史词元的关系。为了模拟注意力中“距离越远、关系越弱”的特点,RWKV还引入了一个可学习的“时间衰减因子”,让久远信息的权重在状态中逐渐降低。

这种滚动更新的机制,使其在保持强大性能的同时,实现了 O(n) 的线性计算复杂度和 O(1) 的内存增长。

五、挑战 (4) :线性注意力 —— 用数学近似绕过瓶颈

核心原理:通过改变注意力计算的顺序或用一个更简单的函数近似Softmax,来避免生成巨大的 n x n 注意力矩阵。

这是对Transformer最直接的“改良派”思路。其核心在于利用矩阵乘法的结合律

  1. 标准计算顺序(Q * K^T) * V。其中 Q * K^T 会产生一个 n x n 的大矩阵,这是瓶颈所在。
  2. 理想计算顺序Q * (K^T * V)。如果能先计算 K^T * V,会得到一个 d x d 的小矩阵(d是特征维度),其大小与序列长度 n 无关。整个计算的复杂度就变成了 O(n * d²),对 n 而言是线性的。

问题:标准Transformer无法这么做的原因是中间的 Softmax 操作。Softmax(Q * K^T) 作用于整个 n x n 矩阵,使其无法被拆分。

解决方案:线性注意力的核心就是用一个可以被拆分的、逐点计算的函数 phi(x) 来代替 Softmax
Softmax(QK^T) 被近似为 phi(Q) * phi(K)^T。这样一来,括号就可以自由移动,从而利用 Q * (phi(K)^T * V) 这个高效的计算顺序。

这种方法从“相对排名”(Softmax,每个元素的权重依赖于其他所有元素)转变为“独立贡献”(每个Key独立计算其贡献,然后汇总),从而实现了线性化。

六、 总结与展望

这四条技术路径,从不同角度为突破Transformer的效率瓶颈提供了可行的方案。

架构核心原理复杂度优势局限性
Mamba 选择性状态空间模型 O(n) 极高的推理效率,真正的线性复杂度,擅长连续数据。 状态压缩可能丢失信息,对离散、非结构化数据建模能力待验证。
Hyena 参数化长卷积 + FFT O(n log n) 极强的长程依赖建模能力,已在多个领域证明有效性。 卷积的归纳偏置不适合所有任务,生态系统尚不成熟。
RWKV 将注意力重写为RNN形式 O(n) 兼具Transformer性能和RNN效率,训练和推理都很快。 数学形式相对复杂,时间衰减因子的设计是关键。
线性注意力 近似Softmax以利用矩阵结合律 O(n * d²) 对原始注意力改动最小,易于理解和实现。 近似可能导致性能损失,尤其在需要精确注意力的任务上。

未来展望:未来的AI架构很可能不是“谁取代谁”的零和游戏,而是走向“融合共生”。正如StripedHyena所展示的,将高效的长程建模算子(如Hyena)与强大的局部推理算子(如Attention)相结合的混合架构,有望在实现更高性能与效率的同时,兼顾不同任务的需求,开启后Transformer时代的新篇章。

posted @ 2025-09-22 10:38  Engineblogs  阅读(82)  评论(0)    收藏  举报