Rethinking Expressivity and Efficiency in Test-Time Training¶
会议: NeurIPS 2026
arXiv: 2608.21308
OpenReview: https://openreview.net/forum?id=s9Npo5C1Sz
代码: https://github.com/zeyun-zhong/E2-TTT
领域: 其他
关键词: 测试时训练 (TTT)、长上下文扩展、闭式标量核、快速权重记忆、混合注意力架构
一句话总结¶
针对测试时训练（TTT）在表达能力与硬件效率之间的两难权衡，本文提出 E2-TTT，在固定块首权重的梯度近似下推导闭式标量核，以单次反向传播精确重现逐 token 动量与衰减动力学，在保持高吞吐的同时实现 8 倍上下文长度外推。
研究背景与动机¶
大语言模型在长程智能体与复杂时序任务中面临着随上下文长度爆炸的内存与计算开销，标准 Transformer 维护的键值缓存（KV Cache）使得服务长序列的代价极其高昂。为了摆脱二次复杂度依赖，紧凑状态模型（包括线性注意力、状态空间模型以及循环混合架构）尝试将历史序列压缩至固定尺寸的隐状态中。在这条路线上，测试时训练（Test-Time Training, TTT）将循环状态重构为内部神经网络的“快速权重”（fast weights），在推理阶段通过自监督损失执行在线梯度更新，展现出比固定维度向量更优越的记忆容量与表达上限。
然而，现有的测试时训练范式在更新粒度上陷入了表达能力与硬件利用率的根本矛盾。标准的逐 token 更新（Token-wise TTT）具备精细的在线梯度下降动力学，能够根据输入动态调节学习率、动量和权重衰减，但其严格串行的内部循环导致 GPU 计算利用率通常低于 5%，无法在现代硬件上规模化扩展。相反，最新的分块方法（如 LaCT）将长序列切分为大块（chunk），累加块内损失并执行单次更新，大幅提升了硬件吞吐；但这种做法强行将块内动量和衰减因子折算为均值，抹杀了块内不同 token 的时序重要性差异，导致模型在真实检索和长文本长度外推时性能断崖式崩溃。
既然分块并行计算的吞吐优势不可放弃，而逐 token 的时序优化动态又是长程泛化的关键，是否能在不具体物化块内密集矩阵状态的前提下，在数学上严格保全逐 token 的状态转移？核心 idea：在冻结块首权重求导的标准假设下，通过重排时间步累积求和导出闭式标量核，将逐 token 的动量与解耦衰减动力学压缩为双重标量加权聚合，使大块并行更新在单次反向传播中严格等价于逐 token 循环的块尾状态。
方法详解¶
整体框架¶
E2-TTT 采用分块混合宏观架构，将输入的 token 序列划分为大小为 \(C\)（如 \(C=512\)）的非重叠块。在模型宏观层面上，输入隐状态经过投影得到共享的查询（\(\mathbf{Q}\)）、键（\(\mathbf{K}\)）与值（\(\mathbf{V}\)），并分流至两个并行分支：滑动窗口注意力（Sliding Window Attention, SWA）负责建模块内精细的局部因果交互，而 E2-TTT 状态分支则通过快速权重矩阵跨块压缩和传递长期历史上下文。最后，数据驱动的输入相关门控网络将两路分支的输出进行逐元素融合。
在 E2-TTT 分支内部，模型在每个块的前向阶段使用上一块传递下来的块首冻结权重 \(\mathbf{W}_0\) 计算块内所有 token 的输出与重构损失，完全消除块内串行依赖；在反向传播与状态更新阶段，轻量预测头从输入特征中实时预测每个 token 的标量学习率 \(\eta_t\)、动量因子 \(\beta_t\) 和衰减强度 \(\alpha_t\)，利用闭式标量核直接计算块尾快速权重 \(\mathbf{W}_C\) 与动量状态 \(\mathbf{M}_C\)，随后传递至下一块。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入序列 X 分块 (长度 C)"] --> B["共享 Q, K, V 投影与归一化"]
B --> C["局部滑动窗口注意力 (SWA)"]
B --> D["逐 token 动态参数预测<br/>学习率 η, 动量 β, 衰减 α"]
B --> E["块首权重前向推理与自监督损失<br/>o_t = f_W0(q_t), L_t = L(f_W0(k_t), v_t)"]
D --> F["闭式标量核并行扫描<br/>对数空间累积后缀积与比值和 R_t"]
E --> G["单次反向传播共享梯度激活 g_t"]
F --> H["双核重加权梯度聚合<br/>ΔW 与 ΔM 计算"]
G --> H
H --> I["精确跳跃更新块尾状态<br/>(W_C, M_C) 传递至下一块"]
C --> J["数据相关门控融合与输出"]
E --> J
J --> K["下层网络输入"]关键设计¶
1. 逐 token 动态参数预测：赋予定制化的记忆保留与遗忘粒度
传统分块 TTT 强行将优化器超参数设为常量或块级均值，无法感知关键检索词与无关填充词在记忆价值上的巨大悬殊。E2-TTT 针对这一痛点，通过端到端学习的轻量线性头，使优化动力学具备精细的输入自适应能力： - 学习率 \(\eta_t = \eta_{\text{base}} \cdot \sigma(\text{Linear}_\eta(\bm{x}_t))\) 动态调节当前 token 写入记忆的更新步长； - 动量因子 \(\beta_t = \sigma(\text{Linear}_\beta(\bm{x}_t))^{1/\tau}\) 控制历史梯度的保留惯性，其中时间尺度标量 \(\tau\)（默认 32）将数值偏置在接近 1 的区间，显著增强长程记忆维持； - 权重衰减强度 \(\alpha_t = \alpha_{\text{base}} \cdot \sigma(\text{Linear}_\alpha(\bm{x}_t))\) 充当自适应遗忘门，并与学习率解耦绑定形成乘性衰减因子 \(\gamma_t = 1 - \eta_t \alpha_t\)。当某 token 的更新步长极大时，其遗忘机制被同步强化，防止历史快速权重在长期累积中发生范数爆炸。
2. 闭式标量核推导：跳过中间密集矩阵物化的精确状态转移
逐 token 循环展开式 \(\mathbf{M}_t = \beta_t \mathbf{M}_{t-1} + \eta_t \mathbf{G}_t\) 与 \(\mathbf{W}_t = \gamma_t \mathbf{W}_{t-1} + \mathbf{M}_t\) 若在块内逐步物化，需要存储 \(O(C \cdot d^2)\) 的密集矩阵，在显存和带宽上不可承受。E2-TTT 抓住“矩阵状态昂贵但衰减和动量因子为纯标量”的代数特征，通过时间步代入与和式指标交换，将块尾状态完全解耦为初始状态结转项与各时间步梯度的标量加权和： 定义后缀累积积与比值和： $\(\widetilde{\beta}_t := \prod_{i=t+1}^{C} \beta_i, \quad \widetilde{\gamma}_t := \prod_{i=t+1}^{C} \gamma_i, \quad R_t := \sum_{i=t}^{C} \frac{\widetilde{\gamma}_i}{\widetilde{\beta}_i}\)$ 据此导出两套精准刻画梯度时序贡献的标量核： $\(\mathbf{M}_C = \widetilde{\beta}_0 \mathbf{M}_0 + \sum_{t=1}^{C} \mathcal{K}_t^M \mathbf{G}_t, \quad \text{其中 } \mathcal{K}_t^M := \eta_t \widetilde{\beta}_t\)$ $\(\mathbf{W}_C = \widetilde{\gamma}_0 \mathbf{W}_0 + \widetilde{\beta}_0 R_1 \mathbf{M}_0 + \sum_{t=1}^{C} \mathcal{K}_t^W \mathbf{G}_t, \quad \text{其中 } \mathcal{K}_t^W := \eta_t \widetilde{\beta}_t R_t\)$ 动量核 \(\mathcal{K}_t^M\) 刻画了负梯度 \(\mathbf{G}_t\) 经动量衰减后存活至块尾的份额；权重核 \(\mathcal{K}_t^W\) 则额外加权了比值累加和 \(R_t\)，完整涵盖了梯度注入动量缓冲池后再持续推动后续权重更新的双重累积效应。所有标量序列在对数空间中通过高效并行前缀和/后缀和计算，仅需 \(O(C)\) 开销。
3. 单次反向传播双核重加权聚合：无缝兼容分块硬件效率
计算 \(\mathbf{W}_C\) 和 \(\mathbf{M}_C\) 涉及两个不同的梯度加权和，若执行两次反向传播将严重拖慢吞吐。E2-TTT 利用链式法则分析梯度内部结构：
$\(\mathbf{G}_t = -\nabla_\mathbf{W} \mathcal{L}_t\big|_{\mathbf{W}_0} = - \left( \frac{\partial f_{\mathbf{W}_0}(\bm{k}_t)}{\partial \mathbf{W}} \right)^{\!\top} \bm{g}_t, \quad \text{其中 } \bm{g}_t := \frac{\partial \mathcal{L}_t}{\partial f_{\mathbf{W}_0}(\bm{k}_t)}\)$
由于雅可比矩阵仅取决于冻结的 \((\mathbf{W}_0, \bm{k}_t)\)，与标量核完全解耦，因此模型只需对块内总损失执行单次自动微分，获取形状为 \(C \times d\) 的逐 token 激活梯度中间张量 \(\{\bm{g}_t\}\)。随后分别利用标量核序列 \(\{\mathcal{K}_t^W\}\) 和 \(\{\mathcal{K}_t^M\}\) 进行两次收缩加权累加（Aggregate 操作），其计算复杂度保持在与传统分块 TTT 完全相同的 \(O(C \cdot d^2)\)，将额外算力开销降至微不足道的标量级别。
损失函数 / 训练策略¶
内部快速权重网络 \(f_{\mathbf{W}}\) 采用带有残差连接和 LayerNorm 的双层 GELU MLP 或三矩阵 SwiGLU MLP。在自监督内循环中，输入 token 投影的键向量 \(\bm{k}_t\) 作为输入，以值向量 \(\bm{v}_t\) 为拟合目标，采用均方误差损失（MSE）： $\(\mathcal{L}_t = \frac{1}{2} \| f_{\mathbf{W}_0}(\bm{k}_t) - \bm{v}_t \|_2^2\)$ 在外循环训练中，整个模型使用标准自回归因果语言建模交叉熵损失进行端到端优化，无需任何辅助损失监督。超参数方面，默认块大小 \(C=512\)，基准学习率 \(\eta_{\text{base}} = 0.01\)，动量时间尺度 \(\tau = 32\)，基准衰减强度 \(\alpha_{\text{base}} = 0.1\)。在每个分块边界处对快速权重矩阵执行 \(L_2\) 正则化，以确保长序列训练中的数值稳定性。
实验关键数据¶
主实验¶
模型在 HuggingFace FineWeb-Edu 15B/16B tokens 上从头预训练，序列长度为 2048/2240，模型规模分别为 340M 与 1.3B。评估涵盖通用语言建模困惑度、零样本常识推理、真实检索以及长序列外推基准（如 S-NIAH 针寻海与 LongBench 14 项任务）。
通用语言建模与真实文本检索性能对比 (1.3B 规模)
| 模型架构 | 注意力窗口 | WikiText ppl (↓) | LAMBADA ppl (↓) | 常识推理平均准确率 (↑) | 检索 SWDE (↑) | 检索 SQuAD (↑) | 检索 FDA (↑) | 检索平均 (↑) |
|---|---|---|---|---|---|---|---|---|
| Transformer++ | 2048 (全注意力) | 19.8 | 17.9 | 53.0% | 53.7% | 41.5% | 64.7% | 53.3% |
| DeltaNet | – | 20.6 | 19.9 | 52.3% | 32.9% | 29.9% | 23.6% | 28.8% |
| Mamba2 | – | 20.5 | 17.0 | 53.3% | 33.8% | 37.3% | 26.1% | 32.4% |
| HQLT (SWA+DeltaNet) | 512 | 20.3 | 16.1 | 53.8% | 32.1% | 43.4% | 31.0% | 35.5% |
| LaCT (分块 TTT) | 512 | 20.1 | 16.7 | 53.0% | 39.2% | 44.9% | 26.0% | 36.7% |
| E2-TTT (MLP) | 512 | 19.5 | 15.3 | 54.5% | 37.2% | 45.6% | 28.5% | 37.1% |
| E2-TTT (SwiGLU) | 512 | 19.9 | 15.8 | 53.6% | 47.3% | 44.4% | 39.0% | 43.6% |
消融实验¶
消融实验重点验证了优化器更新机制、闭式核相比分块平均的必要性，以及混合分支对长文本外推的关键作用（在 2K 训练长度下外推至 16K）。
更新规则与架构组件控制消融 (1.3B 规模，S-NIAH 与 LongBench)
| 模型配置 / 更新规则 | S-NIAH-1 2K (内插) | S-NIAH-1 4K (外推) | S-NIAH-1 8K (外推) | S-NIAH-1 16K (8倍外推) | S-NIAH-2 8K | LongBench 平均分 |
|---|---|---|---|---|---|---|
| 仅 SWA (无 TTT 循环分支) | 26.8% | – | 6.8% | 2.8% | – | 11.2% |
| 仅 TTT (无局部 SWA 分支) | 0.0% | 0.0% | 0.0% | 0.0% | 0.0% | 1.5% |
| 分块参数平均基线 (Chunk-averaged) | 25.2% | 10.8% | 6.8% | 3.0% | 7.2% | – |
| LaCT 对齐架构 (LaCT-Matched) | 99.8% | 99.2% | 51.6% | 0.0% | 4.8% | – |
| 完整 E2-TTT (w/o 融合门控) | 99.6% | 95.0% | 88.8% | 86.6% | – | 13.3% |
| 完整 E2-TTT (SwiGLU) | 95.6% | 95.0% | 95.4% | 93.6% | 40.6% | 14.1% |
关键发现¶
- 逐 token 动力学是长度外推的决定性支柱：将动量与衰减退化为分块平均后，模型在 8K 外推时的 S-NIAH-1 准确率从 95.4% 骤降至 6.8%，与仅用局部滑动窗口的模型无异；而 E2-TTT 在 8 倍外推长度（16K）下仍维持超过 93% 的准确率。
- 架构互补性不可或缺：单独使用 TTT 分支因缺乏块内因果路径，在 S-NIAH 上完全瘫痪（0.0%），而 SWA 分支补足了块内局部检索能力，二者结合产生了质变。
- SwiGLU 门控快速权重网络在复杂检索任务中优势显著：1.3B 规模下在三项检索基准上达到 43.6% 平均准确率，大幅领先普通 MLP 变体（37.1%）与强基线 LaCT（36.7%）。
亮点与洞察¶
- 数学解析解消除硬件权衡痛点：巧妙利用冻结块首导数这一广泛采用的工程假设，在纯标量代数层面推导出闭式解，使得逐 token 精细优化与大分块 GPU 并行计算实现了兼得。
- 单次反向传播重用激活梯度：发现不同核加权项共享同一组底层激活导数，仅需微不足道的标量乘加即可产出双重更新量，打破了“多状态更新必翻倍反向传播算力”的定势思维。
- 即插即用的多模态扩展潜力：在无需重新预训练的前提下，作为轻量旁路分支融入 Qwen3VL-2B-Instruct，仅微调 TTT 参数即在长视频理解基准上匹敌全参数微调。
局限与展望¶
- 块内依赖的因果盲区：由于块内前向传播强制共享块首权重，块内后续 token 无法通过 TTT 状态感知块内前序 token，必须依赖并行的局部注意力补偿，在纯粹无注意力序列建模上存在限制。
- 复杂优化器推广难度：当前的解析推导依赖于 SGD 动量与解耦衰减的线性可解构性，若想直接推广至带有非线性自适应二阶动量的 AdamW 或 Muon 优化器，难以直接建立完全等价的闭式标量核。
相关工作与启发¶
- vs. Token-wise TTT (Sun et al., 2024)：两者均探索逐 token 优化动力学，但前者因串行循环无法充分释放现代硬件并发算力；E2-TTT 通过标量核跳跃更新实现等价模拟，达到分块吞吐。
- vs. LaCT (Zhang et al., 2025)：LaCT 率先采用大分块来优化训练效率，但在优化器更新中粗暴抹除了时序动态；E2-TTT 填补了其在精细动量/衰减时序结构上的缺失，彻底解决了 LaCT 在长度外推时发散崩溃的问题。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 巧妙的闭式标量核推导与单次反向传播机制，漂亮地填补了理论与硬件实现的鸿沟。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 340M/1.3B 从头训练、语言建模、高压检索、16K 长度外推、消融以及多模态扩展。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨自洽，实验分析层层推进，逻辑闭环完整。
- 价值: ⭐⭐⭐⭐⭐ 为长上下文循环与混合注意力架构提供了一条极具实用价值的高吞吐理论路径。