Capybara & Friends

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 tokensn 个联合 K/V 实体状态 stateO(T)

1. 一切从“四个 R”开始

现代神经网络中的记忆机制经历了几个不同阶段。

Transformer 重新定义了序列建模:模型不再依赖固定的递归状态,而是通过 attention 在上下文中动态选择相关信息。然而,这种能力依赖于显式上下文窗口。推理阶段的 KV cache 虽然保存了历史信息,但它更接近一种只读的访问结构,而不是能够持续修正自身状态的记忆。

Linear Attention 进一步改变了这一点。通过重新组织 attention 的计算形式,历史 Key/Value 可以被压缩成为一个连续状态:

Mt=f(Mt1,xt)

过去逐 token 保存的上下文,被替换为随时间演化的 state。

随后,SSM、RWKV、DeltaNet 等 recurrent architecture 继续沿着这一方向发展:模型开始维护一个长期存在的内部状态,并通过递推动力学处理无限长度序列。

然而,一个持续更新的状态并不等价于真正的记忆。

大量现有 state memory 最终可以归结为:

Mt=γtMt1+ΔMt

其中 γt 控制旧状态保留程度,而 ΔMt 注入当前信息。

ATPM 并不是首先尝试设计一个新的 attention kernel,而是重新审视 recurrent state 本身的结构:如果把"记忆"当成一个真正的认知动作,它到底应该会做什么?

我认为至少需要四种基本行为:

  1. Read:从过去的状态中找到当前相关的内容。
  2. Remember:没有新证据时,尽量保留已有状态。
  3. Revise:新信息到来时,对已有内容做局部、连续的修正。
  4. Rewrite:旧内容失效时,允许模型明确替换它,而不是一直做指数平均。

但如果所有动作都被压成一个 gate × update,模型很容易退化为保守的平均器:该保留时仍在写入,该重写时又改得不够。Read、Remember、Revise、Rewrite 在数学上没有真正分开,最后只剩下一条连续但含义有限的衰减曲线。

设计原则: 不为 4R 分别增加四个 operator,而让它们成为同一个动力系统在不同区域里的行为。

这条原则也限制了模型复杂度:遇到问题时,不能简单地再加一个 gate。

2. 第一个转折:记忆不是两个表,而是一个“实体”

早期设计中,一个反复出现的问题是 K/V 的角色。

如果把 Key memory 和 Value memory 完全拆成两个状态,精确控制会变得容易;但代价也很明显:模型似乎在维护两个彼此相关、却物理上独立的世界。

我们最终更喜欢另一种解释:一条记忆首先是一个 entity;K 与 V 只是这个 entity 的两个坐标视图。

于是每个 head 只有一个联合状态:

EtRS×(dk+dv)et=[kt;vt]

其中:

  • S 是 address / mode 维度;
  • dkdv 是实体内部两块固定坐标;
  • Et 是唯一持久状态;
  • K/V 的分离只发生在“修改哪一部分实体”时,而不是维护两份独立 memory。

这一步形成了 PSM — Partial Substitution Memory 的核心。

先衰减旧状态:

E¯t=αtEt1

从源地址 st 读取当前预测:

e^t=E¯tst

只计算新实体与旧预测之间的残差:

rt=ete^t

再让 K/V 具有独立的 substitution strength:

Ωt=βtKΠK+βtVΠV

其中 ΠK,ΠV 是联合实体空间上的固定坐标投影。

最终更新:

Et=E¯t+dt(Ωtrt)

这条式子已经带来了一个我们很喜欢的性质:

控制量作用
αt保留多少旧状态
rt当前实体与旧预测相差多少
βtK,βtV实体的哪一部分需要修改
dt残差写到哪里

Remember 不再需要一个叫 remember_gate 的东西。

αt 高、而 βt 小时,旧状态自然留下来。

Revise 也不需要单独的 operator。

当 residual 小、substitution strength 温和时,更新天然就是局部修订。

Rewrite 则是同一条更新在更强 substitution 下的另一端。

不过,PSM 仍然没有回答写入地址的问题。

3. Source 与 destination:读到的地方不一定是写回的地方

绝大多数 delta-style memory 都隐含了一个假设:

source addressdestination address

也就是说:我在哪里找到旧内容,就在哪里修改它。

