- 人工智能
- 深度学习
- 大模型
- 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, ... 🧠
本文是 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_{i=1}$，每个时间步的输入是大小为 $d_{in}$ 的向量 $x \in \mathbb{R}^{d_{in}}$。快速权重模型在每个时间步生成一个权重矩阵，用来产生输出 ${y^{(i)}}^L_{i=1}$（$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_{j=1} \frac{ v^{(j)} \kappa(k^{(j)}, q^{(i)}) }{ \sum^i_{j'=1} \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_{j=1} v^{(j)} \otimes \phi(k^{(j)}) \Big) \phi(q^{(i)}) }{ \Big( \sum^i_{j'=1} \phi(k^{(j')}) \Big) \phi(q^{(i)}) } $$
如果定义 $W^{(i)} = \sum^i_{j=1} v^{(j)} \otimes \phi(k^{(j)})$ 与 $z^{(i)} = \sum^i_{j=1} \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 Projection（DPFP），实现于 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}}_{j=1} \phi(k)_j} $$
对应源码实现如下：
def forward(self, k: torch.Tensor): # 计算 φ(k) k = self.dpfp(k) # 按 ∑ φ(k)_j 归一化，eps 防止除零 return k / (torch.sum(k, dim=-1, keepdim=True) + 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(shifts=i, 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$，控制投影维度扩张倍数，默认1；eps：归一化时避免除零的小常数，默认1e-6。
实现上利用了 PyTorch 的torch.roll(shifts=i, 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, bias=False) self.key = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False) self.value = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False) self.interpolation_weight = nn.Sequential( PrepareForMultiHeadAttention(d_model, heads, 1, bias=False), 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, dim=0)把序列拆成单个 token，然后为每一层维护一份权重列表weights = [None for _ in range(len(self.layers))]，每个 token 依次穿过所有层、逐层更新对应权重，最后torch.stack(res)并做 LayerNorm。这种写法与自回归推理时"一次只喂一个 token、权重跨步长保留"的使用方式天然契合。
组装 Transformer：层与整体模型
FastWeightsAttentionTransformerLayer
FastWeightsAttentionTransformerLayer（init.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）。
FastWeightsAttentionTransformer
FastWeightsAttentionTransformer负责把单层复制成多层并接最终归一化：
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_model | 512 | 模型维度 |
nu | 1 | DPFP 投影超参数 $\nu$ |
heads | 8 | 注意力头数，d_k = d_model // heads = 64 |
dropout | 0.0 | 丢弃概率（该实验关闭 dropout） |
d_ff | 2048 | 前馈网络隐藏维度 |
n_layers | 6 | Transformer 层数 |
模型构造逻辑fast_weights_transformer(c)把上述配置拼装为：
AutoregressiveModel( c.n_tokens, c.d_model, FastWeightsAttentionTransformer( FastWeightsAttentionTransformerLayer( d_model=c.d_model, attn=FastWeightsAttention(c.heads, c.d_model, c.dropout, DPFP(nu=c.nu)), feed_forward=FeedForward(c.d_model, c.d_ff, c.dropout), dropout_prob=c.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.py（
DPFP、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, ... 🧠
相关推荐
黑苹果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窗口注意力 是计算机视觉领域
人工智能计算机视觉深度学习预训练
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
