In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升

关键词测试时训练原位更新大语言模型长上下文

“静态的‘训练后部署’范式,从根本上限制了大语言模型在推理时根据新信息动态调整权重的能力。”这是来自字节跳动 Seed 团队与北京大学联合发表的论文《In-Place Test-Time Training》中的核心论断。

当 GPT-4、Llama 3 等大模型在部署后就“凝固”为静态的知识库,它们便无法像人类一样从连续的上下文流中持续学习——这不仅是认知上的遗憾,更是长上下文任务处理的硬伤。

In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升

为了解决这一困局,研究团队提出了一种名为 In-Place Test-Time Training(原位测试时训练)的全新框架。它的本质思路大胆而优雅:不再引入新结构,而是直接“征用”Transformer 中无处不在的 MLP 模块的最终投影矩阵作为“快速权重”,在推理时原地更新。

这意味着,任何预训练好的大模型都可以像安装“插件”一样,在不改变架构、无需昂贵重新训练的前提下,获得边推理边学习的能力。

效果如何?实验数据显示:

  • In-Place TTT 让 Qwen3-4B 模型在 128k 长上下文任务上的准确率从 74.8% 跃升至 77.0%。 当模型规模扩大到 14B 时,优势依旧稳健。
  • 更令人兴奋的是,该方法在 32k 上下文长度下的预训练中,将传统 Transformer 在 RULER-16k 测试中的得分从 6.58 直接拉升至 19.99。

这篇工作不仅是 Test-Time Training 在大模型时代的里程碑式突破,更为我们勾勒了一幅模型“持续进化”的未来图景。

本文目录

  • 一、静态模型的“阿喀琉斯之踵”
  • 二、In-Place TTT:三管齐下的破局之道
    • 2.1 核心洞察:征用 MLP,而非替换 Attention
    • 2.2 大块更新:让并行计算真正“吃饱”
    • 2.3 语言建模对齐目标:从“记住当下”到“预见未来”
    • 2.4 理论保证:为什么 LM 对齐目标更优?
  • 三、工程实现:上下文并行与因果性保证
    • 3.1 并行扫描:把“串行”变成“并行”
    • 3.2 因果性与边界处理
  • 四、实验:用数据说话
    • 4.1 实验一:让 Qwen3-4B“再进化”
    • 4.2 实验二:从零开始,与 TTT 家族正面交锋
    • 4.3 消融实验:拆解每一个设计选择
  • 五、相关工作:站在巨人的肩膀上
    • 5.1 Test-Time Training 的演进
    • 5.2 高效长上下文架构
    • 5.3 记忆增强神经网络
  • 六、结语:迈向“持续学习”的新范式

In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升

一、静态模型的“阿喀琉斯之踵”

当前大语言模型的成功,建立在一种经典的“训练-部署”范式之上:模型在海量语料上完成预训练后,其权重便被“冻结”,部署时只做前向推理。这种做法的好处是稳定、高效,但问题也同样尖锐——模型失去了根据上下文动态调整内部表征的能力。

当你正在让模型处理一篇长达 10 万字的专业报告,在报告第 2 页出现了一个关键术语的定义,而在第 80 页需要基于这个定义进行推理。传统模型的注意力机制必须将这跨越数万 token 的依赖关系完整保留在 KV Cache 中,这既是对显存的巨大消耗,也对长距离信息检索提出了严苛要求。

当然,In-Context Learning(上下文学习)提供了一种补救思路——把历史信息全塞进上下文窗口。但这不过是把问题转移到了另一个维度:注意力机制的二次复杂度使得无限扩展上下文窗口在计算上不可行。

正是在这样的背景下,Test-Time Training 进入了研究者的视野。TTT 的核心理念是:在推理时,动态更新模型的一小部分参数(称为“快速权重”),让这些参数成为上下文信息的“压缩存储器”。这样一来,模型就不必事无巨细地记住每一个 token,而是将关键信息提炼进权重中。

然而,将 TTT 应用于现代 LLM 面临着三座难以逾越的大山:

  1. 架构不兼容:已有的 TTT 方法大多需要引入专门的循环层来替代注意力机制,这意味着必须“从零开始”预训练,无法复用现有的大模型资产;
  2. 计算效率低下:经典的 TTT 采用逐 token 更新策略,这严重限制了 GPU/TPU 的并行能力;
  3. 学习目标错位:现有 TTT 普遍采用“重构”作为自监督目标——让模型记住当前 token 的表征,但这与语言模型“预测下一个 token”的根本目标并不一致。