对很多任务,这完全合理。

但对我们关心的一类记忆行为,它不够。

想象模型读到:“Alice currently lives in Paris.”

随后又获得:“Alice moved to Berlin.”

模型首先需要通过旧内容找到 Alice / location 这条关系,但新的表示未必应该严格沿原来的地址方向写回。尤其当 memory space 本身承担了关系组织、冲突消解或实体绑定时:

where I found itwhere the corrected state should live

这就是 Transport 出现的原因。

我们开始显式区分:

  • stsource,从哪里预测旧实体;
  • dtdestination,残差最终写到哪里。

于是最基本的 substitution 不再是:

strt

而是:

dtrt

形式上只是替换了左侧向量,但它把读取地址和写入地址分开了,因此改变了 memory geometry。

4. Oblique write:不是“移动地址”,而是在切空间里改变写入方向

如果 dt 可以任意偏离 st,系统很快会变得难以控制。

我们希望 transport 有方向自由度,但又不想丢掉一个关键事实:这次更新仍然是由 source st 找到的那条旧记忆所触发的。

因此 ATPM 不直接预测一个任意的 dt

它先产生一个 transport proposal bt,再把它投影到 st 的切空间:

τt=(Istst)bt

由于 st 已归一化:

stτt=0

随后仅沿这个正交方向移动:

dt=st+ρtτt

其中:

ρt(0,1)

由每个 token、每个 head 自适应预测。

由此可以得到不变量:

stdt=stst+ρtstτt=1

因此:

Pt=dtst

满足:

Pt2=dt(stdt)st=dtst=Pt

也就是说,dtst 是一个 oblique projector

几何解释: ATPM 保留 source 的约束,同时允许更新沿其 null / tangent space 偏转。

这也是为什么我们更愿意叫它 transport,而不是简单的 write offset。

5. Transport 从哪里来:让 K 与 V 共同决定“偏转方向”

接下来又出现一个设计问题:谁来决定 τt

如果再加一个独立 routing network,模型当然能学,但整个结构又开始像“外挂控制器”。

最终 ATPM-v2 使用 entity 自己的 K/V 两个固定坐标块构造 transport。

先做无仿射 RMS normalization:

k^t=RMS0(kt)v^t=RMS0(vt)

映射到一个低秩 transport latent:

utK=AKk^tutV=AVv^t

做逐元素的双线性交互:

zt=utKutV

再投回 address space:

ct=CztR

其中 R 是 transport latent rank。

为了不让 proposal 的范数失控,我们不做硬单位归一化,而是压到 soft sphere:

bt=ct1+ct2

于是:

bt<1

最后:

τt=(Istst)btdt=st+ρtτt

为什么不用 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 与额外的 rh。最后这些都被删掉了。

最终式子回到了:

dt=st+ρtτt

最终版本不再需要额外的 head budget。

6. 第二个坐标:写到哪里,与怎么读回来,不应该是同一件事

有了 transport 之后,一个新的不对称出现了。

我们已经允许:

stdt

但外部读取的 query qt0 仍然可以完全自由。

于是模型可能成功地把状态写向一个新的 destination,却仍然用一个与该 destination 几乎无关的 query 去观察它。

最直接的方法,是让 query 直接等于 dt

但这又太强:它会把内容驱动的自由检索全部抹掉。

于是 ATPM 引入第二个连续坐标:

λt(0,1)

先定义 destination anchor:

at=dtdt

然后只修正 query 在 at 方向上的 residual:

qt=qt0+λt(1(qt0)at)at

它对应以下 residual law:

1qtat=(1λt)(1(qt0)at)

因此:

  • λt=0:保持自由 query;
  • 0<λt<1:只纠正一部分 query-anchor mismatch;
  • λt=1:精确满足 qtat=1

这里的 Proximal 不是指“再加一个相似度 loss”。

它表示:读取不是被硬切换到 destination,而是沿最小的锚点方向修正逐渐靠近它。

7. 两个坐标,四个极限:ATPM 的动态地图

λρ 定义了一个二维控制空间。下面列出四个端点;它们是同一套动力学的极限情况,不是四个离散模式。

四个端点的形式化解释:

