Browse this library

Chinese source article

5.2 · 块级更新与端到端 TTT

比较 LaCT 与 TTT-E2E 如何扩大在线更新粒度,并追踪块内依赖、训练成本和精确回忆代价。

View original Markdown ↗

逐 token 更新 fast weights 在数学上自然, 在 GPU 上却容易被串行的小矩阵运算拖慢. 将若干 token 合成一块, 可以把更新变成更大的矩阵计算;块越大, 硬件利用率越高, 块内 token 越难及时看见彼此的写入. LaCT 与 TTT-E2E 都在处理这个张力, 但修改的对象不同: 前者重新安排序列内训练的计算, 后者把更新推到普通 Transformer 的后部 MLP 并用语言模型目标端到端训练.

设块 Bb=(xs,…,xs+m−1)B_b=(x_s,\dots,x_{s+m-1}). 一种并行近似是在块开始的 WsW_s 上计算所有梯度,

Ws+m=Ws−η∑i=ss+m−1∇Wℓi(Ws).W_{s+m}=W_s-\eta\sum_{i=s}^{s+m-1}\nabla_W\ell_i(W_s).

它保留块间顺序, 块内每个 token 的梯度都看不到前面 token 已造成的状态变化. 精确顺序更新则让第 ii 个梯度依赖 WiW_i, 串行路径变长. 块大小既是性能参数, 也是模型能看到何种因果关系的建模参数.

1. LaCT 如何把块做大

LaCT 将测试时学习组织成长块, 并用窗口 attention 补充块内局部依赖. 对视频, 帧和 patch 具有较强局部冗余, 大块更新可以汇总长范围统计;对语言, token 顺序更敏感, 将块当集合会损失语序. 因而同一种并行策略在不同模态上的适用性不同.

更新算法也影响吞吐. 普通梯度下降结构简单, Muon 等矩阵优化器可能改善大块更新的条件, 同时增加正交化或矩阵运算成本. 论文中的 FLOPs、训练速度和最终 loss 要一起看. 算术量下降但 kernel 利用差, 实际墙钟时间仍可能更长.

窗口必须覆盖块内盲区. 如果窗口小于块, 同块较早 token 既没有写入当前状态, 也超出局部 attention, 信息会暂时不可见. 扩大窗口减少盲区, KV cache 与二次计算又回来. 合理比较应扫描块大小与窗口组合, 而不是只给一个最优配置.

2. TTT-E2E 改变写入目标

早期 TTT layer 在每层构造键值重建损失, 它是语言模型目标的代理. TTT-E2E 直接更新 Transformer 后部的 MLP, 使用下一个 token 预测 loss 作为内循环信号. 写入目标与最终任务一致, 减少代理错位;外循环需要穿过内循环更新, 二阶梯度和激活保存使训练更昂贵.

普通 MLP 承载预训练知识, 测试时直接更新会覆盖原函数. 一种保护方式是在可更新分支旁保留静态分支, 让输出同时包含稳定底座和快速适应. 这与持续学习的稳定—可塑性分工相似, 只是状态在每条序列结束后重置. 若跨序列保留这些更新, 现有训练证据不能保证长期稳定.

训练窗口和推理窗口不能混用. 外循环只在有限长度上反向传播, 更长推理依赖更新规则外推. 128K 上的结果可以显示外推表现, 不等于训练已经给每个远距离依赖分配了完整信用. 改变训练截断长度并测长程任务, 才能看到收益来自规则还是更长训练.

3. 固定状态的能力上限

状态参数量给出存储字节, 不直接等于可用记忆. 多个键映射到相似方向时会碰撞, 新残差覆盖旧关联;非线性 MLP 可以形成更复杂地址, 仍受参数和数值精度限制. 精确复制、实体检索与语言 loss 对状态的要求不同, 需要分别测量.

S-NIAH 一类检索曲线能暴露距离和干扰, 但合成键的分布可能比自然语言容易. 可以控制相似键、相反值、重复写入和覆盖顺序, 观察模型是否按预期更新. 清空 fast state、替换为另一序列状态、冻结写入, 则能确认改进确实来自测试时学习.

状态大小、窗口和 backbone 参数要分开扫描. 总模型变大时三者可能同时增长, loss 改善无法归因. 固定慢参数改变 fast state, 固定状态改变 backbone, 再固定总 FLOPs 对比 attention, 才能画出真正的容量—成本前沿.

4. 如何比较 LaCT 与 TTT-E2E

LaCT 主要改变更新粒度与硬件排布, TTT-E2E 主要改变内循环目标和可更新模块. 两者都使用窗口 attention 处理局部关系, 都在序列边界复位, 也都用更多训练复杂度换取长序列推理效率. 它们的正面结果集中在语言建模、视频建模和特定检索任务, 不能直接外推到指令遵循与复杂推理.

等预算比较至少固定总参数、训练 token、训练 FLOPs、状态字节、窗口和推理长度. 短上下文可能由 full attention 更便宜, 长度超过交叉点后固定状态才显出优势. 报告整条长度—延迟—loss 曲线, 比单个最大长度更能说明方法何时值得使用.

LaCT 与 TTT-E2E 会进一步展开更新掩码、梯度路径和论文结果. 下一条序列模型边界将这些方法与 attention、SSM、RWKV 和 Titans 放到统一的读写坐标中, 看清固定状态、显式缓存和参数记忆各自保存什么.

5. 更新边界本身可以成为模型偏置

固定块切分会让相邻 token 因是否跨界而走不同路径. 同一句话平移几个位置, 某个依赖可能从窗口内变成跨块状态, 输出随之变化. 训练时随机化块起点、使用重叠窗口或多尺度块可以减弱边界效应, 也增加重复计算. 评测应专门移动关键依赖相对边界的位置.

自然文档具有段落、镜头和会话边界, 它们可以为状态更新提供更合理的分块. 语义分块需要额外检测, 错误边界会把相关内容拆开;固定长度简单稳定, 未利用结构. 两者的比较要固定平均块大小与窗口, 否则语义分块可能只是获得更多 token 或计算.

块级 checkpoint 也与算法状态相关. 若只在块末保存 WW, 中途故障需要重算整块;若保存块内位置与优化器状态, 存储和迁移成本上升. 这里关注块大小对最大重算量和状态快照粒度的共同影响, 不展开通用容灾. 选择大块提高吞吐时, 这项代价应进入端到端成本.

块级近似还有一个容易忽略的梯度差异: 同块并行梯度都在块首参数处计算, 精确顺序梯度则在不断更新的参数处计算. 二者只有在步长小、曲率低时接近. 扫描学习率与块大小, 测量最终状态差和输出差, 可以判断吞吐优化何时开始改变算法语义.

多尺度块提供另一种折中. 小块处理局部快速变化, 大块汇总长期统计, 不同状态用不同更新频率. 它增加状态与调度复杂度, 但比单一块同时承担局部因果和长程吞吐更自然. 分别关闭各尺度, 才能确认它们没有重复学习同一信息.

等状态字节的单尺度基线能够判断收益来自多时间尺度, 还是仅仅来自增加了更多状态容量.

块边界随机化后的方差也要报告, 防止只展示一次有利切分.

语言和视频各自的最优块大小也应分别寻找, 不能直接沿用.

因此,块级更新的比较不能只看最终精度,还要同时记录更新次数、每次更新覆盖的上下文和由此引入的推理开销。

参考文献