In-Place TTT 的诞生,正是为了同时搬走这三座大山。

二、In-Place TTT:三管齐下的破局之道

2.1 核心洞察:征用 MLP,而非替换 Attention

In-Place TTT 最关键的洞见在于:TTT 的“快速权重”并不必须是某个新模块的参数,它可以“寄生”在已有模块之中。

回顾 Transformer 的结构,多层感知机(MLP)块占据了模型参数的大头。更重要的是,已有研究表明,Transformer 中的 MLP 层本质上扮演着一种键值记忆的角色——它存储着预训练阶段学到的海量通用知识(作为“慢速权重”)。

那么,为什么不让这同一套参数系统,在推理时也承担起存储上下文信息(作为“快速权重”)的职责呢?

具体来说,In-Place TTT 选择了门控 MLP 结构中的最终投影矩阵作为被更新的对象。门控 MLP 的计算公式如下:

MLP(x) = (SiLU(xW_g) ⊙ (xW_u)) W_o

其中,x 是输入隐状态,SiLU 是激活函数,W_gW_u 是输入投影矩阵,W_o 是输出投影矩阵。

在 In-Place TTT 框架下,W_gW_u 保持冻结(作为慢速权重),而 W_o 则被征用为快速权重,在推理时动态更新。

这种“原位”设计的精妙之处在于:

2.2 大块更新:释放并行计算潜力

传统TTT方法通常采用逐token的串行更新机制:为了处理第 i 个token,必须等待前 i-1 个token的更新完成。这种强序列依赖性严重限制了GPU等硬件的并行计算能力。

In-Place TTT通过架构设计巧妙地规避了这一问题。由于注意力机制被完整保留,负责细粒度token间的信息交互,快速权重的更新便可以采用一种更“粗放”的策略:分块更新。具体流程如下:

In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升
图1:In-Place TTT整体框架。模块顺序处理输入块。对每个块,当前快速权重首先作用于中间激活值,产生输出;随后,这些权重利用当前块的激活值和从token嵌入派生的目标值进行更新。这种“先应用,后更新”的循环,使模型能以严格因果的方式动态适应输入上下文。

对于一个长度为 L 的序列,首先将中间激活值 A 和目标值 V 划分为 K 个大小为 B 的不重叠块。对于第 k 个块:
1. 应用操作:使用当前快速权重 W_k 处理该块,得到输出。
2. 更新操作:基于该块的激活值与目标值计算梯度,并就地更新快速权重至 W_{k+1}

这种块级更新策略具有双重优势:
* 并行友好:块内的所有token可以同时被处理,极大提升了硬件利用率。
* 支持大块:由于注意力层依然负责token混合,In-Place TTT无需依赖极小的块大小来保证信息交互质量。实验表明,块大小 B 为512或1024时效果最佳(参见下文图3b)。

In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升
图3:针对In-Place TTT关键设计选择的消融实验,使用17亿参数模型在RULER基准上评估。图表展示了:(a) 状态维度对性能的影响;(b) 分块大小的性能权衡,中等大小(如512、1024)为最优;(c) 目标函数中卷积与投影操作的必要性。

2.3 语言建模对齐目标:从记忆到预测

在解决了架构和效率问题后,核心问题在于:快速权重应该学习什么?

现有TTT方法普遍采用重构目标,即让模型学习将当前输入的某个表示映射回其自身或相关值。这本质上是在训练模型“记住”已见内容。

In-Place TTT团队认为,对于语言建模任务,一个更有效的目标是让快速权重存储对未来预测有用的信息。为此,他们提出了一个新的语言建模对齐目标

V = Proj(Conv(E))

其中,E 是token嵌入矩阵,Conv 是一维卷积操作,Proj 是可训练的投影矩阵。该设计的关键在于:
* 卷积操作:使目标值 V 能够包含未来token的信息,实现“预见未来”。
* 投影矩阵:允许模型学习从嵌入空间到目标空间的最优映射。