控制坐标 (λ,ρ)查询几何写入几何对应端点
(1,1)地址约束有向写入RCQP / CrossScan-like
(0,1)自由查询有向写入Directed memory
(1,0)地址约束对称写入Stable revision
(0,0)自由查询对称写入GDN-like

这张四象限图不是 4R 的一一映射。

λρ 控制的是 observation geometrydestination geometry

Remember / Revise / Rewrite 还同时受到 αtβtKβtV 和 residual rt 的影响。4R 不是四个手工 mode,而是多组连续坐标共同作用的结果。

8. 4R 在状态方程中的对应关系

我进一步逐项说明 4R 的实现方式。

Read

外部读取仍然只有一次:

yt=Etqt

qt 可以保持内容驱动,也可以由 λt 向 destination anchor 做 proximal correction。

Remember

如果没有必要更新:

αt1βtK,βtV0

状态自然延续。

没有 remember() 分支。

Revise

如果新实体与旧预测接近:

rt1

或者 substitution strength 较小:

0<βtK,V1

同一条 recurrence 就表现为局部修订。

如果 ρt0,这种修订还会稳定发生在原地址附近。

Rewrite

如果 residual 大、substitution 强,同时 ρt 打开:

rtβtρt

模型可以沿新的 destination direction 做强替换,而不是被限制为原地平均。

9. 把 ATPM 压缩成六步

整个模块可以整理为六步:

  1. DecayE¯t=αtEt1。先决定旧状态保留多少。
  2. Predict & comparert=etE¯tst。从源地址读出旧实体,只更新当前实体与旧预测的差。
  3. PolarizeΩt=βtKΠK+βtVΠV。Key 与 Value 属于同一实体,但可以使用不同的修订强度。
  4. Transportdt=st+ρtτt。分离“在哪里发现旧内容”和“在哪里写回新内容”。
  5. SubstituteEt=E¯t+dt(Ωtrt)。通过一次外积完成局部、可极化、可定向的替换。
  6. Recallqt=qt0+λt(1(qt0)at)at。查询既可以保持自由,也可以被近端约束拉向更新后的锚点。

完整形式为:

E¯t=αtEt1rt=etE¯tstΩt=βtKΠK+βtVΠVzt=(AKk^t)(AVv^t)bt=Czt/R1+Czt/R2τt=(Istst)btdt=st+ρtτtEt=E¯t+dt(Ωtrt)at=dtdtqt=qt0+λt(1(qt0)at)atyt=Etqt

其中控制量:

αt,βtK,βtV,ρt,λt

全部由当前 token / hidden state 自适应产生。

性质形式含义
Transport invariantstdt=1transport 不破坏 source 对 update 的约束
Oblique projection(dtst)2=dtst左侧 transport operator 是 oblique projector
Recall residual lawϵq=(1λ)ϵqproximal recall residual 具有精确的一阶收缩律

10. 可并行训练

单步 recurrence 可以写成:

Et=αtEt1+dt[Ωt(etαtEt1st)]

它仍然属于带 rank-1 / block-substitution 结构的线性状态更新。

实际实现保留:

  • exact sequential reference,用来验证语义与数值;
  • exact chunk-WY scan,用来做并行训练;
  • 单个联合状态,而不是为 K/V 分别维护两套完整 memory;
  • 每个 token 对外只产生一次 memory read;
  • state 可以持续带到下一个 chunk / generation step。

训练路径可以概括为:

xtProjection(q0,s,e,α,β)TransportdSubstituteEtProximal Recallyt

当前 ATPM-v2 的参考实现采用 write-then-read / read-after-write 语义:

  1. 用当前 token 构造 source、destination 与 entity;
  2. 更新 Et1Et
  3. 当前 token 的外部 memory read 从更新后的 Et 中产生。

这使当前 observation 能立即看到本步 substitution 的结果。
如果未来实验需要严格的 read-before-write causal memory,也可以把 observation 时序作为独立设计轴研究;它不是 ATPM 几何本身的必要条件。

11. 合成任务上的参数匹配对比

下面比较 ATPM-v2 与 GDN、LSTM、Transformer 三个参数量匹配的模型。所有结果均来自seed 0,并在训练 3000 step 后评估。IID 与训练分布一致;OOD 和 Stress 逐步提高分布偏移与任务难度。

