Browse this library

Chinese source article

5 · TestingTime 训练与长上下文

把推理阶段的梯度更新写进序列模型本身: TTT layer 的隐状态是一个小模型的权重, LaCT 把更新块放大到上千 token, TTT-E2E 直接在 Transformer 的 MLP 上做下一个 token 预测的梯度步. 以及它们和线性注意力, SSM, RWKV, Titans 的边界.

View original Markdown ↗

第 4 章的推理阶段更新面对的是分布偏移, 每个样本是一张图. 这一章面对的是长上下文: 输入是一条上万到上百万 token 的序列, 模型要一边读一边把前文压进一个状态. 全注意力保留全部 KV, 每个 token 的开销随长度增长; 线性 RNN 把历史压进一个固定大小的矩阵, 开销与长度无关, 但装得下的东西有限. TTT 这一支的做法是让状态本身是一个模型的权重, 让「读一段上下文」等于「在这段上下文上训练一步」. 本章三篇: Fast weights 与 TTT layers 给出这种层的定义和它与注意力的等价关系, 块级更新与端到端 TTT 讲之后两条改进, 序列模型边界 把这些方法和线性注意力一族放进同一张表.

1 三篇怎样衔接

1.1 从 TTT layer 到 LaCT 和 TTT-E2E

5.1 的 TTT layer 有三个组件: 初始状态 W0W_0, 更新规则和输出规则. 状态 WW 是一个线性模型 (TTT-Linear) 或两层 MLP (TTT-MLP) 的权重, 更新规则是对键值重建损失走一步梯度, 输出规则是用更新后的模型对 query 做前向. 键, 值, query 三个视角的投影是外循环学到的参数, 内循环只更新 WW. 5.1 手算了 TTT-Linear 的一步, 结果就是 delta rule; 论文的定理 1 说明线性模型配 batch GD 时 TTT 等价于线性注意力, 定理 2 说明换成核回归时等价于自注意力. 为了让它在硬件上跑起来, 论文把序列切成 b=16b=16 的块, 块内梯度都相对块开始时的权重求, 再用 dual form 把外积换成矩阵乘, 额外开销约 6%, 训练快 5 倍以上.

b=16b=16 留下了两个缺口, 5.2 的两篇各补一个. 第一个是硬件: 每 16 个 token 就要读写一次 fast weights, GPU 的 FLOPs 利用率常低于 5%. LaCT 把块放大到 2K 到 1M 个 token, 块内依赖交给窗口注意力, 纯 PyTorch 实现在 A100 上利用率最高到 70%, fast weights 可以占到模型权重的 40%. 第二个是损失: TTT layer 的内循环用键值重建, 与模型真正的目标不一致. TTT-E2E 不再设专门的层, 直接在滑动窗口 Transformer 后 1/4 块的 MLP 上, 对下一个 token 预测损失每 1K token 走一步梯度, 并在训练阶段元学习这些 MLP 的初始化.

1.2 5.3 给出的坐标

5.3 不引入新方法, 它把全注意力, 滑动窗口, 线性注意力, Mamba-2, DeltaNet, Gated DeltaNet, RWKV-7, Titans 和 TTT 系列都写成「一个状态, 一条写入规则, 一条读出规则」. 按写入规则排下来, 线性注意力只写不擦, Mamba-2 和 GLA 先衰减再写, DeltaNet 按误差先擦后写, Gated DeltaNet 和 RWKV-7 把衰减和擦写合在一起, TTT 系列和 Titans 把写入换成一步梯度. 写入规则越强, 固定大小的状态能区分的关联越多.

这个排序也说明了几个名字之间的关系. 无动量, 矩阵记忆版的 Titans 与 Gated DeltaNet 是同一族写入规则, 参数化略有不同, Titans 多出来的是动量和深层 MLP 记忆; TTT-E2E 论文把 Titans 归入与 TTT-KVB 相同的构造, 即每层一个以键值重建为损失的可更新模块加滑动窗口注意力. 5.2 正是从 TTT-KVB 出发一步步推到 TTT-E2E: 760M 模型在 DCLM 8K 上, 滑动窗口基线的损失 2.827, TTT-E2E 为 2.805.

2 共用记号与概念

2.1 状态, 写入与读出