例如,若目标是预测下一个token,可将卷积核设置为在下一token位置为1,其余为0,并将 Proj 设为单位矩阵。

在此目标下,损失函数采用简单的相似性度量,并推导出快速权重的更新规则。该规则不仅计算高效,更重要的是具备结合律特性,为后续实现完全的上下文并行提供了理论基础。

2.4 理论分析:为何预测目标更优?

为了从理论上验证语言建模对齐目标的优越性,研究团队在经典的归纳头情景下进行了分析。

情景设定:序列中先出现一个键值对 (k, v)(例如“奥巴马 -> 米歇尔”),之后在查询位置再次出现键 k(“奥巴马”),模型需要预测出值 v(“米歇尔”)。这是语言模型实现上下文学习的关键机制之一。

基于两个合理假设(不同token的嵌入近似正交;相同token在不同位置的中间激活具有正相关性),定理1 量化比较了两种目标下,正确token的logit(模型输出概率的对数)的期望增益:

| 目标类型 | 正确token的logit期望增益 |
| :— | :— |
| 语言建模对齐目标 | 正比于 |v|^2 * 相关性 |
| 重构目标 | 正比于极小的常数 ε |

直观解释
* 语言建模对齐目标(预测未来):当首次看到“奥巴马”时,模型快速权重存储的是其后续“米歇尔”的信息。当再次遇到“奥巴马”时,由于两处激活相似,存储的信息直接用于提升“米歇尔”的logit,增益显著。
* 重构目标(记忆当下):快速权重存储的是“奥巴马”自身的信息。再次遇到“奥巴马”时,取出的是“奥巴马”的特征,这与需要预测的“米歇尔”的特征几乎正交(内积极小),因此对正确预测的贡献微乎其微。

理论分析表明,让快速权重“预见未来”比单纯“记忆当下”能更有效地提升语言建模性能。

3.2 因果性与边界处理

在追求高效并行计算的同时,In-Place TTT 严格保持了模型的因果性,确保其行为与串行处理完全一致。这主要通过以下三个关键设计实现:

  • 因果填充:在生成目标值时,1D 卷积层采用因果填充模式。这意味着每个位置的计算仅依赖于当前及之前的 token,严格防止了“未来信息”的泄露。
  • 边界重置:在处理文档边界时,快速权重会被重置为预训练初始值。这一机制有效防止了不同文档间的信息污染,确保了任务间的独立性。
  • 数学等价:整个并行更新算法在数学上被证明与严格串行的逐块更新过程完全等价,为框架的高效与正确性提供了理论保障。

四、实验验证

实验部分从多个维度系统性地验证了 In-Place TTT 框架的有效性、通用性与核心设计的重要性,主要包括:
1. 作为“即插即用”模块对现有预训练模型的增强效果。
2. 从零开始预训练时,与主流高效注意力及测试时训练方法的对比。
3. 通过消融实验分析各设计选择的作用。

4.1 实验一:即插即用增强预训练模型

本实验以 Qwen3-4B-Base 模型为基座,采用两阶段继续训练策略:首先进行 20B token 的 32k 上下文训练,随后进行 15B token 的 128k 上下文训练(配合 YaRN 位置编码进行扩展)。对比基线为采用相同训练配置但未加入 In-Place TTT 模块的模型。

In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升
表 1:RULER 基准测试结果,报告平均准确率(%),最优结果以粗体标注。

如表 1 所示,在短上下文(4k-32k)范围内,增强版与基线模型表现接近。然而,随着上下文长度扩展至 64k 和 128k,In-Place TTT 带来的性能优势持续扩大,分别高出基线 4.4 和 2.2 个百分点。尤为重要的是,在 256k 的外推测试场景(超出训练长度)下,增强版模型仍保持 2.2 个百分点的领先优势。这表明快速权重学习到的是一种可泛化的上下文信息压缩与利用能力,而非对特定长度的过拟合。

为验证方法的通用性,研究进一步将 In-Place TTT 应用于 LLaMA-3.1-8B 和 Qwen3-14B-Base 模型。

In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升
表 2:在 RULER 基准上,将 In-Place TTT 扩展至 LLaMA-3.1-8B 与 Qwen3-14B-Base 的结果。

