
ATPM 从记忆操作到记忆动力学
Adaptive Transport–Proximal Memory 设计报告 构建具有 Read、Remember、Revise、Rewrite 能力的连续神经记忆系统
本文记录 ATPM 的设计过程和现有实验。当前证据表明:ATPM 可以作为语言模型中的持续状态模块训练,并已在约 200M 参数、约 20B FineWeb tokens 的规模上完成验证。
本文不讨论更大规模的 scaling law,也不声称已经完成严格 compute-matched 的强基线比较。十亿参数以上的行为仍待验证。
ATPM(Adaptive Transport–Proximal Memory)从 Read、Remember、Revise、Rewrite 四种记忆行为出发,讨论一个具体问题:神经网络里的记忆,能否不只是被动保存过去的矩阵,而是一套可以持续修订的内部状态?
| 已完成规模 | 训练数据 | 持久状态 | 序列复杂度 |
|---|---|---|---|
| 约 200M 参数 | 约 20B FineWeb tokens | n 个联合 K/V 实体状态 state |
1. 一切从“四个 R”开始
现代神经网络中的记忆机制经历了几个不同阶段。
Transformer 重新定义了序列建模:模型不再依赖固定的递归状态,而是通过 attention 在上下文中动态选择相关信息。然而,这种能力依赖于显式上下文窗口。推理阶段的 KV cache 虽然保存了历史信息,但它更接近一种只读的访问结构,而不是能够持续修正自身状态的记忆。
Linear Attention 进一步改变了这一点。通过重新组织 attention 的计算形式,历史 Key/Value 可以被压缩成为一个连续状态:
过去逐 token 保存的上下文,被替换为随时间演化的 state。
随后,SSM、RWKV、DeltaNet 等 recurrent architecture 继续沿着这一方向发展:模型开始维护一个长期存在的内部状态,并通过递推动力学处理无限长度序列。
然而,一个持续更新的状态并不等价于真正的记忆。
大量现有 state memory 最终可以归结为:
其中
ATPM 并不是首先尝试设计一个新的 attention kernel,而是重新审视 recurrent state 本身的结构:如果把"记忆"当成一个真正的认知动作,它到底应该会做什么?
我认为至少需要四种基本行为:
- Read:从过去的状态中找到当前相关的内容。
- Remember:没有新证据时,尽量保留已有状态。
- Revise:新信息到来时,对已有内容做局部、连续的修正。
- Rewrite:旧内容失效时,允许模型明确替换它,而不是一直做指数平均。
但如果所有动作都被压成一个 gate × update,模型很容易退化为保守的平均器:该保留时仍在写入,该重写时又改得不够。Read、Remember、Revise、Rewrite 在数学上没有真正分开,最后只剩下一条连续但含义有限的衰减曲线。
设计原则: 不为 4R 分别增加四个 operator,而让它们成为同一个动力系统在不同区域里的行为。
这条原则也限制了模型复杂度:遇到问题时,不能简单地再加一个 gate。
2. 第一个转折:记忆不是两个表,而是一个“实体”
早期设计中,一个反复出现的问题是 K/V 的角色。
如果把 Key memory 和 Value memory 完全拆成两个状态,精确控制会变得容易;但代价也很明显:模型似乎在维护两个彼此相关、却物理上独立的世界。
我们最终更喜欢另一种解释:一条记忆首先是一个 entity;K 与 V 只是这个 entity 的两个坐标视图。
于是每个 head 只有一个联合状态:
其中:
是 address / mode 维度; 与 是实体内部两块固定坐标; 是唯一持久状态; - K/V 的分离只发生在“修改哪一部分实体”时,而不是维护两份独立 memory。
这一步形成了 PSM — Partial Substitution Memory 的核心。
先衰减旧状态:
从源地址
只计算新实体与旧预测之间的残差:
再让 K/V 具有独立的 substitution strength:
其中
最终更新:
这条式子已经带来了一个我们很喜欢的性质:
| 控制量 | 作用 |
|---|---|
| 保留多少旧状态 | |
| 当前实体与旧预测相差多少 | |
| 实体的哪一部分需要修改 | |
| 残差写到哪里 |
Remember 不再需要一个叫 remember_gate 的东西。
当
Revise 也不需要单独的 operator。
当 residual 小、substitution strength 温和时,更新天然就是局部修订。
Rewrite 则是同一条更新在更强 substitution 下的另一端。
不过,PSM 仍然没有回答写入地址的问题。
3. Source 与 destination:读到的地方不一定是写回的地方
绝大多数 delta-style memory 都隐含了一个假设:
也就是说:我在哪里找到旧内容,就在哪里修改它。
对很多任务,这完全合理。
但对我们关心的一类记忆行为,它不够。
想象模型读到:“Alice currently lives in Paris.”
随后又获得:“Alice moved to Berlin.”
模型首先需要通过旧内容找到 Alice / location 这条关系,但新的表示未必应该严格沿原来的地址方向写回。尤其当 memory space 本身承担了关系组织、冲突消解或实体绑定时:
这就是 Transport 出现的原因。
我们开始显式区分:
:source,从哪里预测旧实体; :destination,残差最终写到哪里。
于是最基本的 substitution 不再是:
而是:
形式上只是替换了左侧向量,但它把读取地址和写入地址分开了,因此改变了 memory geometry。
4. Oblique write:不是“移动地址”,而是在切空间里改变写入方向
如果
我们希望 transport 有方向自由度,但又不想丢掉一个关键事实:这次更新仍然是由 source
因此 ATPM 不直接预测一个任意的
它先产生一个 transport proposal
由于
随后仅沿这个正交方向移动:
其中:
由每个 token、每个 head 自适应预测。
由此可以得到不变量:
因此:
满足:
也就是说,
几何解释: ATPM 保留 source 的约束,同时允许更新沿其 null / tangent space 偏转。
这也是为什么我们更愿意叫它 transport,而不是简单的 write offset。
5. Transport 从哪里来:让 K 与 V 共同决定“偏转方向”
接下来又出现一个设计问题:谁来决定
如果再加一个独立 routing network,模型当然能学,但整个结构又开始像“外挂控制器”。
最终 ATPM-v2 使用 entity 自己的 K/V 两个固定坐标块构造 transport。
先做无仿射 RMS normalization:
映射到一个低秩 transport latent:
做逐元素的双线性交互:
再投回 address space:
其中
为了不让 proposal 的范数失控,我们不做硬单位归一化,而是压到 soft sphere:
于是:
最后:
为什么不用 hard normalize?
当 K/V binding 很弱时,强行把接近 0 的向量归一化到单位球面,会把低置信度变成高幅度的随机方向。soft sphere 则保留 proposal 本身的置信强度。
为什么使用 bilinear K×V?
Transport 不应只由地址或内容中的一方决定。它来自同一 entity 的两个坐标视图之间的交互:写入的对象与对象当前表达的内容共同决定偏转方向。
这里的 transport_rank = R 只是这个双线性映射的内部低秩维度。
它不是一个行为 gate,也不是给不同 head 手工分配 transport 配额的超参数。
早期版本中我们尝试过全局 transport budget、per-head softmax 与额外的
最终式子回到了:
最终版本不再需要额外的 head budget。
6. 第二个坐标:写到哪里,与怎么读回来,不应该是同一件事
有了 transport 之后,一个新的不对称出现了。
我们已经允许:
但外部读取的 query
于是模型可能成功地把状态写向一个新的 destination,却仍然用一个与该 destination 几乎无关的 query 去观察它。
最直接的方法,是让 query 直接等于
但这又太强:它会把内容驱动的自由检索全部抹掉。
于是 ATPM 引入第二个连续坐标:
先定义 destination anchor:
然后只修正 query 在
它对应以下 residual law:
因此:
:保持自由 query; :只纠正一部分 query-anchor mismatch; :精确满足 。
这里的 Proximal 不是指“再加一个相似度 loss”。
它表示:读取不是被硬切换到 destination,而是沿最小的锚点方向修正逐渐靠近它。
7. 两个坐标,四个极限:ATPM 的动态地图
四个端点的形式化解释:
| 控制坐标 | 查询几何 | 写入几何 | 对应端点 |
|---|---|---|---|
| 地址约束 | 有向写入 | RCQP / CrossScan-like | |
| 自由查询 | 有向写入 | Directed memory | |
| 地址约束 | 对称写入 | Stable revision | |
| 自由查询 | 对称写入 | GDN-like |
这张四象限图不是 4R 的一一映射。
与 控制的是 observation geometry 和 destination geometry。 Remember / Revise / Rewrite 还同时受到
、 、 和 residual 的影响。4R 不是四个手工 mode,而是多组连续坐标共同作用的结果。
8. 4R 在状态方程中的对应关系
我进一步逐项说明 4R 的实现方式。
Read
外部读取仍然只有一次:
Remember
如果没有必要更新:
状态自然延续。
没有 remember() 分支。
Revise
如果新实体与旧预测接近:
或者 substitution strength 较小:
同一条 recurrence 就表现为局部修订。
如果
Rewrite
如果 residual 大、substitution 强,同时
模型可以沿新的 destination direction 做强替换,而不是被限制为原地平均。
9. 把 ATPM 压缩成六步
整个模块可以整理为六步:
- Decay:
。先决定旧状态保留多少。 - Predict & compare:
。从源地址读出旧实体,只更新当前实体与旧预测的差。 - Polarize:
。Key 与 Value 属于同一实体,但可以使用不同的修订强度。 - Transport:
。分离“在哪里发现旧内容”和“在哪里写回新内容”。 - Substitute:
。通过一次外积完成局部、可极化、可定向的替换。 - Recall:
。查询既可以保持自由,也可以被近端约束拉向更新后的锚点。
完整形式为:
其中控制量:
全部由当前 token / hidden state 自适应产生。
| 性质 | 形式 | 含义 |
|---|---|---|
| Transport invariant | transport 不破坏 source 对 update 的约束 | |
| Oblique projection | 左侧 transport operator 是 oblique projector | |
| Recall residual law | proximal recall residual 具有精确的一阶收缩律 |
10. 可并行训练
单步 recurrence 可以写成:
它仍然属于带 rank-1 / block-substitution 结构的线性状态更新。
实际实现保留:
- exact sequential reference,用来验证语义与数值;
- exact chunk-WY scan,用来做并行训练;
- 单个联合状态,而不是为 K/V 分别维护两套完整 memory;
- 每个 token 对外只产生一次 memory read;
- state 可以持续带到下一个 chunk / generation step。
训练路径可以概括为:
当前 ATPM-v2 的参考实现采用 write-then-read / read-after-write 语义:
- 用当前 token 构造 source、destination 与 entity;
- 更新
; - 当前 token 的外部 memory read 从更新后的
中产生。 这使当前 observation 能立即看到本步 substitution 的结果。
如果未来实验需要严格的 read-before-write causal memory,也可以把 observation 时序作为独立设计轴研究;它不是 ATPM 几何本身的必要条件。
11. 合成任务上的参数匹配对比
下面比较 ATPM-v2 与 GDN、LSTM、Transformer 三个参数量匹配的模型。所有结果均来自seed 0,并在训练 3000 step 后评估。IID 与训练分布一致;OOD 和 Stress 逐步提高分布偏移与任务难度。
Exact match 和 Token Accuracy 采用 autoregressive evaluation,CE 采用 teacher-forced evaluation。mixed_all 表示九类任务混合训练后的宏平均,不重复计入其中的单项结果。
| 任务 | ATPM-v2 | GDN | LSTM | Transformer |
|---|---|---|---|---|
| 100.0% | 100.0% | 1.8% | 12.3% | |
| 99.7% | 0.4% | 0.0% | 0.2% | |
| 100.0% | 100.0% | 100.0% | 100.0% | |
| 99.6% | 100.0% | 2.9% | 15.3% | |
| 41.9% | 58.1% | 1.5% | 13.4% | |
| 98.9% | 97.8% | 0.0% | 5.0% | |
| 99.8% | 99.9% | 0.0% | 97.3% | |
| 93.0% | 89.8% | 0.0% | 82.0% | |
| 80.8% | 91.0% | 0.0% | 52.3% | |
| 73.1% | 79.6% | 1.0% | 26.1% |
如果对十组运行做简单平均,ATPM-v2 在 IID 上的 autoregressive exact match 为 88.7%,高于这里保留的三个对照模型。它在 KV recall、overwrite 和 rebind 上较稳定:前两项在三个 split 中均达到 100%,rebind 在 Stress 下仍为 92.4%。这并不意味着每项任务都领先,例如 pointer chase 的 IID exact match 只有 41.9%,低于 GDN 的 58.1%。
Stress 下,ATPM-v2 的平均 token accuracy 为 43.8%,exact match 为 33.0%,在这四个模型中均最高。不过,两者之间的差距说明模型在困难样本上仍能保留一部分 token 级信息,却经常无法完整生成整个答案。因此,这组实验支持“ATPM-v2 已能学习多种写入与检索操作”,但不能据此认为它在每种合成记忆任务上都占优。
12. 从合成记忆到真实语言建模
一个 memory operator 在精心设计的 synthetic task 上工作,并不能说明它真的适合作为语言模型的一部分。
目前我们完成的最大规模验证大约是:
| 模型规模 | 训练数据 | 目标 | 结果 |
|---|---|---|---|
| 约 200M 参数 | 约 20B FineWeb tokens | next-token prediction | 可端到端训练 |
这次实验不用于比较 benchmark,而是回答一个可行性问题:
一个包含联合实体状态、K/V partial substitution、oblique transport、proximal recall 与 recurrent scan 的系统,到了真实语言建模规模,会不会因为优化、数值稳定性或动力学退化而直接失败?
至少在这个规模,我们观察到答案是:不会。
ATPM 可以被端到端训练,并表现出正常的语言建模学习过程。
Training traces
图:约 200M 参数 ATPM 的完整语言建模训练轨迹。主图使用断轴同时展示训练初期的快速下降与后期的缓慢收敛,局部视图放大了稳定训练阶段的变化。
13. ATPM 与它的几个边界情况
ATPM 保留了清晰的退化端点,因此可以与已有 memory dynamics 直接比较。
| 边界条件 | 结果 |
|---|---|
| Transport 消失,destination 回到 source,写入退化为对称的局部 substitution | |
| Proximal recall 消失,query 保持原本的内容驱动几何 | |
| K/V polarization 消失,整个 entity 使用相同强度更新 | |
| ATPM 的额外几何自由度全部关闭,走向 Gated DeltaNet-like endpoint |
这意味着 ATPM 是在已有 memory dynamics 上增加两个连续自由度,而不是另起一套无法比较的计算图。
14. 开放问题
下一阶段需要回答的问题包括:
- Scaling:200M 只是第一个真实门槛。1B 以上的 optimization 与 state capacity 仍然未知。
- Mechanism:
与 的平均值不足以解释行为,还需要观察 token、层、head 和语义事件上的 trajectory 与 phase portrait。 - Capacity:ATPM 可以 revise,但 persistent state 的容量仍然有限。覆盖与干扰在什么条件下发生,需要单独刻画。
- Online learning:如果生成时的持续状态进一步承担 latent learning,它可能连接到更广义的 test-time / online adaptation。
- Transfer:联合实体状态能否在模型之间迁移,目前没有答案。
- Kernel:exact chunk-WY 说明该结构可以并行,但吞吐、显存和硬件友好度仍有优化空间。
15. 结语
研究开始时,我们想要的是:
ATPM 没有为每种行为分别设计一个机制,而是只维护一个持续状态:
只有一条 substitution recurrence:
这条 recurrence 由几个具有明确几何含义的连续坐标控制:
现有实验表明,这套动力学可以进入真实语言模型,并在约 200M 参数、约 20B FineWeb tokens 的尺度上完成训练。
这些结果还不足以给“记忆”下一个最终定义,但可以把问题收窄到更具体的计算形式。
Appendix A · Symbols
| Symbol | Meaning |
|---|---|
| per-head joint persistent entity state, | |
| current entity written into memory | |
| normalized source address | |
| transport-aware destination address | |
| unconstrained external recall query | |
| proximal-corrected recall query | |
| state retention / decay coordinate | |
| independent K/V substitution strengths | |
| fixed K/V coordinate projectors | |
| adaptive directed-transport coordinate | |
| adaptive proximal-recall coordinate | |
| source-orthogonal transport tangent | |
| low-rank bilinear transport latent dimension |
Appendix B · Implementation invariants worth testing
如果你复现 ATPM,比对最终 loss 之前,先检查这些结构性不变量:
- source
must be unit-normalized must be orthogonal to because of the soft-sphere map - lambda residual law:
- sequential reference
chunk-WY parallel path controllers begin from the neutral 0.5 point