本章记序列为 x1,…,xTx_1,\dots,x_T, 第 tt 步的状态为 WtW_t (5.3 里线性 RNN 一族写作 StS_t), 键, 值, query 为 kt,vt,qtk_t,v_t,q_t. 线性注意力的写入是 St=St−1+vtkt⊤S_t=S_{t-1}+v_tk_t^\top, DeltaNet 是 St=St−1(I−βtktkt⊤)+βtvtkt⊤S_t=S_{t-1}(I-\beta_tk_tk_t^\top)+\beta_tv_tk_t^\top, TTT 的写入是 Wt=Wt−1−η∇ℓ(Wt−1;xt)W_t=W_{t-1}-\eta\nabla\ell(W_{t-1};x_t), 读出都是 yt=f(qt;Wt)y_t=f(q_t;W_t). TTT-Linear 取 ℓ=∥Wkt−vt∥2\ell=\lVert W k_t-v_t\rVert^2 时, 梯度步展开后就是 DeltaNet 的形式, 这是 5.1 手算的内容, 也是 5.3 能把它们放进一张表的原因.

状态大小决定了压缩的程度. 5.3 算过一个例子: 每层 32 个 head 的矩阵状态, 32 层合计 32 MB, 与上下文长度无关; 同一模型的 KV cache 每个 token 0.5 MB, 到 128K 时约为矩阵状态的 2048 倍. 反过来看, 128K 个 token 的全部信息要挤进 32 MB, 平均每个 token 只分到 256 字节. LaCT 和 TTT-E2E 把状态做大, 都是为了在这个约束下多存一些.

2.2 内循环, 外循环与状态的生命周期

本章的「训练」分两层. 内循环发生在推理阶段, 只更新状态 WW, 一条序列结束就丢弃; 外循环是普通的训练, 更新投影矩阵, 学习率参数和初始状态 W0W_0. TTT-E2E 的元学习初始化就是外循环: 训练时对内循环的更新求二阶梯度, 让 W0W_0 成为「走几步梯度后效果最好」的起点. 这也是它训练贵的原因.

状态活多久, 决定了它和持续学习的关系. 5.3 把状态分成三层: 训练结束后固定的网络权重, 一条序列内的状态, 跨样本延续的适应状态. 本章所有方法的推理阶段学习都在第二层, Titans 的「长期记忆」是相对注意力窗口而言, TTT-E2E 说的「持续学习」也指序列内. 第 4 章的 TENT online, CoTTA, EATA 在第三层, 所以有误差积累, 遗忘和跨用户影响; 本章的方法在序列结束时复位, 没有这些问题, 也不会把一次对话里学到的东西带到下一次.

2.3 预测误差怎样写成状态

把序列状态统一写成 sts_t, 当前 token 产生键 ktk_t、值 vtv_t 和查询 qtq_t. 一般的状态模型包含写入与读出:

st=Uψ(st−1,kt,vt),yt=Rψ(st,qt).s_t=U_\psi(s_{t-1},k_t,v_t),\qquad y_t=R_\psi(s_t,q_t).

attention 把历史的键值对显式保留,读出时让查询和全部键比较;SSM、RWKV 与线性注意力把历史递推进固定状态;TTT layer 则令 sts_t 本身是一组快速权重 WtW_t, 写入规则来自预测或重建误差:

Wt=Wt−1−ηt∇Wℓt(Wt−1;kt,vt),yt=fWt(qt).W_t=W_{t-1}-\eta_t\nabla_W\ell_t(W_{t-1};k_t,v_t),\qquad y_t=f_{W_t}(q_t).

这里的「训练」是一个可微状态转移. 外循环决定 kt,vt,qtk_t,v_t,q_t 怎样构造、损失 ℓt\ell_t 测什么、学习率多大以及初始状态 W0W_0 在哪里. 内循环在读取上下文时执行这些已经学好的规则. 状态虽然以权重形式存在,角色更接近 RNN hidden state,而非发布后永久修改的模型参数.

当 fW(k)=Wkf_W(k)=Wk,损失取平方误差 ℓt=12∥Wkt−vt∥2\ell_t=\frac12\lVert Wk_t-v_t\rVert^2,梯度为 (Wkt−vt)kt⊤(Wk_t-v_t)k_t^\top,更新展开为

Wt=Wt−1(I−ηtktkt⊤)+ηtvtkt⊤.W_t=W_{t-1}(I-\eta_tk_tk_t^\top)+\eta_tv_tk_t^\top.

第一项沿当前键方向擦除旧预测,第二项写入新值,这就是 delta rule. 线性注意力只有外积累加 vtkt⊤v_tk_t^\top,缺少由当前预测误差决定的擦除;Gated DeltaNet 与 RWKV 再加入数据依赖的衰减或门控;更深的 TTT-MLP、Titans 记忆把 WW 从矩阵扩展成非线性模型. 5.1 与 5.3 正是沿这条更新式建立共同坐标.

误差驱动写入的价值在于「已经记住的关联少改,预测错的关联多改」. 若 Wt−1kt≈vtW_{t-1}k_t\approx v_t,梯度接近零;若误差大,更新更强. 这比无条件累加更能控制容量冲突,也会把错误的自监督目标写进状态. 因而 fast weights 的能力取决于损失是否与未来读出相关,和第 4 章代理目标对主任务风险的关系具有相似数学外形,episode 与任务语义却不同.