如表 2 所示,在不同模型家族和参数量级上,In-Place TTT 均带来了全上下文长度范围内的性能提升,尤其在 64k 长上下文下增益显著。这有力地证明了该框架是一种不依赖于特定模型架构的普适性长上下文增强方案。

4.2 实验二:从零预训练与主流方法对比

本实验在 500M 和 1.5B 参数量规模上,从零开始预训练模型,并与多种主流高效注意力及测试时训练方法进行对比,包括:滑动窗口注意力(SWA)、门控线性注意力(GLA)、DeltaNet 以及另一种 TTT 变体 LaCT。所有模型均在 32k 上下文长度下进行训练。

In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升
图 2:在 Pile 数据集上,500M(左)与 1.5B(右)参数模型在不同上下文长度的滑动窗口困惑度。

如图 2 所示,In-Place TTT 在两个参数量级上的困惑度始终低于其他基线方法,并且其性能随着上下文长度的增加持续改善直至 32k。这表明其动态压缩机制能够有效利用更长的历史信息,而其他方法在达到一定长度后性能趋于饱和。

在 4B 规模模型的进一步实验中,结果展示了更全面的能力提升。

In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升
表 3:4B 模型在常识推理与长上下文评估基准的测试结果。

表 3 揭示了几个关键发现:
1. 基础能力无损:加入 In-Place TTT 后,模型在 HellaSwag、MMLU 等常识推理任务上的表现普遍略有提升,说明动态更新机制并未扰乱预训练知识,反而可能通过改善上下文理解产生了正向作用。
2. 长上下文能力飞跃:使用全注意力机制的模型在 RULER-16k 任务上的得分从 6.58 大幅提升至 19.99;使用滑动窗口注意力的模型在 RULER-8k 上的得分也从 9.91 提升至 26.80。
3. 效率与效果的平衡:一个有趣的观察是,“滑动窗口注意力 + In-Place TTT”的组合在 RULER-8k 上的表现(26.80)甚至超越了未增强的“全注意力”模型(9.91),这为在有限计算资源下实现高效的长上下文处理提供了有前景的路径。

4.3 消融实验:拆解设计选择

为验证各核心组件的必要性,研究以 1.7B 模型在 RULER 基准上进行了系统的消融实验。

状态规模的影响
In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升
图 3a:状态规模消融实验。
随着更多层的 MLP 被征用为快速权重(即状态规模增大),模型在不同上下文长度下的性能一致提升,验证了“更大的快速权重容量能存储更多上下文信息”的直觉。

块大小的影响
In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升
图 3b:块大小消融实验。
对比块大小(C)为 128、256、512、1024 的配置,发现 C=512 和 1024 在长上下文(16k)上表现最优。较小的块(如 C=128)在短上下文上略有优势,但在长上下文上性能明显下降,表明 In-Place TTT 适合采用较大的块以兼顾效率与效果。

语言模型对齐目标的关键组件
In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升
图 3c:LM 对齐目标消融实验。
对比了无卷积无投影、仅有投影、仅有卷积、以及卷积与投影结合(完整版)四种配置。结果显示:卷积操作对于利用长上下文信息至关重要,而投影层则对短上下文性能有重要贡献。两者结合才能在所有上下文长度上达到最优性能。

In-Place Test-Time Training:让大语言模型在推理时原地进化,长上下文任务准确率显著提升
图 4:吞吐量与内存效率。 对比了带/不带 In-Place TTT 的 4B 模型在预填充阶段的吞吐量(a、b)和峰值内存(c、d)。结果表明,In-Place TTT 引入的额外开销非常有限——在 128k 上下文长度下,吞吐量下降不到 10%,内存增加不到 5%。考虑到模型长上下文处理能力的显著提升,这一性能代价几乎可以忽略。

五、相关工作:站在巨人的肩膀上

In-Place TTT 并非凭空产生,它建立在多条研究脉络的基础之上。

5.1 测试时训练的演进

测试时训练(TTT)的概念最早在计算机视觉领域被提出,旨在通过自监督任务(如图像旋转预测)在测试时更新模型部分参数,以应对数据分布偏移。随后,这一思想被引入语言领域,TTT-LM 等工作尝试采用类似机制来增强语言模型的上下文处理能力。

然而,早期的 TTT 实现存在两个主要瓶颈:
1. 逐 Token 更新的串行性:严重限制了计算并行效率。
2. 重构目标的局限性:与语言建模的因果预测本质不完全匹配。

