过去两年里大模型领域的迭代速度快到让人应接不暇但有一个事实很容易被忽略真正处于地基位置的模型架构其实一直没有变。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 TransformerViT、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. 五大候选方向它们凭什么挑战 Transformer3.1 状态空间模型MambaMamba 是这两年最受关注的新架构之一。它建立在状态空间模型 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.xCPU 环境即可运行有 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, maskNone): # 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, keepdimTrue).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, batch2, d_model128, n_heads4, warmup3, repeats10): 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_model128, n_heads4) linear_attn LinearAttention(d_model128, n_heads4) t1 bench_attention(self_attn, seq_len) t2 bench_attention(linear_attn, seq_len) print(fseq_len{seq_len:5d} SelfAttention{t1:.4f}s fLinearAttention{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 的原理和新架构的改进点一起吃透。下次看到新模型发布时你只需要问一句它到底在哪一层降低了复杂度又为此付出了什么代价能答出这个问题你就已经比大多数只会追热点的人走得更远了。 Transformer会被取代吗？解析自注意力与Mamba等新型架构
过去两年里大模型领域的迭代速度快到让人应接不暇但有一个事实很容易被忽略真正处于地基位置的模型架构其实一直没有变。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 TransformerViT、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. 五大候选方向它们凭什么挑战 Transformer3.1 状态空间模型MambaMamba 是这两年最受关注的新架构之一。它建立在状态空间模型 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.xCPU 环境即可运行有 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, maskNone): # 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, keepdimTrue).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, batch2, d_model128, n_heads4, warmup3, repeats10): 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_model128, n_heads4) linear_attn LinearAttention(d_model128, n_heads4) t1 bench_attention(self_attn, seq_len) t2 bench_attention(linear_attn, seq_len) print(fseq_len{seq_len:5d} SelfAttention{t1:.4f}s fLinearAttention{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 的原理和新架构的改进点一起吃透。下次看到新模型发布时你只需要问一句它到底在哪一层降低了复杂度又为此付出了什么代价能答出这个问题你就已经比大多数只会追热点的人走得更远了。