第 20 天 · 让记忆槽互相注意，RNN 才会比较过去

一句话

论文提出关系记忆核心（Relational Memory Core）：把 RNN 的隐状态换成多个记忆槽，每一步用多头点积注意力让记忆槽彼此交互。它在需要比较历史信息的任务、强化学习和语言建模上超过 LSTM。

论文讲了什么

结论先行：LSTM 能记住信息，但不擅长把不同时刻记住的信息拿来互相比较。作者把隐状态拆成多个记忆槽，并在每一步对记忆做一次自注意力。这样模型在时间维度上也能做第 19 天那种关系推理。

问题。LSTM 的记忆是一个向量。第 3 步看到的点和第 9 步看到的点被压进同一个向量里。之后要问"这两个点谁离原点更远"，模型必须先把它们拆开再比较。神经图灵机（第 17 天）有多个存储位置，但这些位置之间不直接交互。

方法：关系记忆核心（RMC）。记忆是一个矩阵 M，有 N 行，每行是一个记忆槽。每个时间步做 4 件事：

• 拼接：把新输入 x 拼到记忆下面，得到 [M; x]。
• 注意力：查询（query）只来自 M，键（key）和值（value）来自 [M; x]。每个记忆槽因此能读到其他记忆槽和新输入。计算方式是 softmax(Q K^T / sqrt(d_k)) V，并使用多个头。
• 逐行 MLP：对每个记忆槽单独做一个 MLP，并加残差连接。这一步和 Transformer 的前馈层相同。
• 门控：对每个记忆槽使用类似 LSTM 的输入门和遗忘门，决定保留多少旧记忆、写入多少新内容。

可以把 RMC 理解为：把一个 Transformer 块放进 RNN 的循环里，记忆槽充当固定数量的 token。

关键结果。

• 第 N 远任务（N-th farthest）：依次输入 8 个向量，再问"离第 m 个向量第 n 远的是哪个"。这个任务必须比较所有向量对。论文报告 RMC 约 91%，LSTM 约 30%。
• 程序求值：读入一段简单代码的字符序列，输出执行结果。RMC 在多数设置下优于 LSTM。
• 强化学习：在 Mini PacMan 等环境中，把智能体的 LSTM 换成 RMC 后得分更高。
• 语言建模：在 WikiText-103、Project Gutenberg 和 GigaWord 上降低了困惑度。WikiText-103 测试集困惑度为 31.6。

代价。每一步都要做一次注意力，计算比 LSTM 重。记忆槽数量和每槽维度都是新的超参数。

关键概念

关系记忆核心（Relational Memory Core, RMC）

论文提出的循环单元。隐状态是 N 个记忆槽组成的矩阵，每一步用多头注意力更新，再用门控决定保留多少旧记忆。

记忆槽（memory slot）

记忆矩阵的一行，是一个固定长度的向量。多个记忆槽可以分别存放不同时刻的信息，并通过注意力互相读取。

多头点积注意力（multi-head dot-product attention, MHDPA）

把查询、键、值投影到多个子空间，分别计算 softmax(QK^T/√d) V 再拼接。每个头可以关注不同类型的关系。

门控（gating）

用 sigmoid 输出的 0 到 1 之间的系数，按记忆槽决定遗忘多少旧内容、写入多少新内容。它让长序列训练更稳定。

第 N 远任务（N-th farthest）

论文设计的诊断任务：依次给出若干向量，再问离某个向量第 n 远的是哪一个。答对必须记住全部向量并逐对比较距离。

历史位置：为什么值得读

这篇论文站在 3 条路线的交叉点上。

• 第 17 天的神经图灵机给 RNN 加了外部存储，但存储单元之间不交互。
• 第 19 天的关系网络能比较物体，但只能处理一次给定的静态集合。
• 第 13、14 天的 Transformer 用自注意力处理整个序列，但没有跨时间步的循环状态。