后续研究如 Titans 和 LaCT 开始探索分块更新策略,部分缓解了效率问题。但在这些工作中,TTT 模块通常被设计为注意力机制的替代品,这导致它们必须使用较小的处理块以保证 Token 间的信息交互质量,且难以直接复用现有的预训练模型架构。

In-Place TTT 的关键突破在于,将 TTT 从“替代者”转变为“增强者”。通过征用模型中原有的 MLP 层而非替换注意力模块,它既保留了预训练模型的核心能力,又获得了使用更大分块进行高效更新的自由度。

5.2 高效长上下文架构

另一条密切相关的研究脉络是高效长上下文架构的设计,主要包括三大方向:
* 稀疏注意力:如 Longformer、Sparse Transformer,通过限制注意力范围(如滑动窗口、扩张窗口)来降低计算复杂度。优点是简单直接,缺点在于可能牺牲了全局的交互能力。
* 线性注意力与状态空间模型:如 Linear Transformer、Mamba、GLA,将注意力近似为核函数形式,或使用结构化状态空间来压缩历史信息。优点是实现了线性计算复杂度,缺点是在某些任务上的表达能力可能不及标准注意力。
* 基于 Delta 规则的记忆增强:如 DeltaNet,将快速权重更新形式化为 Delta 规则,实现了高效的并行化处理。

In-Place TTT 与这些工作并非竞争关系,而是互补关系。事实上,In-Place TTT 可以自然地集成到任何包含 MLP 层的架构中——无论是标准 Transformer、其滑动窗口变体,还是 Mamba 类的状态空间模型。这种“架构无关”的特性使其具有广泛的适用前景。

5.3 记忆增强神经网络

从更宏观的视角看,TTT 可以被理解为一种特殊的记忆增强机制——将模型参数本身作为动态记忆的载体。

这与 Neural Turing Machines、Memory Networks 等经典工作一脉相承,也与近期为 LLM 增加外部记忆(如 MemoryLLM)的尝试共享着相似的设计精神。

区别在于:In-Place TTT 的记忆是内嵌于模型权重之中的,而非外挂的独立记忆库。这使其实现更加轻量,并能与模型原有的计算流程无缝融合。

六、结语:迈向“持续学习”的新范式

In-Place TTT 的出现,为大语言模型从“静态部署”走向“动态进化”提供了新的可能性。其贡献可以总结为三个层面:
* 方法层面:提出了一种“原位”TTT 设计,无需改动模型架构或重新训练,即可使现有预训练 LLM 获得在推理时进行学习的能力。
* 理论层面:首次从归纳头的角度分析了不同 TTT 训练目标对语言建模的影响,为后续研究提供了可参考的理论框架。
* 工程层面:实现了与上下文并行计算兼容的高效分块更新算法,在几乎不影响推理效率的前提下,显著提升了模型的长上下文表现。

当然,这项工作也引出了若干值得进一步探索的问题:
* 优化器的选择:当前仅使用最简单的单步梯度下降,更复杂的优化器(如带动量的更新)是否能带来进一步的性能提升?
* 遗忘机制:当前的快速权重更新只会“累加”而不会“遗忘”,这是否会在处理极长序列(如百万 Token 级别)时导致信息过载?
* 多模态扩展:这种“原位 TTT”的设计思想能否成功迁移到视觉-语言等多模态模型中?

正如论文所言,这项工作“向 LLM 的持续学习新范式迈出了充满希望的一步”。在可预见的未来,一个能够边推理边学习、并与用户及环境持续交互进化的模型,或许将不再遥不可及。

对于 AI 基础设施的从业者而言,In-Place TTT 也预示着一个有趣的趋势:未来的推理系统可能需要支持“带状态的推理”——模型权重在推理过程中不再是只读的,而是会进行局部、动态的更新。 这对推理引擎的设计、分布式部署策略乃至硬件架构,都提出了新的挑战与机遇。


关注“鲸栖”小程序,掌握最新AI资讯

本文来自网络搜集,不代表鲸林向海立场,如有侵权,联系删除。转载请注明出处:https://www.itsolotime.com/archives/29482

(0)
上一篇 13小时前
下一篇 13小时前

相关推荐