Exact matchToken Accuracy 采用 autoregressive evaluation,CE 采用 teacher-forced evaluation。mixed_all 表示九类任务混合训练后的宏平均,不重复计入其中的单项结果。

任务ATPM-v2GDNLSTMTransformer
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 tokensnext-token prediction可端到端训练

这次实验不用于比较 benchmark,而是回答一个可行性问题:

一个包含联合实体状态、K/V partial substitution、oblique transport、proximal recall 与 recurrent scan 的系统,到了真实语言建模规模,会不会因为优化、数值稳定性或动力学退化而直接失败?

至少在这个规模,我们观察到答案是:不会。

ATPM 可以被端到端训练,并表现出正常的语言建模学习过程。

Training traces

约 200M 参数 ATPM 在约 20B FineWeb tokens 上的语言建模训练损失曲线

图:约 200M 参数 ATPM 的完整语言建模训练轨迹。主图使用断轴同时展示训练初期的快速下降与后期的缓慢收敛,局部视图放大了稳定训练阶段的变化。

13. ATPM 与它的几个边界情况

ATPM 保留了清晰的退化端点,因此可以与已有 memory dynamics 直接比较。

边界条件结果
ρ0Transport 消失,destination 回到 source,写入退化为对称的局部 substitution
λ0Proximal recall 消失,query 保持原本的内容驱动几何
βK=βVK/V polarization 消失,整个 entity 使用相同强度更新
ρ0, λ0, βK=βVATPM 的额外几何自由度全部关闭,走向 Gated DeltaNet-like endpoint

这意味着 ATPM 是在已有 memory dynamics 上增加两个连续自由度,而不是另起一套无法比较的计算图。

14. 开放问题

下一阶段需要回答的问题包括:

  1. Scaling:200M 只是第一个真实门槛。1B 以上的 optimization 与 state capacity 仍然未知。
  2. Mechanismρλ 的平均值不足以解释行为,还需要观察 token、层、head 和语义事件上的 trajectory 与 phase portrait。
  3. Capacity:ATPM 可以 revise,但 persistent state 的容量仍然有限。覆盖与干扰在什么条件下发生,需要单独刻画。
  4. Online learning:如果生成时的持续状态进一步承担 latent learning,它可能连接到更广义的 test-time / online adaptation。
  5. Transfer:联合实体状态能否在模型之间迁移,目前没有答案。
  6. Kernel:exact chunk-WY 说明该结构可以并行,但吞吐、显存和硬件友好度仍有优化空间。

15. 结语

研究开始时,我们想要的是:

ReadRememberReviseRewrite

ATPM 没有为每种行为分别设计一个机制,而是只维护一个持续状态:

Et

只有一条 substitution recurrence:

Et=αtEt1+dt[Ωt(etαtEt1st)]

这条 recurrence 由几个具有明确几何含义的连续坐标控制:

αt, βtK, βtV, ρt, λt

现有实验表明,这套动力学可以进入真实语言模型,并在约 200M 参数、约 20B FineWeb tokens 的尺度上完成训练。

这些结果还不足以给“记忆”下一个最终定义,但可以把问题收窄到更具体的计算形式。

Appendix A · Symbols

SymbolMeaning
Etper-head joint persistent entity state, S×(dk+dv)
et=[kt;vt]current entity written into memory
stnormalized source address
dttransport-aware destination address
qt0unconstrained external recall query
qtproximal-corrected recall query
αtstate retention / decay coordinate
βtK,βtVindependent K/V substitution strengths
ΠK,ΠVfixed K/V coordinate projectors
ρtadaptive directed-transport coordinate
λtadaptive proximal-recall coordinate
τtsource-orthogonal transport tangent
Rlow-rank bilinear transport latent dimension

Appendix B · Implementation invariants worth testing

如果你复现 ATPM,比对最终 loss 之前,先检查这些结构性不变量:

  1. source s must be unit-normalized
  2. τ must be orthogonal to s
  3. sTd1
  4. ||b||<1 because of the soft-sphere map
  5. ρ,λ(0,1)
  6. lambda residual law: 1qTa=(1λ)(1q0Ta)
  7. sequential reference chunk-WY parallel path
  8. ρ/λ controllers begin from the neutral 0.5 point