过去两年里，大模型领域的迭代速度快到让人应接不暇，但有一个事实很容易被忽略：真正处于地基位置的模型架构，其实一直没有变。2017 年论文《Attention Is All You Need》提出的 Transformer，至今仍是 GPT、BERT、LLaMA、Qwen 等主流模型的核心骨架。不过最近 Mamba、RWKV、RetNet 等新架构频繁出现在技术社区，“下一代架构会不会取代 Transformer”的讨论也越来越多。尤其当长文本、端侧推理、实时交互成为业务刚需后，Transformer 在训练并行性和推理效率上的矛盾，被进一步放大了。
本文不打算替哪种架构“站队”，而是把 Transformer 的核心优势、真实瓶颈，以及那些被称为“替代者”的候选架构放在一起拆解。为了有直观感受，我会用一段可运行的 PyTorch 代码，对比自注意力与线性注意力在计算方式上的本质差异。无论你是算法工程师、后端开发，还是刚入门 AI 的学习者，这篇文章都能帮你建立一条判断线索：新架构到底新在哪里，它凭什么挑战 Transformer，又付出了哪些代价。
1. 为什么大家都在讨论“取代 Transformer”
1.1 一个容易被忽略的事实：Transformer 已经红了近八年
Transformer 从提出到现在，已经有近八年时间。这八年里，它先是取代了循环神经网络 RNN 和长短期记忆网络 LSTM，成为自然语言处理的主流架构；接着又通过 Vision Transformer（ViT）、Swin Transformer 等衍生模型，把注意力机制带进了计算机视觉领域；再后来，多模态模型、语音识别、推荐系统，甚至蛋白质结构预测，都开始使用 Transformer 或它的变体。
为什么它能火这么久？核心在于自注意力机制解决了此前序列模型最头疼的两个问题：一是长距离依赖，RNN 家族在处理长文本时，很容易把早期信息“遗忘”掉，Transformer 的每个 token 都能直接和序列中任意位置的 token 建立连接；二是并行计算，RNN 必须按时间步逐个计算，而 Transformer 可以把整个序列当作一个矩阵并行处理，训练效率高出好几个量级。正因如此，“为什么最后是 Transformer”才会成为很多学习者的共同疑问，也催生了大量“手撕 Transformer”“Transformer 源码解析”类的教程。
1.2 讨论替代者的三条核心线索
既然 Transformer 这么强，为什么还会有人想“取代”它？目前社区里的讨论，基本围绕三条线索展开。
第一条是计算复杂度。经典自注意力的时间复杂度和空间复杂度都是 O(n²)，n 是序列长度。当序列从 2048 涨到 128K 甚至更长时，注意力矩阵会呈平方级膨胀，数据量非常惊人。
第二条是推理阶段的效率问题。训练时 Transformer 可以高度并行，但在自回归生成时，每生成一个 token 都要重新计算或读取之前的 Key/Value 缓存（KV Cache），序列越长，KV Cache 越大，显存占用和延迟也随之上升。这直接影响到长对话、长文档生成、实时语音交互等实际业务场景。
第三条是架构本身的探索空间。既然注意力的完整矩阵计算这么贵，能不能用其他机制替代？于是出现了状态空间模型、线性注意力、稀疏注意力、混合架构等一系列新方向。它们不是简单优化某一个算子，而是从顶层设计上改变“信息如何传递”。
1.3 本文要解决的问题
这篇文章不打算给出一个“谁取代谁”的绝对结论，因为技术演进从来不是二选一。我更想做的，是把 Transformer 的“不可替代之处”和“被挑战的原因”讲清楚，再带大家看几个典型候选架构，并用代码展示其中关键的计算差异。
读完本文，你会理解：Mamba 的 O(n) 复杂度是怎么来的，线性注意力为什么能降低开销，混合架构为什么可能是未来方向，以及在实际工程中，模型选型到底该关注哪些指标。
2. Transformer 的核心优势与真实瓶颈
2.1 自注意力机制到底解决了什么
先来看自注意力的基本思想。给定一个长度为 n 的输入序列，每个 token 会生成三个向量：Query（查询）、Key（键）、Value（值）。注意力分数通过 Query 和 Key 的内积计算，表示“当前 token 应该多关注另一个 token”，再经过 Softmax 归一化，最后对 Value 做加权求和。
这听起来不复杂，但它带来的效果非常关键。以前的 RNN 想获取某个历史信息，需要沿着时间步一步一步“传递”，路径很长，信息容易丢失和扭曲。而 Transformer 的每个 token 都可以直接和序列中任意位置的 token 互动，相当于一条“信息直达通道”。这让模型在理解长距离依赖时，不再依赖信息在链式结构中层层传递。
同时，这个操作全部是矩阵乘法和 Softmax，非常适合 GPU 加速。换句话说，Transformer 之所以能在过去几年快速扩张，不只是结构上的创新，还因为它和现代硬件高度契合，这为大规模分布式训练提供了天然便利。
2.2 长序列场景下的三个硬伤
如果说短序列场景里 Transformer 是无冕之王，那么长序列场景它就是“高成就高负担”的代表。
第一是计算复杂度高。标准的自注意力需要计算一个 n×n 的注意力矩阵，每一层都是 O(n²) 级别的计算量。序列长度翻倍，计算量变成四倍；长度增长到 10 倍，计算量就是 100 倍。对 128K 甚至 1M 上下文来说，这是一个非常昂贵的代价。
第二是显存占用高。注意力矩阵本身要保存在显存里，训练时还要保存中间梯度。即使只用推理模式，KV Cache 也会随着序列长度线性增长。很多团队在做长文档问答时，明明模型支持长上下文，实际部署却被显存卡住，这就是原因之一。
第三是推理阶段串行。虽然训练可以并行，但在自回归生成时，模型必须一个 token 接一个 token 地生成，无法批量预知未来内容。每生成一个 token，都需要从 KV Cache 中读取历史信息，这导致“训练快、推理慢”的典型现象。对于实时聊天、语音助手这类低延迟场景，这个矛盾尤其突出。
2.3 为什么“并行计算能力强”也会成为束缚
这里有一个容易被新手忽略的细节：Transformer 的并行优势主要体现在训练阶段，尤其是在非自回归任务中。一旦进入自回归解码，模型在时间维度上依然是一个串行过程。你可以把注意力计算看作一个“读取全部历史再做加权”的操作，这个操作本身可以并行，但“必须等前一个 token 生成完才知道下一个 token 的输入”这件事，是无法用并行解决的。
所以业内很多优化工作都集中在两个方向：一是减少单次注意力计算的代价，比如 FlashAttention 通过分块计算降低显存读写；二是减少需要重复读取的历史信息，比如各种线性注意力、状态空间模型。理解了这一点，再看后面的候选架构，思路就会清晰很多。
3. 五大候选方向：它们凭什么挑战 Transformer
3.1 状态空间模型：Mamba
Mamba 是这两年最受关注的新架构之一。它建立在状态空间模型 S4 的基础上，核心思想是用一个固定维度的隐状态 h 来压缩序列信息，而不是维护 n×n 的注意力矩阵。每一步计算只依赖当前输入和上一步的隐状态，复杂度降为 O(n)。
更关键的是，Mamba 引入了一种“选择性扫描”机制，让模型可以根据当前输入动态决定“记住什么、遗忘什么”，解决了原始状态空间模型在内容感知能力上的不足。从公开实验看，Mamba 在长序列任务上展示出非常有竞争力的表现，同时推理时的内存占用更稳定，不会像 Transformer 那样随序列长度线性膨胀。
不过 Mamba 也有代价。它看起来像 RNN，时间步之间存在依赖，要发挥硬件效率，需要依赖并行扫描算法和高度定制的 CUDA 算子。这导致它不像 Transformer 那样开箱即用，很多优化技巧需要重新积累。如果你只在 CPU 或者普通 GPU 上跑，效果未必能超过优化得很好的 Transformer。
3.2 线性注意力：RWKV、RetNet
线性注意力是另一条重要路线。它的核心思路很简单：把注意力矩阵中的 Softmax 展开成特征映射的形式。原来计算 Q 和 K 的相似度需要完整矩阵乘法，复杂度 O(n²)；而如果可以把 Q 和 K 映射到高维空间，让相似度近似为内积，就可以调整矩阵乘法的顺序，变成“先算 KV，再乘 Q”，复杂度降到 O(n)。
RWKV 是这条路线里很有代表性的开源项目，它把 Transformer 的训练并行性和 RNN 的推理高效性结合起来，用线性注意力替代标准自注意力。RetNet 则提出了 retention 机制，既能像 Transformer 一样并行训练，也能像 RNN 一样循环推理，还支持分块递归的折中方案。
这条路线的问题在于，Softmax 的近似不是免费的，单纯换成线性核函数后，模型可能在部分任务上精度下降、训练不稳定。因此实际实现往往要配合额外的位置编码、门控机制、归一化策略，复杂度并不比 Transformer 低太多。
3.3 混合架构：Jamba 等
既然纯 Transformer 和纯 Mamba 各有优势，自然有人想到“把两者组合在一起”。AI21 实验室发布的 Jamba 就是代表性案例。它把 Transformer 层和 Mamba 层交替堆叠，同时引入混合专家（MoE）机制来扩大参数规模、控制计算量。
混合架构的逻辑很清楚：让 Transformer 负责全局信息交互和成熟稳定的建模能力，让 Mamba 负责高效的序列压缩和低延迟推理。两者互补。从工程角度看，混合架构可能是未来最容易被落地的一类方案，因为团队不需要完全抛弃过去围绕 Transformer 积累的优化经验和生态工具。
可以预见，未来开源大模型里会出现越来越多“Transformer + SSM”或者“Transformer + 线性注意力”的组合结构，而不是单一架构的彻底替换。
3.4 稀疏注意力与滑动窗口
除了改变注意力本身，另一个思路是“不把所有位置都纳入注意力范围”。Longformer、BigBird 用稀疏注意力模式，让每个 token 只关注固定窗口内的邻近 token，再用少量全局 token 负责远程信息汇总。Swin Transformer 在视觉任务里也采用了类似思路，把注意力限制在局部窗口内，再用移位窗口实现跨窗口信息交换。
这类方法严格来说不是“取代 Transformer”，而是“改造注意力”。它的优点是能兼容现有 Transformer 生态，损失相对可控；缺点是全局信息的获取需要额外机制兜底，当序列特别长时，如何设计稀疏模式依然是一个需要经验和实验的问题。
3.5 为什么目前还没有一个“完全替代者”
既然新架构这么多，为什么它们还没有真正取代 Transformer？主要原因是“架构能力”和“工程生态”是两码事。
Transformer 积累了整整八年的工程红利：FlashAttention、分布式训练框架、各种量化工具、推理引擎、硬件适配，几乎全部优先支持它。一个新架构即使理论上更高效，也要重新适配这些基础设施。另外，大模型的训练效果不只取决于架构，还取决于数据规模、训练策略、对齐技术。架构只是其中一个变量，很难单独说明“谁比谁强”。
所以在现阶段，我更愿意把 Mamba、RWKV 这些方案看作 Transformer 的“补充者”和“竞争者”，而不是“即刻替代者”。
4. 用代码理解：自注意力 vs 线性注意力
前面讲的都是概念，下面用代码把“为什么线性注意力复杂度低”这件事说清楚。这里不会实现一个完整可用的 Mamba，而是用两个最小版本展示核心计算差异。
4.1 环境准备
本文示例代码需要以下环境：
- Python 3.9 或更高版本。
- PyTorch 2.x，CPU 环境即可运行，有 NVIDIA GPU 会更快。
- 无需额外数据集，使用随机张量做演示。
项目只用一个文件，建议命名为test_attention.py。版本可以根据你的实际环境调整，重点是理解计算逻辑。
4.2 标准 SelfAttention 最小实现
先实现一个简化版的多头自注意力。为了便于理解，这里去掉了残差连接和层归一化，只保留核心路径。
import math import torch import torch.nn as nn class SelfAttention(nn.Module): def __init__(self, d_model, n_heads): super().__init__() self.d_model = d_model self.n_heads = n_heads self.d_k = d_model // n_heads self.w_q = nn.Linear(d_model, d_model) self.w_k = nn.Linear(d_model, d_model) self.w_v = nn.Linear(d_model, d_model) self.w_o = nn.Linear(d_model, d_model) def forward(self, x, mask=None): # x 形状: (batch, seq_len, d_model) batch, seq_len, _ = x.shape Q = self.w_q(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2) K = self.w_k(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2) V = self.w_v(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2) # 注意力分数: (batch, n_heads, seq_len, seq_len) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores = scores.masked_fill(mask == 0, float("-inf")) attn = torch.softmax(scores, dim=-1) out = torch.matmul(attn, V) out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1) return self.w_o(out)这里最值得注意的就是scores这一行。Q和K做矩阵乘法之后，形状是(batch, n_heads, seq_len, seq_len)，也就是每个 head 都维护一个 n×n 的注意力矩阵。序列长度 n 越大，这个矩阵占用的显存就越大，复杂度也呈平方级增长。
/ math.sqrt(self.d_k)的作用是缩放，防止内积数值过大导致 Softmax 梯度消失，这是 Transformer 原论文里的细节。
4.3 LinearAttention 最小实现
线性注意力的核心，是避免显式构造 n×n 的注意力矩阵。我们用一个非线性函数elu(x) + 1当作核函数，把 Q 和 K 变换到非负特征空间，然后交换矩阵乘法顺序。
class LinearAttention(nn.Module): def __init__(self, d_model, n_heads): super().__init__() self.d_model = d_model self.n_heads = n_heads self.d_k = d_model // n_heads self.w_q = nn.Linear(d_model, d_model) self.w_k = nn.Linear(d_model, d_model) self.w_v = nn.Linear(d_model, d_model) self.w_o = nn.Linear(d_model, d_model) @staticmethod def _feature_map(x): # 使用 elu + 1 保证输出非负，替代 softmax 的近似核函数 return torch.nn.functional.elu(x) + 1.0 def forward(self, x): batch, seq_len, _ = x.shape Q = self._feature_map(self.w_q(x)).view( batch, seq_len, self.n_heads, self.d_k ).transpose(1, 2) K = self._feature_map(self.w_k(x)).view( batch, seq_len, self.n_heads, self.d_k ).transpose(1, 2) V = self.w_v(x).view( batch, seq_len, self.n_heads, self.d_k ).transpose(1, 2) # 先计算 KV: (batch, n_heads, d_k, d_k) KV = torch.matmul(K.transpose(-2, -1), V) # 再计算 Q @ KV: (batch, n_heads, seq_len, d_k) out = torch.matmul(Q, KV) # 归一化分母: (batch, n_heads, seq_len, 1) z = torch.matmul(Q, K.sum(dim=-2, keepdim=True).transpose(-2, -1)) out = out / (z + 1e-6) out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1) return self.w_o(out)关键区别在第 33 行左右：先算K.transpose(-2, -1) @ V，得到一个(d_k, d_k)的小矩阵，再让所有 Q 去乘这个矩阵。整个过程不再需要构造 n×n 的注意力矩阵，复杂度从 O(n²) 降到 O(n)。
需要说明的是，这只是教学演示版。真实线性注意力还要解决因果掩码、位置编码、数值稳定性、训练收敛等问题，不能用这段代码直接替代成熟实现。
4.4 对比实验脚本
把上面两个类放到同一个文件中，然后运行下面的脚本，观察不同序列长度下的耗时变化。
import time def bench_attention(attn, seq_len, batch=2, d_model=128, n_heads=4, warmup=3, repeats=10): x = torch.randn(batch, seq_len, d_model) for _ in range(warmup): attn(x) if torch.cuda.is_available(): torch.cuda.synchronize() start = time.time() for _ in range(repeats): attn(x) if torch.cuda.is_available(): torch.cuda.synchronize() return (time.time() - start) / repeats if __name__ == "__main__": for seq_len in [64, 256, 1024, 2048]: self_attn = SelfAttention(d_model=128, n_heads=4) linear_attn = LinearAttention(d_model=128, n_heads=4) t1 = bench_attention(self_attn, seq_len) t2 = bench_attention(linear_attn, seq_len) print(f"seq_len={seq_len:5d} SelfAttention={t1:.4f}s " f"LinearAttention={t2:.4f}s")代码里的warmup是为了让显存分配、CUDA 内核加载等预热完成，减少偶然波动。如果你在 CPU 上运行，耗时规律同样可以参考。
4.5 结果怎么看
预期结果应该是：序列较短时，两种注意力耗时差别不大；随着 seq_len 增加，SelfAttention 的耗时增长越来越快，而 LinearAttention 更平缓。
不过这里要特别提醒：这个实验只对比了两个单独的注意力模块，不代表完整大模型。真正的大模型还有 FFN、LayerNorm、嵌入层、采样解码等大量环节。Mamba 的“O(n) 复杂度”也需要配合高度优化的算子才能发挥出来。所以千万不要因为一个玩具实验，就断定新架构一定更快。
但这个代码能帮助我们建立最核心的直觉：自注意力的瓶颈在于 n×n 矩阵，线性注意力的思路就是“绕开”这个矩阵。
5. 常见问题与排查思路
5.1 训练不收敛怎么办
使用线性注意力或状态空间模型时，比 Transformer 更容易遇到训练不稳定、Loss 突刺、不收敛等问题。
常见原因有三个：一是核函数和位置编码实现不对，导致信息无法有效区分；二是数值范围没有控制好，内积结果差异过大；三是学习率策略不适合新架构。
排查时可以先用小规模数据、小模型复现论文结果，确认实现正确；接着检查注意力输出和梯度范数是否出现异常值；最后调整学习率，线性注意力往往需要更小的峰值学习率或更长的 warmup。
5.2 显存不够怎么办
如果是训练 Transformer 长序列导致显存不够，优先考虑梯度累积、降低 batch size、使用 FlashAttention 或梯度检查点。如果模型本身已经切换到线性注意力，但显存依然不够，常见原因是 FFN 层或 KV Cache 的优化没有做。
另外，序列长度并不是唯一影响显存的因素，batch size、head 数量、隐层维度都会放大占用。排查时先用一个最小配置跑通，再逐步增大，定位是哪一部分开始爆显存。
5.3 推理速度仍然慢怎么办
很多人在测试 Mamba 或 RWKV 时发现，实际推理速度没有理论预期那么快，甚至比优化后的 Transformer 还慢。
问题往往出在算子实现上。RNN 式的循环结构在短序列上无法充分利用 GPU 并行能力，而 Transformer 的矩阵乘法已经被优化到非常成熟。要让新架构真正跑出优势，需要配套的 CUDA kernel、批处理策略和内存布局优化。
建议直接用官方开源实现，而不是自己重写算子。使用前先跑官方 benchmark，确认在你的硬件和后端版本下能达到预期，再接入业务。
5.4 常见问题速查表
| 问题现象 | 常见原因 | 解决思路 |
|---|---|---|
| 训练 Loss 不下降 | 位置编码或核函数实现错误 | 用小模型复现论文，检查核心代码 |
| Loss 突然变为 NaN | 数值溢出或学习率过大 | 加梯度裁剪、降低学习率、检查归一化 |
| 推理显存增长快 | KV Cache 未优化 | 使用 PagedAttention、换混合架构 |
| 显式注意力矩阵 OOM | 序列过长、head 太多 | 降低 batch、用 FlashAttention、加梯度检查点 |
| 新架构实测速度慢 | 算子未融合或硬件适配差 | 使用官方 kernel，先跑 benchmark |
| 长文本下游任务掉效果 | 线性近似损失精度 | 改用混合架构或加全局 token 机制 |
6. 工程实践建议：选型、部署与未来
6.1 什么场景继续用 Transformer
在绝大多数短中长度序列任务里，Transformer 加上 FlashAttention 仍然是最稳妥的选择。原因很直接：生态成熟、资料多、硬件支持好、踩坑成本低。比如常见的 4K 到 8K 上下文的对话系统、RAG 问答、文本分类，用标准的 Transformer 架构没有明显短板。
如果你的团队时间紧、任务重，我建议不要为了“追新”而贸然切换到新架构。架构选型是成本很高的事，数据规范、训练代码、推理链路、评测体系都要跟着变。
6.2 什么场景可以尝试新型架构
如果你面临下面几类情况，可以认真考虑 Mamba、RWKV 或混合架构：
- 超长上下文：比如处理几十万 token 的代码仓库、整本书级别的长文档。
- 低延迟推理：语音助手、实时翻译等对每 token 生成时间很敏感的业务。
- 端侧部署：内存和算力受限，希望模型随序列增长时显存更可控。
即便如此，也建议先在离线任务上做小规模验证，确认精度损失在可接受范围内，再进入线上。
6.3 部署时关注什么
部署层面不能只看参数量和理论复杂度。要把下面几个问题纳入评估：
- 模型格式兼容性：ONNX、TensorRT、llama.cpp 等工具对 Mamba 这类新架构的支持可能滞后，需要确认是否支持导出和量化。
- 量化敏感性：RNN 式递归结构的量化难度通常比 Transformer 更高，小比特量化后效果跌落需要额外测试。
- 评测体系：不要只盯困惑度，至少要在长文本检索、多轮对话、代码生成等下游任务上对比效果和延迟。
6.4 给学习者的建议
如果你还在学习阶段，我建议坚持做两件事。
第一，把 Transformer 源码吃透。理解自注意力、多头机制、位置编码、KV Cache，这些都是后续所有架构的基础。不要因为大家都在讨论“取代 Transformer”，就跳过这门基本功。
第二，主动读新架构的论文和官方代码。Mamba、RWKV、RetNet 的实现并不算特别长，但每个细节都浓缩了研究者的工程思考。读源码时多问一句：它把复杂度转移到了哪里？训练快还是推理快？需要什么硬件支持？
7. 写在最后：不是取代，而是交替迭代
回到开头那个问题，“下一场 AI 革命，要取代 Transformer 吗？”如果只看新闻标题，总觉得有一种架构要终结另一种架构。但从论文、代码和工程实践来看，更准确的说法可能是：下一代模型不会只靠某一个单一架构吃遍天，而是会在不同层级把自注意力、线性注意力、状态空间模型、稀疏化组合起来。
对于普通开发者，短期内能直接感受到的，仍然是 API 和开源模型的变化，而不是底层某个算子的替换。与其焦虑架构会不会被淘汰，不如把 Transformer 的原理和新架构的改进点一起吃透。下次看到新模型发布时，你只需要问一句：它到底在哪一层降低了复杂度，又为此付出了什么代价？能答出这个问题，你就已经比大多数只会追热点的人，走得更远了。