RMC 把这三者合在一起：用 Transformer 的注意力，在 RNN 的记忆里，做关系推理。同一时期，同一团队发表了关系深度强化学习（Relational Deep RL），把类似的自注意力模块用在 BoxWorld 和星际争霸 II 小游戏上。

它为什么值得读：它是"固定大小记忆 + 注意力更新"的早期完整设计。后来的 Transformer-XL（2019）把上一段的状态缓存给下一段使用。Recurrent Memory Transformer（2022）用少量记忆 token 在段之间传递信息。Perceiver（2021）用一小组潜变量反复读取长输入。这些设计都可以在 RMC 里找到雏形。读完第 19 天马上读这篇，可以看清"关系推理"如何从图像搬到时间序列。

学前准备

• 理解 LSTM 的输入门和遗忘门（先看第 2 天）
• 会手写缩放点积注意力（先看第 13、14 天）
• 理解关系网络的逐对比较思想（先看第 19 天）
• 了解神经图灵机的外部记忆（建议看第 17 天）

在原仓库怎么学

1. 安装 requirements.txt 里的依赖。目录里的 setup.py 是打包脚本，可以用 pip install -e . 安装成本地包。
2. 阅读 implementation.py：找出拼接 [M; x]、多头注意力、逐行 MLP、门控 4 个部分。
3. 按顺序做 exercises/ 里的 5 个练习：exercise_01_attention.py 写注意力，exercise_02_gating.py 写门控，exercise_03_rmc_step.py 组装单步 RMC，exercise_04_nth_farthest.py 生成第 N 远任务数据，exercise_05_train.py 完成训练。
4. 对照 exercises/solutions/ 检查自己的实现。
5. 运行 train_minimal.py：默认同时训练 RMC 和 LSTM 做对比。可以改 --seq_len 增加难度，改 --slots、--mem_size、--heads 观察记忆容量的影响，加 --model rmc --visualize 输出注意力图。
6. 用 visualization.py 画图，并查看 plots/ 里已有的注意力热力图和损失对比图。
7. 最后用 notebook.ipynb 复习全流程。

今天再看

RMC 这个具体模块没有流行起来，Transformer 赢了。当算力足够时，直接对整段序列做自注意力，比逐步更新一个小记忆更简单，也更容易并行训练。

但它提出的问题在 2026 年更重要了。大语言模型的 KV 缓存随上下文线性增长，长对话和长文档的推理成本很高。业界因此重新研究"固定大小的状态"：Mamba 等状态空间模型、RWKV、各种线性注意力，以及在段之间传递记忆 token 的循环 Transformer。它们要回答的正是 RMC 的问题：用有限的槽位，怎样保留并比较很久以前的信息。

智能体（agent）的长期记忆也在走类似路线：把记忆分成多条，检索后再与当前输入做注意力。

容易卡住的地方

• 查询、键、值的来源搞混：查询只来自记忆 M，键和值来自 [M; x]。如果查询也包含 x，输出行数会多出来，下一步记忆形状就对不上。
• 第 N 远任务数据生成出错：先算出全部距离并排序，再取第 n 个的下标。用 3 到 4 个向量的小例子手算验证标签。
• RMC 训练不收敛：先把序列长度和记忆槽数调小，确认能过拟合一个小批次，再逐步放大。门控偏置的初始化也会影响稳定性。
• 和 LSTM 对比不公平：两者参数量差距大时，结论没有意义。打印两边的参数总数，尽量调到同一量级。

延伸阅读

• 原论文：Relational recurrent neural networks
• 原仓库第 20 天目录：20_Relational_RNNs/ —— 代码、笔记本、练习与答案（英文）
• Relational Deep Reinforcement Learning —— 同期姊妹篇，把自注意力用于强化学习
• Attention Is All You Need —— RMC 所用多头注意力的出处
• Transformer-XL —— 在 Transformer 里加入段级循环记忆
• Recurrent Memory Transformer —— 用记忆 token 跨段传递信息

← 返回 30 天目录