In-Place Test-Time Training
作者: Guhao Feng, Shengjie Luo, Kai Hua, Ge Zhang, Di He, Wenhao Huang, Tianle Cai
分类: cs.LG, cs.AI, cs.CL, stat.ML
发布日期: 2026-04-07
备注: ICLR 2026 Oral Presentation; Code is released at https://github.com/ByteDance-Seed/In-Place-TTT
💡 一句话要点
提出In-Place TTT,使LLM在推理时动态适应新信息,提升长上下文任务性能。
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 测试时训练 大型语言模型 持续学习 上下文学习 快速权重
📋 核心要点
- 现有LLM的“训练后部署”模式无法动态适应真实世界任务中的新信息流,限制了其性能。
- In-Place TTT将MLP块的最终投影矩阵作为可适应的快速权重,并设计了与Next-Token-Prediction任务对齐的目标。
- 实验表明,In-Place TTT能显著提升LLM在长上下文任务上的性能,且优于其他TTT方法。
📝 摘要(中文)
静态的“训练后部署”模式从根本上限制了大型语言模型(LLMs)动态调整其权重以响应现实世界任务中固有的连续新信息流。测试时训练(TTT)提供了一种引人注目的替代方案,它在推理时更新模型参数的一个子集(快速权重),但其在当前LLM生态系统中的潜力受到关键障碍的阻碍,包括架构不兼容、计算效率低下以及与语言建模不一致的快速权重目标。在这项工作中,我们引入了In-Place测试时训练(In-Place TTT),该框架无缝地赋予LLM测试时训练能力。In-Place TTT将无处不在的MLP块的最终投影矩阵视为其可适应的快速权重,从而实现LLM的“即插即用”增强,而无需从头开始进行昂贵的重新训练。此外,我们用一个定制的、理论上合理的、明确与控制自回归语言建模的下一个token预测任务对齐的目标取代了TTT的通用重建目标。这种有原则的目标,结合高效的块状更新机制,产生了一种与上下文并行性高度兼容的可扩展算法。广泛的实验验证了我们框架的有效性:作为一种就地增强,它使一个40亿参数的模型能够在上下文高达128k的任务上实现卓越的性能,并且当从头开始预训练时,它始终优于有竞争力的TTT相关方法。消融研究结果进一步提供了对我们设计选择的更深入的见解。总的来说,我们的结果将In-Place TTT确立为LLM持续学习范式的一个有希望的步骤。
🔬 方法详解
问题定义:现有大型语言模型(LLMs)采用静态的“训练后部署”模式,无法在推理阶段动态适应新的数据和任务,导致在不断变化的环境中性能下降。传统的Test-Time Training (TTT)方法虽然尝试解决这个问题,但存在架构不兼容、计算效率低以及目标函数与语言建模任务不一致等问题。
核心思路:In-Place TTT的核心思想是在不改变LLM原有架构的基础上,将模型中的部分参数(MLP块的最终投影矩阵)作为可快速适应的“快速权重”,并通过一个专门设计的、与语言建模任务对齐的目标函数来更新这些权重。这种方法旨在实现LLM的“即插即用”式增强,使其能够在推理时动态学习和适应新信息。
技术框架:In-Place TTT框架主要包含以下几个关键组成部分:1) 选择MLP块的最终投影矩阵作为快速权重;2) 设计与Next-Token-Prediction任务对齐的损失函数,用于指导快速权重的更新;3) 采用高效的chunk-wise更新机制,以提高计算效率并支持上下文并行;4) 将更新后的模型应用于下游任务进行推理。
关键创新:In-Place TTT的关键创新在于:1) 提出了“In-Place”的概念,即在不改变模型架构的前提下实现TTT,避免了重新训练的成本;2) 设计了与Next-Token-Prediction任务对齐的损失函数,解决了传统TTT方法目标函数与语言建模任务不一致的问题;3) 采用了chunk-wise更新机制,提高了计算效率并支持上下文并行。
关键设计:In-Place TTT的关键设计包括:1) 选择MLP块的最终投影矩阵作为快速权重,这是一种轻量级的参数更新方式,易于实现且计算成本较低;2) 设计的损失函数基于Next-Token-Prediction任务,能够更好地指导快速权重的更新,提高模型的性能;3) chunk-wise更新机制将输入序列分成多个chunk,并行地更新快速权重,从而提高计算效率。
🖼️ 关键图片
📊 实验亮点
实验结果表明,In-Place TTT能够显著提升LLM在长上下文任务上的性能。例如,一个40亿参数的模型在上下文长度为128k的任务上取得了优异的性能。此外,从头开始预训练的模型在使用In-Place TTT后,性能始终优于其他TTT相关方法。消融实验也验证了该框架中各个设计选择的有效性。
🎯 应用场景
In-Place TTT具有广泛的应用前景,例如:在线学习、个性化推荐、智能对话系统、机器翻译等。它可以使LLM在实际应用中更好地适应用户需求和环境变化,提高模型的性能和鲁棒性。此外,该方法还可以应用于持续学习领域,使LLM能够不断地学习新的知识和技能。
📄 摘要(原文)
The static
train then deploy" paradigm fundamentally limits Large Language Models (LLMs) from dynamically adapting their weights in response to continuous streams of new information inherent in real-world tasks. Test-Time Training (TTT) offers a compelling alternative by updating a subset of model parameters (fast weights) at inference time, yet its potential in the current LLM ecosystem is hindered by critical barriers including architectural incompatibility, computational inefficiency and misaligned fast weight objectives for language modeling. In this work, we introduce In-Place Test-Time Training (In-Place TTT), a framework that seamlessly endows LLMs with Test-Time Training ability. In-Place TTT treats the final projection matrix of the ubiquitous MLP blocks as its adaptable fast weights, enabling adrop-in" enhancement for LLMs without costly retraining from scratch. Furthermore, we replace TTT's generic reconstruction objective with a tailored, theoretically-grounded objective explicitly aligned with the Next-Token-Prediction task governing autoregressive language modeling. This principled objective, combined with an efficient chunk-wise update mechanism, results in a highly scalable algorithm compatible with context parallelism. Extensive experiments validate our framework's effectiveness: as an in-place enhancement, it enables a 4B-parameter model to achieve superior performance on tasks with contexts up to 128k, and when pretrained from scratch, it consistently outperforms competitive TTT-related approaches. Ablation study results further provide deeper insights on our design choices. Collectively, our results establish In-Place TTT as a promising step towards a paradigm of continual learning in LLMs.