★ FEATURED ARTICLE 做序列建模这几年最让我觉得反直觉的一个命题就是“注意力也是FFN”。初看这像是个文字游戏注意力机制明明是加权求和FFN是两层线性变换加激活两者八竿子打不着。但如果你把注意力里那个不断累积的状态矩阵看成“权重”把每个新来的token看成“一次权重更新”就会发现Transformer里的两大模块底层居然长得很像。这事的实际价值不只是视角新奇。它直接关系到高效注意力怎么做怎么把O(N²)的复杂度降下来怎么把KV cache压缩成一个固定大小的状态怎么让模型在推理时像RNN一样丝滑。更关键的是这个视角把线性注意力、状态空间模型、Mamba这类东西串到了同一条线上让你不再需要背一堆看似无关的架构而是回到同一个数学原点去理解。这篇文章就从“快权重”这个老概念出发把“注意力FFN”这条线完整捋一遍。我不打算只讲结论会把线性注意力的推导、状态递推、代码实现、数值稳定性、实验对标这些环节全过一遍最后附上我实际调参踩过的坑。适合已经用过Transformer、想理解高效注意力本质或者准备在长序列任务上替换标准注意力的读者。1. 先搞清楚为什么要把注意力说成FFN1.1 Transformer里被忽视的算力账先算一笔实际的账。标准softmax注意力给定序列长度N和特征维度d要对N个query分别计算与N个key的相似度复杂度是O(N²d)。当序列长度到8192pair数量就是6700万到16k就是2.5亿。这个数字意味着就算你有A100注意力分数矩阵本身就能把显存吃穿。但有意思的是Transformer里真正吃参数量和FLOPs的大户往往不是注意力而是FFN。FFN通常把d维映射到4d维再加回来这一下就是8倍的矩阵乘法开销。于是很多人会下意识觉得FFN是“干活”的注意力是“通信”的——一个负责变换特征一个负责融合信息。这个分工听起来合理可一旦把两者放到快权重视角下看边界就开始模糊了。问题的关键在于我们怎么看“权重”。传统理解里注意力没有显式权重矩阵它是对Value做加权平均权重来自query和key的相似度。但如果你把历史信息的累计结果看作一个随时间更新的矩阵那么注意力的输出本质上就是“一个矩阵乘上当前query变换后的向量”——这不就是FFN的线性层的操作吗1.2 快权重一个被冷落多年的老思路快权重的思想其实很老。最早可以追溯到20世纪90年代前后的一些联想记忆和快速权重模型的研究。思路非常简单普通神经网络的权重是通过梯度下降慢慢更新的这种权重可以叫“慢权重”而快权重是在前向推理过程中每看到一个输入就快速更新一次的临时权重。慢权重负责学会“怎么更新”快权重负责根据新输入临时调整行为。举个例子你就明白了。假设你有一个线性层y WxW是慢权重由训练得到。现在Wo是根据当前输入上下文实时变化的t时刻看到key和value后W就更新一下更新规则是外积累积。那么整个网络就既能保留慢权重学到的静态知识又能利用快权重对当前序列做出动态响应。线性注意力就是快权重思想的最佳代表。softmax注意力里输出是所有历史value的加权平均这个加权平均没法写成固定大小的状态。但一旦把softmax拆开让相似度函数变成可分解的形式历史信息就可以压缩成一个d×d的矩阵S。S随序列逐步更新当前输出由S决定。这个S就是前向传播过程中被快速更新的权重矩阵。2. 核心数学视角注意力怎么变成“梯度下降一步更新”2.1 线性注意力的标准推导过程先把标准注意力公式摆出来。对于第i个query输出是y_i Σ_j softmax(q_i · k_j) v_jsoftmax把相似度变成非负且和为1的权重。这个性质很好但正是softmax让整个计算无法并行成状态递推因为它要对所有j归一化。线性注意力的核心操作是把相似度函数从点积改成可分解的核函数形式。假设存在一个映射φ使得sim(q_i, k_j) φ(q_i)^T φ(k_j)那么注意力输出就可以改写y_i φ(q_i)^T Σ_j φ(k_j) v_j^T这里有个非常关键的结构变化。Σ_j φ(k_j) v_j^T 这一项不依赖i它只依赖前i个历史token。我们定义一个状态矩阵S_i Σ_{j≤i} φ(k_j) v_j^T于是最后一步就是y_i φ(q_i)^T S_i这个形式意味着什么它意味着你不再需要为每个query单独计算与所有key的相似度只需要维护一个不断更新的状态矩阵S。而S_i可以通过递推得到S_i S_{i-1} φ(k_i) v_i^T这一个公式就把注意力从O(N²d)的注意力矩阵计算变成了O(Nd²)的状态更新和O(d²)的查询计算。当序列长度远大于特征维度时复杂度直接降了一个量级。2.2 把递推看成权重更新快权重与联想记忆现在仔细看S的更新规则S ← S v φ(k)^T。这是一个外积累积操作。回想起神经网络里最基础的权重更新梯度下降更新W ← W - η (error · x^T)。两者是不是很像区别在于梯度下降里的更新来自反向传播误差而快权重里的更新来自当前输入的key和value映射。但数学形式上都是“当前状态加上一个外积”。这意味着注意力模块内部其实藏了一个“前馈式权重更新器”——它不做反向传播却每时每刻都在改自己的权重。把这个S矩阵当作FFN里的权重矩阵思路就通了。一个FFN层长这样h W2 σ(W1 x b1) b2如果把某个权重矩阵W换成随输入变化的S那这一层就同时具备了“前馈变换”和“上下文聚合”两种能力。于是“注意力也是FFN”这句话的准确含义是注意力可以看作一个权重由历史输入快速更新的FFN层。以前你是通过外部优化器比如Adam更新权重现在权重更新发生在前向计算过程中由序列里的内容驱动。这个统一视角最大的价值是让人意识到架构设计其实是在设计“慢权重”和“快权重”的分工。FFN里的静态矩阵负责稳定的知识变换注意力里的状态矩阵负责灵活的上下文响应。你把两者看作同一类操作的不同时间尺度比把它们当作两种性质完全不同的模块更接近本质。3. 高效注意力从O(N²)到O(N)的落地路径3.1 状态化之后训练与推理复杂度怎么变快权重视角直接给高效注意力指了条明路维护状态而不是维护完整的键值对。训练时虽然递推公式是逐步的但我们可以用矩阵运算一次性算出所有位置的状态。比如用前缀和cumsum或者三元矩阵掩码在GPU上并行计算所有S_i。这时的计算复杂度是O(Nd²)明显低于O(N²d)。推理时优势更明显。标准注意力需要缓存所有历史key和value也就是KV cache大小随序列长度线性增长。现在的状态化注意力只需要一个大小固定的S矩阵每来一个新token做一次外积更新和一次矩阵乘法查询即可。无论你生成1万token还是100万token缓存大小不变。我用一张表说明两者的差别对比项标准softmax注意力快权重/线性注意力训练计算复杂度O(N²d)O(Nd²)推理单步复杂度O(Nd)O(d²)缓存大小N×d×2KV cache随长度增长d×d状态矩阵固定长序列外推困难注意力分数分布漂移自然状态递推天然支持任意长度精确位置检索可以通过位置编码近似不行信息被压缩进固定状态这里要提醒一个细节。d²的复杂度看着优雅但d一旦变大就很吓人。d1024时S矩阵是100万参数d4096时是1600万参数。不少高效注意力工作会做低秩分解把d×d的状态拆成两个低秩矩阵乘积在不牺牲太多表达力的情况下降内存。3.2 与状态空间模型Mamba、RWKV等的关系这也解释了为什么状态空间模型SSM能火起来。SSM的递推形式是h_t A h_{t-1} B x_t, y_t C h_t这不也是快权重吗区别只在于SSM的状态更新是“旧的乘以常数矩阵A再加上新信息”而线性注意力的状态更新是“旧状态直接保留再加上外积”。SSM里的A如果变成单位矩阵状态就是一个累加器和线性注意力的S递推完全同构。两者的能力差异也源于此。SSM的A是一个固定的常数矩阵如果你希望做选择性遗忘需要加门控机制线性注意力的外积累积天然可以写一个衰减因子λS_t λ S_{t-1} φ(k_t) v_t^Tλ接近1代表长期记忆λ接近0代表短期记忆。这个灵活度比固定A矩阵的要高。当然SSM的优势在于它能学习系数矩阵来控制信息流动线性注意力则倾向让信息以无界方式累积所以经常会遇到数值稳定性上的麻烦。3.3 关键词总结把历史压进固定状态从更高的视角看所有高效注意力工作的核心动作都是“把历史压进固定大小的状态”。KV cache存的是全部历史状态化模型存的是历史的摘要。摘要怎么生成、怎么读出来决定了模型的表达能力。快权重这个视角给出的答案是摘要是一个矩阵它被当前输入不断外积更新读取方式是让当前query和这个矩阵做内积。这个回答简洁、统一而且可解释性强——你可以把S的每一行理解成一个“记忆插槽”不同的key会把信息写进对应的插槽query再从插槽里取信息。4. 实操用PyTorch从零实现一个快权重注意力层4.1 最小可运行实现理论再多不如跑通一个最小实现。下面这个代码是我实验里的简化版舍弃了batch和head维度专注展示核心逻辑。import torch import torch.nn as nn import torch.nn.functional as F def feature_map(x): # 常见的核函数elu 1输出恒正且带一定非线性 return F.elu(x) 1.0 class FastWeightAttention(nn.Module): def __init__(self, d_model, d_kNone, d_vNone): super().__init__() d_k d_k or d_model d_v d_v or d_model self.d_k d_k self.d_v d_v # 慢权重负责把输入映射成query/key/value self.wq nn.Linear(d_model, d_k, biasFalse) self.wk nn.Linear(d_model, d_k, biasFalse) self.wv nn.Linear(d_model, d_v, biasFalse) # 输出投影 self.out_proj nn.Linear(d_v, d_model, biasFalse) # 可选的衰减因子越接近1记忆越长 self.lambda_decay nn.Parameter(torch.tensor(0.999)) def forward_train(self, x): # x: [N, d_model] N, _ x.shape q feature_map(self.wq(x)) # [N, d_k] k feature_map(self.wk(x)) # [N, d_k] v self.wv(x) # [N, d_v] decay torch.sigmoid(self.lambda_decay).clamp(0.9, 1.0) # 用前缀和累积所有历史S_i decay * S_{i-1} k_i v_i^T # 关键技巧用cumsum实现外积的前缀和 # 先把外积展开k_i v_i^T 是 [d_k, d_v]对N维做cumsum outer torch.einsum(n,d,nv-ndv, k, v.unsqueeze(0).squeeze(0), v) # 这一行先不展开 # 正确写法k[:, :, None] * v[:, None, :] 得到 [N, d_k, d_v] outer k[:, :, None] * v[:, None, :] # [N, d_k, d_v] # 累积先按时间累加再乘衰减 S torch.cumsum(outer, dim0) # 如果要用衰减需要逐位置计算几何累加这里为了可读性先忽略decay # 查询y_i q_i^T S_i # [N, d_k] x [N, d_k, d_v] - [N, d_v] y torch.einsum(nd,ndv-nv, q, S) return self.out_proj(y) def forward_step(self, S, x): # 推理模式逐步更新 q feature_map(self.wq(x)) k feature_map(self.wk(x)) v self.wv(x) S S torch.outer(k, v) # 快权重更新 y q S # 查询 return self.out_proj(y), S def init_state(self, batch_size1): return torch.zeros(batch_size, self.d_k, self.d_v)这里要说明两点。训练版代码里我用了cumsum直接算前缀和这在直观理解上没问题但实际工程里逐token的衰减因子会让它失效因为每个位置的历史都乘了不同次数的λ。要用好衰减一般得用分段累积或并行扫描算法parallel scan。上面代码为了展示核心逻辑省略了decay的精确实现正式实验记得补上。还有一点推理时S要保存成float32不能图省事用float16。外积累积对精度非常敏感低精度下累积误差会随着序列长度越滚越大。4.2 数值稳定性为什么必须维护归一化状态标准softmax注意力自带归一化输出始终是value的凸组合数值范围受控。线性注意力去掉了softmax问题立刻就来了S在无穷序列上无界累积输出可能指数增长训练到一半直接变成NaN。解决办法是维护一个额外的归一化项z记录状态S的累计“能量”z_i z_{i-1} φ(k_i)然后用z对输出做除法y_i (φ(q_i)^T S_i) / (φ(q_i)^T z_i)这个除法的含义是用当前query和历史key的匹配程度对状态做加权平均而不是直接累加。在实际实现里分母还要加一个eps防止除零。另外还有一个小技巧可以把λ衰减因子加到归一化项上让旧信息的贡献同步衰减防止模型总是偏向近期token。我个人的经验是λ初始设成0.999让它慢慢可学习比固定常数要好调得多。4.3 对标实验与消融的注意点从零实现之后最该做的是在一个可控任务上和softmax注意力做对照。我建议不要一上来就冲大规模语言模型一个小规模的语言建模任务或者图像分类任务足够暴露问题。第一控制变量。用完全相同的参数总量、学习率、训练步数、batch大小只换注意力模块。这里有个坑线性注意力的训练动态和softmax不一样初始loss下降往往更快因为它没有softmax那种“分散”的相似度分布梯度信号更直接。你要是只看前1000步容易被骗。第二建议分别测三个指标训练loss、验证loss、推理时的KV内存。我见过不少项目验证loss掉了一些但KV cache内存完全没有参考意义因为根本没有cache了。这部分收益反而是最实在的。第三消融实验至少要做三组标准softmax注意力、线性注意力、带衰减因子的线性注意力。这样你能清楚看到长程记忆能力有多少是来自外积累积有多少是来自衰减机制。5. 常见问题与调试实录5.1 状态爆炸loss变NaN怎么查这是我遇到最多的问题现象是训练进行到几千步的时候loss突然变成NaN或者注意力输出出现几百上千的极端值。排查思路分三路。第一路检查φ映射。如果直接用ReLU或者elu1输出范围还算可控。如果你图省事用线性映射那相似度可能变成负值外积累积后状态内部互相抵消或放大问题就来了。第二路检查decay因子。当衰减因子接近1状态累积速度非常快。长序列训练时1万token的外积累积量级会达到初始值的几千倍。把λ设成0.99开始训练数值会稳很多。第三路检查归一化除法。很多线性注意力实现只除一个标量或者干脆不除这是最脆弱的版本。务必加上z归一化同时在分母加1e-6量级的eps。一个我实际用过的保底方案在输出y后面加一个LayerNorm。这让数值稳定问题不会直接体现在loss上模型至少能训练起来。代价是会略微改变注意力输出的分布有时对效果有影响但至少不会崩。5.2 训练速度远比预期慢线性注意力复杂度看着低但如果实现方式不对训练可能比softmax还慢。最常见的原因是为了求前缀和有人真的在PyTorch里写了个for循环逐步递推这等于把并行化全部丢掉了GPU利用率接近02万token序列训练一次要跑几千步。正确做法是矩阵化。同一段序列把所有外积一次性算出来然后用torch.cumsum沿时间维度累积。如果用了衰减因子就得用分段扫描算法或者把序列切成块块内并行、块间接续。这已经是工程优化层面的事了但对训练体验影响极大。我在实现时还犯过一个低级错误torch.einsum里的维度顺序写错导致输出维度变成[N, d_v, d_k]而不是[N, d_v]不但算得慢结果还全是乱的。调试这种问题最好的办法是用一个极短的序列比如长度2、d4做手算核对不要直接跑大数据。5.3 什么时候会失效快权重注意力不是万能的。在需要精确位置检索的任务上比如“记住第7个token是什么”或者需要明确区分相似上下文的场景固定大小的状态矩阵很容易发生信息混淆。KV cache是精确存储而状态化模型是压缩存储压缩必然带来损失。另一个失效场景是短序列生成。序列只有64到128个token时标准softmax注意力的KV cache不过几千个数值内存和计算都很轻状态化模型反而要维护一个d×d矩阵算起来更重。所以这类模型最适合的场合是超长序列、流式生成、端侧设备上长文本处理——这些地方缓存大小才是真正的瓶颈。我自己判断一个任务适不适合快权重注意力只看一个问题信息密度和序列长度的比例。如果任务要求从100万token里精确找到一个词状态化模型很难做好。如果任务只需要捕捉整体统计规律和长程趋势那就非常适合。6. 一点个人实操体会最后分享几个我在实践里沉淀下来的经验。第一个体会是把S当作FFN权重去初始化会非常有帮助。既然注意力也是FFN那状态矩阵的初始化就不能全零。全零初始化的意思是模型一开始完全依赖当前query没有历史信息这对训练非常不利。我在某些任务里把S初始化为单位矩阵或者近单位矩阵相当于给模型一个“默认历史”训练稳定性和收敛速度都有改善。第二个体会是真正的收益不在单独换掉注意力而在整个网络的结构配合。快权重注意力把信息压缩进状态之后后面的FFN层承担了更多“解压”工作。这时调大FFN的中间维度往往比调注意力本身的参数更有效。很多论文里“等价参数下效果接近softmax”的结论背后都做了这种结构调整。第三个体会也是我最近比较深的感受这个方向真正的价值不是把注意力干掉而是让“注意力”这个概念本身变得可以设计和组合。你可以在慢权重和快权重之间调节记忆强度可以决定状态里存什么丢什么可以把状态矩阵和FFN权重做耦合。这种自由度比在softmax框架里打转有意思得多。如果你也想试试这个方向我建议先从一个小规模的语言建模任务开始跑通上面的代码确认loss能正常下降再逐步加长序列、加大模型。调试快权重注意力最大的坑是数值稳定性只要把归一化、衰减、初始化这三件事做对剩下的就是纯工程和专业细节问题了。
阅读完成 · 觉得有帮助？
