让大模型在推理时动态调整权重,无需重训即可适应新信息。
In-Place Test-Time Training
- 将MLP层的投影矩阵作为可调快参数,实现即插即用的测试时训练。
- 在128k上下文任务中,40亿参数模型性能超越现有方法。
- 针对语言建模设计目标函数,适合需要持续学习的场景。
静态的‘训练后部署’范式限制了大语言模型(LLMs)对真实世界任务中持续流入的新信息做出动态权重调整。测试时训练(TTT)通过在推理时更新部分模型参数(快参数)提供了一种替代方案,但其在当前LLM生态中的应用受限于架构不兼容、计算效率低及快参数目标与语言建模不匹配等问题。本文提出一种名为‘原位测试时训练’(In-Place TTT)的框架,使LLMs具备测试时训练能力。该框架将通用MLP块的最终投影矩阵视为可调快参数,实现无需从头训练的‘即插即用’增强。同时,我们以针对自回归语言建模中下一个词预测任务的理论基础目标替代传统重建目标,并结合高效的分块更新机制,使算法高度可扩展且兼容上下文并行。大量实验验证了该框架的有效性:作为原位增强,40亿参数模型在长达128k上下文的任务上表现更优;若从头预训练,亦持续优于竞争性TTT方法。消融实验进一步揭示设计选择的深层影响。综合来看,本工作为大模型持续学习范式迈出关键一步。
原文摘要 · Abstract (English)
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 a ``drop-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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。