2.4 attention、SSM、RWKV 与 Titans 的状态语义

全 attention 的状态是一组不断增长的 KV 项. 它保留每个位置的独立寻址线索,精确检索强,读取成本与历史长度一起增长. 滑动窗口只保留最近片段,把更早历史直接丢弃. 这两者的状态含义最直观:存了哪些 token,就能回看哪些 token.

SSM 将历史投影进固定维度递推状态. 线性时可写为 st=Atst−1+Btxts_t=A_ts_{t-1}+B_tx_t,读出为 yt=Ctsty_t=C_ts_t. 衰减矩阵决定旧信息保留多久,固定容量要求不同历史在同一向量里叠加. Mamba 的输入依赖参数让选择性随 token 改变,但状态仍是压缩后的激活,不是通过显式梯度训练的小模型.

RWKV 同样以递推形式实现类似注意力的加权汇总,时间衰减和门控决定写入、保留与读出. RWKV-7 更接近带误差修正和数据依赖衰减的矩阵状态族. 它可以与 DeltaNet、Titans 放在统一更新表中比较,不能仅凭「像 RNN」或「像 attention」概括状态能力.

Titans 将长期记忆写成可训练模块,利用惊讶度相关信号、动量和遗忘进行更新,并与短期注意力结合. 它的记忆也是 fast weights,通常在序列或规定上下文边界内工作. 「长期」指相对局部注意力窗口保存更远信息,不自动意味着跨会话永久学习. 若实现把记忆跨用户延续,第 4 章的污染、权限与回滚问题会立即出现.

TTT layer 与 Titans 的共同点是用优化过程写状态,差异在损失、模块结构、动量与主架构位置. TTT-KVB 用键值重建训练专门的内循环层;TTT-E2E 将下一个 token 预测损失直接施加到 Transformer 部分 MLP;Titans 保留专门长期记忆模块并与注意力组合. 名称相近时,先问快速权重在哪里、由哪个损失更新、何时复位,便能避免把不同系统混成一类.

LaCT 的重点又不同. 它没有改变「快速权重由块内上下文更新」这个核心,而是把更新块从很短的 token 组扩展到数千乃至更长,让大矩阵乘法提高硬件利用率. 块内 token 的细粒度依赖交给窗口 attention,块间信息进入 fast weights. 这是一种计算分工:窗口负责近邻顺序,TTT 状态负责跨块压缩.

TTT-E2E 则移动了训练位置. KVB 路线先人为设计键值重建目标,再期待它服务语言建模;E2E 路线直接用 next-token loss 更新后部 MLP,并在外循环中元学习适合被快速更新的初始化. 代理目标与最终目标更一致,训练需要穿过内循环求梯度,二阶计算和显存随之增加. 5.2 讨论的重点正是这种目标对齐换来的成本.

这几类方法都可以混合. 窗口 attention 提供精确近邻,SSM 或矩阵状态提供便宜递推,TTT/Titans 提供更强写入. 混合架构不能只按模块名称判断谁贡献长上下文,需要做关闭各分支、匹配状态大小与计算的消融. 精确 needle retrieval、平均语言模型损失和推理任务分别检验不同能力,某一项领先不代表状态在所有意义上更强.

状态容量也不等于参数数量. 一个 d×dd\times d 矩阵有 d2d^2 个标量,但连续写入的关联会互相干扰;深层 MLP 参数更多,优化一步能访问的有效方向仍受梯度、学习率和数据块限制. 衡量容量需要看写入多少关联后旧关联还能否读出、相似键是否混淆,以及跨块推理是否保持,而非只比较 MB.

读写次序是另一个边界. 若先用 xtx_t 更新再输出 yty_t,当前 token 已参与自己的状态,任务要避免信息泄漏;常见设计用由过去 token 构造的监督,或让更新服务后续位置. 若先读后写,语义更接近因果 RNN,当前信息只能影响未来. 5.1 的公式推导与 5.2 的块级训练都要明确这一次序,否则困惑度比较可能包含不同可见信息.

块化还引入延迟. 块内并行更新常相对块起点计算梯度,块内 token 看不到逐 token 最新状态,硬件效率提高、在线性减弱. 块越大,矩阵乘法越饱满,状态刷新越慢. LaCT 用窗口 attention 补块内依赖,TTT-E2E 用固定更新窗口承担同样取舍. 块大小因此既是系统参数,也是模型能看到什么历史的语义参数.

跨序列复位必须进入服务接口. batch 中若包含多条独立文档,每条都要有自己的 W0W_0 与状态 mask;拼接训练时 EOS 后是否清空,直接决定信息能否越过文档. 不复位可能降低连续语料损失,也会造成不同用户或样本间泄漏. 论文中的长文档实验与多会话服务不能默认使用同一生命周期.

