人工智能深度学习大模型NLP计算机视觉强化学习LoRA【免费下载链接】annotated_deep_learning_paper_implementations 60 Implementations/tutorials of deep learning papers with side-by-side notes ; including transformers (original, xl, switch, feedback, vit, ...), optimizers (adam, adabelief, sophia, ...), gans(cyclegan, stylegan2, ...), reinforcement learning (ppo, dqn), capsnet, distillation, ... 项目地址https://gitcode.com/gh_mirrors/an/annotated_deep_learning_paper_implementations点击查看免费下载本文是 annotated_deep_learning_paper_implementations 仓库中 Fast Weights Transformer 模块的技术指南对应论文《Linear Transformers Are Secretly Fast Weight Memory Systems in PyTorch》arXiv:2102.11174。文章以仓库内的 labml_nn/transformers/fast_weights/init.py 注释实现为主体结合 训练代码 与 Tiny Shakespeare 数据集上的自回归实验逐层讲解快速权重模型、线性自注意力核替换、DPFP 投影函数与新的权重更新规则并给出可直接运行的训练配置。读完本文你将能够理解快速权重记忆系统与线性注意力之间的等价关系掌握在 PyTorch 中实现并训练一个 Fast Weights Transformer 的完整方案。背景快速权重系统与 Transformer 自注意力的关联快速权重Fast Weights模型考虑一个长度为 L 的输入序列 ${x^{(i)}}^L_{i1}$每个时间步的输入是大小为 $d_{in}$ 的向量 $x \in \mathbb{R}^{d_{in}}$。快速权重模型在每个时间步生成一个权重矩阵用来产生输出 ${y^{(i)}}^L_{i1}$$y \in \mathbb{R}^{d_{out}}$其递推过程为$$ \begin{aligned} a^{(i)}, b^{(i)} W_a x^{(i)},\ W_b x^{(i)} \ W^{(i)} \sigma \Big( W^{(i-1)} a^{(i)} \otimes b^{(i)} \Big) \ y^{(i)} W^{(i)} x^{(i)} \end{aligned} $$其中 $\otimes$ 是外积$a \otimes b a b^\top$即两个向量的元素两两相乘得到一个矩阵$\sigma$ 是激活函数$W_a$ 与 $W_b$ 是可训练参数橙色标记$W^{(i)}$ 就是每一步动态生成的快速权重青色标记。核心思想是权重不是一次学习得到的静态参数而是随输入序列逐部递推更新的记忆。线性自注意力的快速计算形式标准 Transformer 自注意力省略 $\frac{1}{d_k}$ 缩放因子可以写成$$ y^{(i)} \sum^i_{j1} \frac{ v^{(j)} \kappa(k^{(j)}, q^{(i)}) }{ \sum^i_{j1} \kappa(k^{(j)}, q^{(i)}) }, \quad \kappa(k, q) \exp(k \cdot q) $$线性化自注意力的核心思想是把 softmax 核 $\kappa$ 替换成另一个核 $\kappa$使得分母可以更快地计算$$ \kappa(k, q) \phi(k)^\top \phi(q) $$代入后得到$$ y^{(i)} \frac{ \Big( \sum^i_{j1} v^{(j)} \otimes \phi(k^{(j)}) \Big) \phi(q^{(i)}) }{ \Big( \sum^i_{j1} \phi(k^{(j)}) \Big) \phi(q^{(i)}) } $$如果定义 $W^{(i)} \sum^i_{j1} v^{(j)} \otimes \phi(k^{(j)})$ 与 $z^{(i)} \sum^i_{j1} \phi(k^{(j)})$就可以递推计算$$ \begin{aligned} W^{(i)} W^{(i-1)} v^{(i)} \otimes \phi(k^{(i)}) \ z^{(i)} z^{(i-1)} \phi(k^{(i)}) \ y^{(i)} \frac{1}{z^{(i)} \cdot \phi(q^{(i)})} W^{(i)} \phi(q^{(i)}) \end{aligned} $$对比快速权重模型的更新式 $W^{(i)} \sigma(W^{(i-1)} a^{(i)} \otimes b^{(i)})$可以看到两者高度相似——这正是论文标题Transformer 其实是快速权重记忆系统的由来。基于这一观察论文对自注意力做了三处修改引入新的线性注意力投影函数 $\phi$DPFP、为 $W^{(i)} f(W^{(i-1)})$ 设计新的更新规则、并改进归一化项 $\frac{1}{z^{(i)} \cdot \phi(q^{(i)})}$。DPFP无参数的确定性投影函数投影公式与超参数论文引入的投影函数称为Deterministic Parameter Free ProjectionDPFP实现于 labml_nn/transformers/fast_weights/init.py 中的DPFP类。它把维度为 $d_{key}$ 的向量 $k$ 投影到维度 $d_{dot} 2 d_{key} \nu$其中 $\nu \in {1, 2, \dots, 2 d_{key} - 1}$ 是超参数。投影定义为$$ \phi_{2 d_{key} (i - 1) j}(k) \text{ReLU}\Big(\big[k, -k\big]\Big){j} \cdot \text{ReLU}\Big(\big[k, -k\big]\Big){i j} $$其中 $\big[k, -k\big]$ 是 $k$ 与 $-k$ 的拼接得到一个大小为 $2 d_{key}$ 的向量$i \in {1, 2, \dots, \nu}$$j \in {1, 2, \dots, 2 d_{key}}$下标超出向量长度时采用循环回绕roll取值。简单来说DPFP 通过把 $[k, -k]$ 的元素按位移 $i$ 相乘来构造新向量全程不含任何可学习参数。这种构造方式带来两个性质稀疏性$\phi$ 中只有少数元素非零近似正交性$\phi(k^{(i)}) \cdot \phi(k^{(j)}) \approx 0$ 对大多数 $i, j$ 成立除非 $k^{(i)}$ 与 $k^{(j)}$ 非常相似。归一化论文还为 $\phi$ 引入了一个简单的归一化推导细节见论文$$ \phi(k) \frac{\phi(k)}{\sum^{d_{dot}}_{j1} \phi(k)_j} $$对应源码实现如下def forward(self, k: torch.Tensor): # 计算 φ(k) k self.dpfp(k) # 按 ∑ φ(k)_j 归一化eps 防止除零 return k / (torch.sum(k, dim-1, keepdimTrue) self.eps) def dpfp(self, k: torch.Tensor): # x ReLU([k, -k]) x self.relu(torch.cat([k, -k], dim-1)) # 按 i ∈ {1, 2, ..., ν} 位移并回绕 x_rolled [x.roll(shiftsi, dims-1) for i in range(1, self.nu 1)] x_rolled torch.cat(x_rolled, dim-1) # 复制 x 拼接与位移版本逐元素相乘 x_repeat torch.cat([x] * self.nu, dim-1) return x_repeat * x_rolled构造函数DPFP(nu: int 1, eps: float 1e-6)中nu超参数 $\nu$控制投影维度扩张倍数默认1eps归一化时避免除零的小常数默认1e-6。实现上利用了 PyTorch 的torch.roll(shiftsi, dims-1)完成循环位移与论文中下标超出即回绕的定义完全对应。FastWeightsAttention新的权重更新规则更新公式FastWeightsAttention同样定义于init.py实现了论文提出的新更新规则。模型首先用快速权重检索当前 key $k^{(i)}$ 对应的已有值 $\bar{v}^{(i)}$然后把检索值与输入 $v^{(i)}$ 做插值组合再写回权重矩阵$$ \begin{aligned} k^{(i)}, v^{(i)}, q^{(i)} W_k x^{(i)},\ W_v x^{(i)},\ W_q x^{(i)} \ \bar{v}^{(i)} W^{(i-1)} \phi(k^{(i)}) \ \beta^{(i)} \sigma \Big(W_\beta x^{(i)} \Big) \ v^{(i)}{new} \beta^{(i)} v^{(i)} \Big(1 - \beta^{(i)} \Big) \bar{v}^{(i)} \ W^{(i)} W^{(i-1)} v^{(i)}{new} \otimes \phi(k^{(i)}) \ W^{(i-1)} \beta^{(i)} \Big( v^{(i)} - \bar{v}^{(i)} \Big ) \otimes \phi(k^{(i)}) \ y^{(i)} W^{(i)} \phi(q^{(i)}) \end{aligned} $$其中 $W_\beta$ 是可训练参数$\sigma$ 是 sigmoid 函数。注意由于 $\phi$ 本身已经归一化这里不再需要线性注意力中的归一化项 $z$。模块结构与源码解读FastWeightsAttention的构造参数为(heads, d_model, dropout_prob, phi)self.d_k d_model // heads # 每个头的特征数 self.heads heads self.query PrepareForMultiHeadAttention(d_model, heads, self.d_k, biasFalse) self.key PrepareForMultiHeadAttention(d_model, heads, self.d_k, biasFalse) self.value PrepareForMultiHeadAttention(d_model, heads, self.d_k, biasFalse) self.interpolation_weight nn.Sequential( PrepareForMultiHeadAttention(d_model, heads, 1, biasFalse), nn.Sigmoid() ) self.phi phi self.output nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout_prob)各组件作用query/key/value三个PrepareForMultiHeadAttention线性变换见 labml_nn/transformers/mha.py其内部是nn.Linear(d_model, heads * d_k)后按头拆分这里关闭 bias分别投影出 $\phi(q^{(i)})$、$\phi(k^{(i)})$、$v^{(i)}$interpolation_weight插值权重网络 $\sigma(W_\beta x^{(i)})$即 $\beta^{(i)}$ 的生成器输出维度为每个头一个标量phi传入的DPFP实例output输出层把多头拼接结果映射回 $d_{model}$dropout对注意力输出做随机失活。前向过程按序列步长逐步递推对应文档第 230-267 行源码def forward(self, x: torch.Tensor): seq_len x.shape[0] query self.phi(self.query(x)) # φ(q^(i))所有步与头 key self.phi(self.key(x)) # φ(k^(i))所有步与头 value self.value(x) # v^(i)所有步与头 beta self.interpolation_weight(x) # β^(i)所有步与头 # W^(0) 初始化为零矩阵形状 [batch, heads, d_value, d_key] weights key.new_zeros((key.shape[1], key.shape[2], value.shape[3], key.shape[3])) outputs [] for i in range(seq_len): # 检索v̄^(i) W^(i-1) φ(k^(i)) value_existing torch.einsum(bhvk,bhk-bhv, weights, key[i]) # 写回W^(i) W^(i-1) β^(i)(v^(i) − v̄^(i)) ⊗ φ(k^(i)) weights weights torch.einsum(bhv,bhk-bhvk, beta[i] * (value[i] - value_existing), key[i]) # 读取y^(i) W^(i) φ(q^(i)) y torch.einsum(bhvk,bhk-bhv, weights, query[i]) # 合并多头 outputs.append(y.reshape(y.shape[0], -1)) x torch.stack(outputs) return self.output(x)每一步都对应公式中的一个操作先用einsum(bhvk,bhk-bhv, ...)完成矩阵-向量检索得到 $\bar{v}^{(i)}$再用外积einsum(bhv,bhk-bhvk, ...)将插值残差 $\beta^{(i)}(v^{(i)} - \bar{v}^{(i)})$ 写回权重矩阵最后用更新后的权重读取输出。这样权重矩阵 $W^{(i)}$ 就在序列推进过程中持续演化充当了记忆。token_wise 变体逐 token 推理仓库还在 labml_nn/transformers/fast_weights/token_wise.py 提供了逐 token 处理的变体FastWeightsAttentionTransformer。它的forward先用torch.unbind(x_seq, dim0)把序列拆成单个 token然后为每一层维护一份权重列表weights [None for _ in range(len(self.layers))]每个 token 依次穿过所有层、逐层更新对应权重最后torch.stack(res)并做 LayerNorm。这种写法与自回归推理时一次只喂一个 token、权重跨步长保留的使用方式天然契合。组装 Transformer层与整体模型FastWeightsAttentionTransformerLayerFastWeightsAttentionTransformerLayerinit.py 第 270 行起是一个标准 Transformer 层组合了快速权重自注意力与前馈网络self.size d_model self.attn attn # FastWeightsAttention 实例 self.feed_forward feed_forward self.dropout nn.Dropout(dropout_prob) self.norm_self_attn nn.LayerNorm([d_model]) self.norm_ff nn.LayerNorm([d_model]) def forward(self, x: torch.Tensor): attn self.attn(x) x x self.dropout(attn) # 自注意力残差 z self.norm_ff(x) # 前馈前归一化 ff self.feed_forward(z) x x self.dropout(ff) # 前馈残差 return x结构上与常规 Transformer 层一致自注意力残差连接 → LayerNorm → 前馈网络 → 残差连接只是把 MHA 换成了FastWeightsAttention。前馈网络复用 labml_nn/transformers/feed_forward.py 的FeedForward(d_model, d_ff, dropout)两层线性 激活 dropout默认 ReLU。FastWeightsAttentionTransformerFastWeightsAttentionTransformer负责把单层复制成多层并接最终归一化def __init__(self, layer, n_layers): self.layers clone_module_list(layer, n_layers) # 深拷贝 n_layers 份 self.norm nn.LayerNorm([layer.size]) def forward(self, x): for layer in self.layers: x layer(x) return self.norm(x)层复制使用了仓库工具 labml_nn/utils/init.py 中的clone_module_list(module, n)即nn.ModuleList([copy.deepcopy(module) for _ in range(n)])确保每层拥有独立参数。训练实验Tiny Shakespeare 上的自回归语言模型模型与默认配置训练入口位于 labml_nn/transformers/fast_weights/experiment.py。AutoregressiveModel由三部分组成nn.Embedding(n_vocab, d_model)词嵌入、Fast Weights Transformer 主干、nn.Linear(d_model, n_vocab)输出层前向输出下一个 token 的 logits。Configs继承自 labml_nn/experiments/nlp_autoregression.py 的NLPAutoRegressionConfigs核心默认参数如下参数默认值说明d_model512模型维度nu1DPFP 投影超参数 $\nu$heads8注意力头数d_k d_model // heads 64dropout0.0丢弃概率该实验关闭 dropoutd_ff2048前馈网络隐藏维度n_layers6Transformer 层数模型构造逻辑fast_weights_transformer(c)把上述配置拼装为AutoregressiveModel( c.n_tokens, c.d_model, FastWeightsAttentionTransformer( FastWeightsAttentionTransformerLayer( d_modelc.d_model, attnFastWeightsAttention(c.heads, c.d_model, c.dropout, DPFP(nuc.nu)), feed_forwardFeedForward(c.d_model, c.d_ff, c.dropout), dropout_probc.dropout), c.n_layers))启动训练main()使用 labml 实验框架创建并运行实验覆盖的关键配置包括experiment.configs(conf, { tokenizer: character, # 字符级分词 text: tiny_shakespeare, # Tiny Shakespeare 数据集 optimizer.learning_rate: 1.0, optimizer.optimizer: Noam, # Noam 学习率调度优化器 prompt: It is, # 采样起始提示 prompt_separator: , train_loader: shuffled_train_loader, valid_loader: shuffled_valid_loader, seq_len: 128, # 上下文长度 epochs: 128, batch_size: 16, inner_iterations: 25})运行时执行python labml_nn/transformers/fast_weights/experiment.py即可开始训练仓库同时提供了 experiment.ipynb 笔记可在 Colab 中直接训练同一个模型。训练过程中使用交叉熵损失与准确率指标并支持梯度裁剪grad_norm_clip 1.0见NLPAutoRegressionConfigs。推理与采样得益于快速权重机制模型在生成阶段可以天然地记忆更长的历史token_wise.py中每个 token 输入时各层快速权重矩阵在上一步的基础上继续更新无需重新计算整段历史的注意力。训练时用prompt启动采样配合prompt_separator字符级分词时为空字符串控制提示的拼接方式。源码地图与延伸阅读核心实现labml_nn/transformers/fast_weights/init.pyDPFP、FastWeightsAttention、FastWeightsAttentionTransformerLayer、FastWeightsAttentionTransformer逐 token 变体labml_nn/transformers/fast_weights/token_wise.py训练脚本与笔记labml_nn/transformers/fast_weights/experiment.py、labml_nn/transformers/fast_weights/experiment.ipynb依赖组件多头准备层 labml_nn/transformers/mha.py、前馈网络 labml_nn/transformers/feed_forward.py、模块克隆工具 labml_nn/utils/init.py、自回归训练基类 labml_nn/experiments/nlp_autoregression.py在线文档仓库 docs 目录下的 docs/transformers/fast_weights/index.html、docs/transformers/fast_weights/token_wise.html 与 docs/transformers/fast_weights/experiment.html 与本文章节一一对应总结Fast Weights Transformer 的核心洞察在于线性自注意力中逐步累积的 $W^{(i)} \sum v^{(j)} \otimes \phi(k^{(j)})$ 本质上就是一个快速权重记忆系统。论文及其在本仓库的注释实现通过三点改造提升了该系统的表达能力一是用无参数的 DPFP 投影 $\phi$ 替代 softmax 核得到稀疏、近似正交且自带归一化的投影二是把朴素的直接写回更新规则改为先检索、再插值、后写回用可学习的 $\beta^{(i)}$ 门控控制新信息与旧记忆的融合三是利用 $\phi$ 的归一化省去了额外的归一化项 $z$。从init.py 的每一步 einsum 到 experiment.py 的完整训练配置读者可以沿着本文的路径复现并扩展这个将记忆直接编码进权重演化的高效 Transformer 变体。赞分享人工智能深度学习大模型NLP计算机视觉强化学习LoRA【免费下载链接】annotated_deep_learning_paper_implementations 60 Implementations/tutorials of deep learning papers with side-by-side notes ; including transformers (original, xl, switch, feedback, vit, ...), optimizers (adam, adabelief, sophia, ...), gans(cyclegan, stylegan2, ...), reinforcement learning (ppo, dqn), capsnet, distillation, ... 项目地址https://gitcode.com/gh_mirrors/an/annotated_deep_learning_paper_implementations点击查看免费下载相关推荐黑苹果EFI配置终极指南OpCore Simplify工具快速入门教程黑苹果EFI配置终极指南OpCore Simplify工具快速入门教程 还在为黑苹果复杂的EFI配置而头疼吗OpCore Simplify是一款专门为简化O开发工具CLINYU-DLSP20 注意力机制与 Transformer从自注意力数学推导到 PyTorch 完整实现NYU DLSP20 注意力机制与 Transformer从自注意力数学推导到 PyTorch 完整实现 本文基于 NYU Deep Learning Spr示例工程深入理解视觉Transformer窗口注意力pytorch-image-models中的高效注意力机制深入理解视觉Transformer窗口注意力pytorch image models中的高效注意力机制 视觉Transformer窗口注意力 是计算机视觉领域人工智能计算机视觉深度学习预训练上一篇AutoClip 投稿状态页面使用指南B 站投稿任务的监控、管理与问题排查下一篇从晶体管到 CPU数字电路基础与抽象层次全解析Easy-Vibe 计算机基础篇创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/10/10 6:07:11 · 阅读 1.2k 深入解析 Fast Weights Transformer：线性自注意力与快速权重记忆系统的 PyTorch 实现
人工智能深度学习大模型NLP计算机视觉强化学习LoRA【免费下载链接】annotated_deep_learning_paper_implementations 60 Implementations/tutorials of deep learning papers with side-by-side notes ; including transformers (original, xl, switch, feedback, vit, ...), optimizers (adam, adabelief, sophia, ...), gans(cyclegan, stylegan2, ...), reinforcement learning (ppo, dqn), capsnet, distillation, ... 项目地址https://gitcode.com/gh_mirrors/an/annotated_deep_learning_paper_implementations点击查看免费下载本文是 annotated_deep_learning_paper_implementations 仓库中 Fast Weights Transformer 模块的技术指南对应论文《Linear Transformers Are Secretly Fast Weight Memory Systems in PyTorch》arXiv:2102.11174。文章以仓库内的 labml_nn/transformers/fast_weights/init.py 注释实现为主体结合 训练代码 与 Tiny Shakespeare 数据集上的自回归实验逐层讲解快速权重模型、线性自注意力核替换、DPFP 投影函数与新的权重更新规则并给出可直接运行的训练配置。读完本文你将能够理解快速权重记忆系统与线性注意力之间的等价关系掌握在 PyTorch 中实现并训练一个 Fast Weights Transformer 的完整方案。背景快速权重系统与 Transformer 自注意力的关联快速权重Fast Weights模型考虑一个长度为 L 的输入序列 ${x^{(i)}}^L_{i1}$每个时间步的输入是大小为 $d_{in}$ 的向量 $x \in \mathbb{R}^{d_{in}}$。快速权重模型在每个时间步生成一个权重矩阵用来产生输出 ${y^{(i)}}^L_{i1}$$y \in \mathbb{R}^{d_{out}}$其递推过程为$$ \begin{aligned} a^{(i)}, b^{(i)} W_a x^{(i)},\ W_b x^{(i)} \ W^{(i)} \sigma \Big( W^{(i-1)} a^{(i)} \otimes b^{(i)} \Big) \ y^{(i)} W^{(i)} x^{(i)} \end{aligned} $$其中 $\otimes$ 是外积$a \otimes b a b^\top$即两个向量的元素两两相乘得到一个矩阵$\sigma$ 是激活函数$W_a$ 与 $W_b$ 是可训练参数橙色标记$W^{(i)}$ 就是每一步动态生成的快速权重青色标记。核心思想是权重不是一次学习得到的静态参数而是随输入序列逐部递推更新的记忆。线性自注意力的快速计算形式标准 Transformer 自注意力省略 $\frac{1}{d_k}$ 缩放因子可以写成$$ y^{(i)} \sum^i_{j1} \frac{ v^{(j)} \kappa(k^{(j)}, q^{(i)}) }{ \sum^i_{j1} \kappa(k^{(j)}, q^{(i)}) }, \quad \kappa(k, q) \exp(k \cdot q) $$线性化自注意力的核心思想是把 softmax 核 $\kappa$ 替换成另一个核 $\kappa$使得分母可以更快地计算$$ \kappa(k, q) \phi(k)^\top \phi(q) $$代入后得到$$ y^{(i)} \frac{ \Big( \sum^i_{j1} v^{(j)} \otimes \phi(k^{(j)}) \Big) \phi(q^{(i)}) }{ \Big( \sum^i_{j1} \phi(k^{(j)}) \Big) \phi(q^{(i)}) } $$如果定义 $W^{(i)} \sum^i_{j1} v^{(j)} \otimes \phi(k^{(j)})$ 与 $z^{(i)} \sum^i_{j1} \phi(k^{(j)})$就可以递推计算$$ \begin{aligned} W^{(i)} W^{(i-1)} v^{(i)} \otimes \phi(k^{(i)}) \ z^{(i)} z^{(i-1)} \phi(k^{(i)}) \ y^{(i)} \frac{1}{z^{(i)} \cdot \phi(q^{(i)})} W^{(i)} \phi(q^{(i)}) \end{aligned} $$对比快速权重模型的更新式 $W^{(i)} \sigma(W^{(i-1)} a^{(i)} \otimes b^{(i)})$可以看到两者高度相似——这正是论文标题Transformer 其实是快速权重记忆系统的由来。基于这一观察论文对自注意力做了三处修改引入新的线性注意力投影函数 $\phi$DPFP、为 $W^{(i)} f(W^{(i-1)})$ 设计新的更新规则、并改进归一化项 $\frac{1}{z^{(i)} \cdot \phi(q^{(i)})}$。DPFP无参数的确定性投影函数投影公式与超参数论文引入的投影函数称为Deterministic Parameter Free ProjectionDPFP实现于 labml_nn/transformers/fast_weights/init.py 中的DPFP类。它把维度为 $d_{key}$ 的向量 $k$ 投影到维度 $d_{dot} 2 d_{key} \nu$其中 $\nu \in {1, 2, \dots, 2 d_{key} - 1}$ 是超参数。投影定义为$$ \phi_{2 d_{key} (i - 1) j}(k) \text{ReLU}\Big(\big[k, -k\big]\Big){j} \cdot \text{ReLU}\Big(\big[k, -k\big]\Big){i j} $$其中 $\big[k, -k\big]$ 是 $k$ 与 $-k$ 的拼接得到一个大小为 $2 d_{key}$ 的向量$i \in {1, 2, \dots, \nu}$$j \in {1, 2, \dots, 2 d_{key}}$下标超出向量长度时采用循环回绕roll取值。简单来说DPFP 通过把 $[k, -k]$ 的元素按位移 $i$ 相乘来构造新向量全程不含任何可学习参数。这种构造方式带来两个性质稀疏性$\phi$ 中只有少数元素非零近似正交性$\phi(k^{(i)}) \cdot \phi(k^{(j)}) \approx 0$ 对大多数 $i, j$ 成立除非 $k^{(i)}$ 与 $k^{(j)}$ 非常相似。归一化论文还为 $\phi$ 引入了一个简单的归一化推导细节见论文$$ \phi(k) \frac{\phi(k)}{\sum^{d_{dot}}_{j1} \phi(k)_j} $$对应源码实现如下def forward(self, k: torch.Tensor): # 计算 φ(k) k self.dpfp(k) # 按 ∑ φ(k)_j 归一化eps 防止除零 return k / (torch.sum(k, dim-1, keepdimTrue) self.eps) def dpfp(self, k: torch.Tensor): # x ReLU([k, -k]) x self.relu(torch.cat([k, -k], dim-1)) # 按 i ∈ {1, 2, ..., ν} 位移并回绕 x_rolled [x.roll(shiftsi, dims-1) for i in range(1, self.nu 1)] x_rolled torch.cat(x_rolled, dim-1) # 复制 x 拼接与位移版本逐元素相乘 x_repeat torch.cat([x] * self.nu, dim-1) return x_repeat * x_rolled构造函数DPFP(nu: int 1, eps: float 1e-6)中nu超参数 $\nu$控制投影维度扩张倍数默认1eps归一化时避免除零的小常数默认1e-6。实现上利用了 PyTorch 的torch.roll(shiftsi, dims-1)完成循环位移与论文中下标超出即回绕的定义完全对应。FastWeightsAttention新的权重更新规则更新公式FastWeightsAttention同样定义于init.py实现了论文提出的新更新规则。模型首先用快速权重检索当前 key $k^{(i)}$ 对应的已有值 $\bar{v}^{(i)}$然后把检索值与输入 $v^{(i)}$ 做插值组合再写回权重矩阵$$ \begin{aligned} k^{(i)}, v^{(i)}, q^{(i)} W_k x^{(i)},\ W_v x^{(i)},\ W_q x^{(i)} \ \bar{v}^{(i)} W^{(i-1)} \phi(k^{(i)}) \ \beta^{(i)} \sigma \Big(W_\beta x^{(i)} \Big) \ v^{(i)}{new} \beta^{(i)} v^{(i)} \Big(1 - \beta^{(i)} \Big) \bar{v}^{(i)} \ W^{(i)} W^{(i-1)} v^{(i)}{new} \otimes \phi(k^{(i)}) \ W^{(i-1)} \beta^{(i)} \Big( v^{(i)} - \bar{v}^{(i)} \Big ) \otimes \phi(k^{(i)}) \ y^{(i)} W^{(i)} \phi(q^{(i)}) \end{aligned} $$其中 $W_\beta$ 是可训练参数$\sigma$ 是 sigmoid 函数。注意由于 $\phi$ 本身已经归一化这里不再需要线性注意力中的归一化项 $z$。模块结构与源码解读FastWeightsAttention的构造参数为(heads, d_model, dropout_prob, phi)self.d_k d_model // heads # 每个头的特征数 self.heads heads self.query PrepareForMultiHeadAttention(d_model, heads, self.d_k, biasFalse) self.key PrepareForMultiHeadAttention(d_model, heads, self.d_k, biasFalse) self.value PrepareForMultiHeadAttention(d_model, heads, self.d_k, biasFalse) self.interpolation_weight nn.Sequential( PrepareForMultiHeadAttention(d_model, heads, 1, biasFalse), nn.Sigmoid() ) self.phi phi self.output nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout_prob)各组件作用query/key/value三个PrepareForMultiHeadAttention线性变换见 labml_nn/transformers/mha.py其内部是nn.Linear(d_model, heads * d_k)后按头拆分这里关闭 bias分别投影出 $\phi(q^{(i)})$、$\phi(k^{(i)})$、$v^{(i)}$interpolation_weight插值权重网络 $\sigma(W_\beta x^{(i)})$即 $\beta^{(i)}$ 的生成器输出维度为每个头一个标量phi传入的DPFP实例output输出层把多头拼接结果映射回 $d_{model}$dropout对注意力输出做随机失活。前向过程按序列步长逐步递推对应文档第 230-267 行源码def forward(self, x: torch.Tensor): seq_len x.shape[0] query self.phi(self.query(x)) # φ(q^(i))所有步与头 key self.phi(self.key(x)) # φ(k^(i))所有步与头 value self.value(x) # v^(i)所有步与头 beta self.interpolation_weight(x) # β^(i)所有步与头 # W^(0) 初始化为零矩阵形状 [batch, heads, d_value, d_key] weights key.new_zeros((key.shape[1], key.shape[2], value.shape[3], key.shape[3])) outputs [] for i in range(seq_len): # 检索v̄^(i) W^(i-1) φ(k^(i)) value_existing torch.einsum(bhvk,bhk-bhv, weights, key[i]) # 写回W^(i) W^(i-1) β^(i)(v^(i) − v̄^(i)) ⊗ φ(k^(i)) weights weights torch.einsum(bhv,bhk-bhvk, beta[i] * (value[i] - value_existing), key[i]) # 读取y^(i) W^(i) φ(q^(i)) y torch.einsum(bhvk,bhk-bhv, weights, query[i]) # 合并多头 outputs.append(y.reshape(y.shape[0], -1)) x torch.stack(outputs) return self.output(x)每一步都对应公式中的一个操作先用einsum(bhvk,bhk-bhv, ...)完成矩阵-向量检索得到 $\bar{v}^{(i)}$再用外积einsum(bhv,bhk-bhvk, ...)将插值残差 $\beta^{(i)}(v^{(i)} - \bar{v}^{(i)})$ 写回权重矩阵最后用更新后的权重读取输出。这样权重矩阵 $W^{(i)}$ 就在序列推进过程中持续演化充当了记忆。token_wise 变体逐 token 推理仓库还在 labml_nn/transformers/fast_weights/token_wise.py 提供了逐 token 处理的变体FastWeightsAttentionTransformer。它的forward先用torch.unbind(x_seq, dim0)把序列拆成单个 token然后为每一层维护一份权重列表weights [None for _ in range(len(self.layers))]每个 token 依次穿过所有层、逐层更新对应权重最后torch.stack(res)并做 LayerNorm。这种写法与自回归推理时一次只喂一个 token、权重跨步长保留的使用方式天然契合。组装 Transformer层与整体模型FastWeightsAttentionTransformerLayerFastWeightsAttentionTransformerLayerinit.py 第 270 行起是一个标准 Transformer 层组合了快速权重自注意力与前馈网络self.size d_model self.attn attn # FastWeightsAttention 实例 self.feed_forward feed_forward self.dropout nn.Dropout(dropout_prob) self.norm_self_attn nn.LayerNorm([d_model]) self.norm_ff nn.LayerNorm([d_model]) def forward(self, x: torch.Tensor): attn self.attn(x) x x self.dropout(attn) # 自注意力残差 z self.norm_ff(x) # 前馈前归一化 ff self.feed_forward(z) x x self.dropout(ff) # 前馈残差 return x结构上与常规 Transformer 层一致自注意力残差连接 → LayerNorm → 前馈网络 → 残差连接只是把 MHA 换成了FastWeightsAttention。前馈网络复用 labml_nn/transformers/feed_forward.py 的FeedForward(d_model, d_ff, dropout)两层线性 激活 dropout默认 ReLU。FastWeightsAttentionTransformerFastWeightsAttentionTransformer负责把单层复制成多层并接最终归一化def __init__(self, layer, n_layers): self.layers clone_module_list(layer, n_layers) # 深拷贝 n_layers 份 self.norm nn.LayerNorm([layer.size]) def forward(self, x): for layer in self.layers: x layer(x) return self.norm(x)层复制使用了仓库工具 labml_nn/utils/init.py 中的clone_module_list(module, n)即nn.ModuleList([copy.deepcopy(module) for _ in range(n)])确保每层拥有独立参数。训练实验Tiny Shakespeare 上的自回归语言模型模型与默认配置训练入口位于 labml_nn/transformers/fast_weights/experiment.py。AutoregressiveModel由三部分组成nn.Embedding(n_vocab, d_model)词嵌入、Fast Weights Transformer 主干、nn.Linear(d_model, n_vocab)输出层前向输出下一个 token 的 logits。Configs继承自 labml_nn/experiments/nlp_autoregression.py 的NLPAutoRegressionConfigs核心默认参数如下参数默认值说明d_model512模型维度nu1DPFP 投影超参数 $\nu$heads8注意力头数d_k d_model // heads 64dropout0.0丢弃概率该实验关闭 dropoutd_ff2048前馈网络隐藏维度n_layers6Transformer 层数模型构造逻辑fast_weights_transformer(c)把上述配置拼装为AutoregressiveModel( c.n_tokens, c.d_model, FastWeightsAttentionTransformer( FastWeightsAttentionTransformerLayer( d_modelc.d_model, attnFastWeightsAttention(c.heads, c.d_model, c.dropout, DPFP(nuc.nu)), feed_forwardFeedForward(c.d_model, c.d_ff, c.dropout), dropout_probc.dropout), c.n_layers))启动训练main()使用 labml 实验框架创建并运行实验覆盖的关键配置包括experiment.configs(conf, { tokenizer: character, # 字符级分词 text: tiny_shakespeare, # Tiny Shakespeare 数据集 optimizer.learning_rate: 1.0, optimizer.optimizer: Noam, # Noam 学习率调度优化器 prompt: It is, # 采样起始提示 prompt_separator: , train_loader: shuffled_train_loader, valid_loader: shuffled_valid_loader, seq_len: 128, # 上下文长度 epochs: 128, batch_size: 16, inner_iterations: 25})运行时执行python labml_nn/transformers/fast_weights/experiment.py即可开始训练仓库同时提供了 experiment.ipynb 笔记可在 Colab 中直接训练同一个模型。训练过程中使用交叉熵损失与准确率指标并支持梯度裁剪grad_norm_clip 1.0见NLPAutoRegressionConfigs。推理与采样得益于快速权重机制模型在生成阶段可以天然地记忆更长的历史token_wise.py中每个 token 输入时各层快速权重矩阵在上一步的基础上继续更新无需重新计算整段历史的注意力。训练时用prompt启动采样配合prompt_separator字符级分词时为空字符串控制提示的拼接方式。源码地图与延伸阅读核心实现labml_nn/transformers/fast_weights/init.pyDPFP、FastWeightsAttention、FastWeightsAttentionTransformerLayer、FastWeightsAttentionTransformer逐 token 变体labml_nn/transformers/fast_weights/token_wise.py训练脚本与笔记labml_nn/transformers/fast_weights/experiment.py、labml_nn/transformers/fast_weights/experiment.ipynb依赖组件多头准备层 labml_nn/transformers/mha.py、前馈网络 labml_nn/transformers/feed_forward.py、模块克隆工具 labml_nn/utils/init.py、自回归训练基类 labml_nn/experiments/nlp_autoregression.py在线文档仓库 docs 目录下的 docs/transformers/fast_weights/index.html、docs/transformers/fast_weights/token_wise.html 与 docs/transformers/fast_weights/experiment.html 与本文章节一一对应总结Fast Weights Transformer 的核心洞察在于线性自注意力中逐步累积的 $W^{(i)} \sum v^{(j)} \otimes \phi(k^{(j)})$ 本质上就是一个快速权重记忆系统。论文及其在本仓库的注释实现通过三点改造提升了该系统的表达能力一是用无参数的 DPFP 投影 $\phi$ 替代 softmax 核得到稀疏、近似正交且自带归一化的投影二是把朴素的直接写回更新规则改为先检索、再插值、后写回用可学习的 $\beta^{(i)}$ 门控控制新信息与旧记忆的融合三是利用 $\phi$ 的归一化省去了额外的归一化项 $z$。从init.py 的每一步 einsum 到 experiment.py 的完整训练配置读者可以沿着本文的路径复现并扩展这个将记忆直接编码进权重演化的高效 Transformer 变体。赞分享人工智能深度学习大模型NLP计算机视觉强化学习LoRA【免费下载链接】annotated_deep_learning_paper_implementations 60 Implementations/tutorials of deep learning papers with side-by-side notes ; including transformers (original, xl, switch, feedback, vit, ...), optimizers (adam, adabelief, sophia, ...), gans(cyclegan, stylegan2, ...), reinforcement learning (ppo, dqn), capsnet, distillation, ... 项目地址https://gitcode.com/gh_mirrors/an/annotated_deep_learning_paper_implementations点击查看免费下载相关推荐黑苹果EFI配置终极指南OpCore Simplify工具快速入门教程黑苹果EFI配置终极指南OpCore Simplify工具快速入门教程 还在为黑苹果复杂的EFI配置而头疼吗OpCore Simplify是一款专门为简化O开发工具CLINYU-DLSP20 注意力机制与 Transformer从自注意力数学推导到 PyTorch 完整实现NYU DLSP20 注意力机制与 Transformer从自注意力数学推导到 PyTorch 完整实现 本文基于 NYU Deep Learning Spr示例工程深入理解视觉Transformer窗口注意力pytorch-image-models中的高效注意力机制深入理解视觉Transformer窗口注意力pytorch image models中的高效注意力机制 视觉Transformer窗口注意力 是计算机视觉领域人工智能计算机视觉深度学习预训练上一篇AutoClip 投稿状态页面使用指南B 站投稿任务的监控、管理与问题排查下一篇从晶体管到 CPU数字电路基础与抽象层次全解析Easy-Vibe 计算机基础篇创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考