外循环训练还要覆盖将来内循环会遇到的长度和更新次数. 只在短序列上元学习 W0W_0,推理时运行百万 token,快速权重可能漂出训练分布. 长度外推不仅考验位置编码,也考验优化动力学的稳定性. 应画状态范数、梯度、检索与损失随 token 的曲线,检查性能提升是否伴随状态逐渐饱和.

进入 5.1 时,关键问题是梯度步怎样退化为 delta rule,以及 mini-batch dual form 怎样恢复并行;进入 5.2 时,关键问题转为更新块多大、损失放在哪里、外循环怎样穿过内循环;进入 5.3 时,再把这些选择映射为统一的状态、写入、读出与生命周期. 三个视角共同回答「上下文怎样成为可计算的记忆」.

3 方法对照

3.1 写入规则与代价

方法 状态 写入 每次更新的块 主要代价
线性注意力, DeltaNet d×dd\times d 矩阵 加外积, 或按误差擦后写 1 个 token (可分块并行) 状态小, 装满后互相干扰
TTT-Linear / TTT-MLP 线性模型或 MLP 的权重 一步重建损失的梯度 16 个 token 硬件利用率低
Titans 长期记忆 MLP 权重 梯度加动量加遗忘 分块并行 主架构仍要注意力
LaCT SwiGLU 权重, 可达参数的 40% 大块梯度, Muon, L2 归一化 2K 到 1M 个 token 块内顺序交给窗口注意力
TTT-E2E 后 1/4 块 MLP 权重 一步下一个 token 预测梯度 1K 个 token 训练要二阶梯度, 8K 时慢 3.4 倍

表中从上到下, 状态越来越大, 写入越来越接近普通训练. 5.1 引用的 125M 消融从线性注意力出发, 逐项加入可学习的 W0W_0, ff 内的 LN 与残差, mini-batch, 可学习学习率和 Mamba backbone, 困惑度从 15.91 降到 11.09, 其中从 batch GD 改到 b=16b=16 的贡献最大, 值 1.7 个困惑度. 块越小梯度步越多, 效果越好, 但并行度越低; LaCT 选择相反的方向, 用大块换硬件效率, 再用窗口注意力补回块内的结构.

表里的系统代价落在不同位置. TTT-Linear 在 TPU 上 2K 上下文训练一步 0.27 秒, Transformer 0.30 秒; LaCT 的 Muon 版吞吐低约 14%; TTT-E2E 在 128K 时每个 token 的处理时间比全注意力快 2.7 倍, 但短上下文训练更慢, 而训练算力大部分花在短上下文预训练上. 这些数字分属训练与推理, 上下文长度也不同, 不能直接并成一条快慢排名.

3.2 压缩与检索

这些方法在困惑度上可以追平全注意力, 在精确检索上仍然落后. 3B 模型从 8K 扩到 128K, 其他方法与全注意力的损失差随长度变大, TTT-E2E 的差基本不变且为负; 同一个模型在 RULER 的 S-NIAH 任务上, 全注意力大幅领先. Gated DeltaNet 的 S-NIAH 案例也显示, 不同写入规则在「长期保持」和「装满后的区分」上各有短板, 只是都在固定状态内部比较.

所以平均预测质量和精确检索是两种能力. 检索, 复制和在长文里定位原句依赖精确回看, 这也是混合架构普遍保留注意力的原因: LaCT 本身就是窗口注意力加 TTT, Titans 的主架构里有注意力, Mamba-2 和 Gated DeltaNet 在大规模模型里也常与滑动窗口注意力混合. 两篇的正面结果主要落在语言建模损失; 下游任务, 推理和指令跟随没有测, 因而「与全注意力相当」只覆盖论文实际报告的指标.

4 内容依赖

5.1 建立本章共用的更新式: TTT-Linear 的一步更新可以写成 delta rule, mini-batch 又把状态变化分成块内无擦除与块间覆盖. 5.2 沿着两个缺口继续推进, LaCT 处理大块更新的硬件效率, TTT-E2E 则让快速权重直接承受语言模型损失. 两者仍共同面对精确检索和训练成本问题. 5.3 把这些设计放回线性注意力, Mamba 与 DeltaNet 的坐标中, 用状态生命周期解释它们为何相近, 又为何不能简单互换.

这个生命周期也是本章与 04 · TestingTime 适应 的分界. 长上下文 TTT 的快速权重服务于一条序列, 通常随序列结束而复位; 持续 TestingTime 适应的状态沿测试流积累, 必须处理跨样本泄漏, 隔离与何时重置. 如果把序列内状态延长到跨请求, 问题就从长上下文记忆滑向了持续适应, 第 4 章的约束也会随之出现.