原位测试时训练 In-Place Test-Time Training
通过复用MLP块实现测试时动态参数更新，增强LLM长上下文适应性
前置知识
Test-Time Training (TTT)
TTT是一种在推理时动态更新模型参数的范式。它引入可快速更新的fast weights，通过更新和应用两个核心操作，使模型能够持续压缩和内化上下文信息。
本文直接基于TTT范式，理解它的工作机制是理解In-Place TTT如何实现动态适应的基础。
Fast Weights
Fast weights是TTT中可动态更新的参数集合，与固定的slow weights相对。它们通过梯度下降步骤快速更新，充当动态记忆，持续存储和检索序列的上下文信息。
本文的核心创新是将现有MLP块的投影矩阵复用为fast weights，理解fast weights的概念是理解这个设计选择的关键。
Next-Token Prediction (NTP)
NTP是自回归语言模型的核心训练目标，要求根据上下文预测序列中的下一个token。它强制模型学习token间的因果依赖关系，是上下文学习和文本生成的基础。
本文提出与NTP对齐的新目标函数，理解NTP是理解为何这个目标比传统reconstruction目标更适合语言模型的关键。
Context Parallelism
Context parallelism是一种高效处理长序列的并行策略，将序列分割成多个chunk同时处理。通过parallel scan等算法维护因果依赖，显著提升长序列处理的吞吐量。
本文的实现完全兼容context parallelism，这是其能够高效处理长达128k上下文的关键技术之一。
研究动机
当前大语言模型采用的静态'训练后部署'范式存在根本性限制。一旦部署，模型权重无法更新，这阻止了模型动态适应流式输入token提供的特定上下文，导致在测试时处理长期演进任务的能力受限，无法像人类一样从无界经验流中持续学习。现有的上下文学习通过在context中维护所有过去的输入token来缓解这个问题，但其效果受限于模型的context window，而事实上的attention机制的二次方复杂度构成了瓶颈。TTT作为一种新的范式出现，但它面临三个关键障碍：架构不兼容，现有TTT方法依赖标准Transformer块之外的专门层，通常需要从头预训练；计算效率低，规范的TTT机制本质上是顺序的；以及用于语言建模的fast weight目标不对齐，普遍使用的通用重构目标没有显式为自回归语言建模的因果Next-Token Prediction任务量身定制。
本文的目标是本文旨在通过直接解决上述三个障碍，设计一个无缝地为LLM赋予Test-Time Training能力的框架。具体目标包括：实现架构兼容性，使方法能够从预训练checkpoint热启动，不需要昂贵的从头预训练；确保计算效率，在现代并行加速器上高效运行；设计专门为语言模型定制的目标函数，显式与Next-Token Prediction目标对齐。最终目标是将TTT转化为一个实用的'drop-in'增强模块，能够无缝集成到现有的预训练LLM中。
与已有工作不同的是，本文的独特切入点在于不是引入新的专门层或替代attention机制，而是直接复用Transformer中无处不在的MLP块。核心洞察是：MLP块在Transformer中可以被视为一种键值记忆，作为在预训练期间获取的大量通用知识的'slow weights'。将同一组件也用作自适应的'fast weights'，在推理时动态内化瞬时的上下文信息，是一个自然的扩展。这种in-place设计消除了引入新层的需求，保持了模型的架构完整性，将TTT从破坏性重构转变为轻量级的增强模块。这种设计使得方法能够完全兼容现有的预训练权重，使LLM能够以零架构修改的方式在测试时动态适应，这是前所未有的。
核心方法
In-Place TTT的整体思路是复用现有MLP块的最终投影矩阵作为fast weights，在推理时进行in-place更新。方法采用chunk-wise更新机制确保计算效率，并提出与Next-Token Prediction对齐的目标函数以优化语言建模性能。直觉上，MLP块已经学习了大量的通用知识，通过将其投影矩阵变为可更新的fast weights，模型能够在不破坏原有知识的基础上，动态地适应当前上下文。技术路线上，方法分为三个关键组件：MLP块复用机制，将Wdown作为fast weights，保持Wup和Wgate冻结；chunk-wise更新规则，将输入序列分成非重叠块，对每个块先应用当前的fast weights，然后使用梯度下降步骤更新；LM-aligned目标函数，设计包含未来token信息的target，使其显式与NTP目标对齐。
核心创新点是将Transformer中无处不在的MLP块的最终投影矩阵Wdown复用为可更新的fast weights，而不是引入新的专门层。这与已有方法的本质区别在于：已有TTT方法通常作为独立的循环层设计来替代attention，而本文的方法与attention机制互补操作，保留了attention层的学习属性。具体而言，对于gated MLP架构O = ϕ(HWgate⊤) ⊙ (HWup⊤)Wdown⊤，本文将Wup和Wgate视为冻结的slow weights，而将Wdown处理为可适应的fast weights，在推理时in-place更新。这种'drop-in'设计不需要修改模型架构，保持预训练权重的完整性，实现了无需昂贵从头预训练的即时适应。另一个关键创新是提出LM-aligned目标：target定义为V̂ = Conv1D(X0)Wtarget，其中X0是token嵌入，Conv1D是1D卷积算子，Wtarget是可训练投影矩阵。通过参数化Wtarget和Conv1D的核权重，可以控制target中未来token信息的数量，例如Next-Token target可以通过将Wtarget参数化为恒等变换并将Conv1D的核权重设置为下一个token为1、其他token为0来实现。
方法步骤详情
方法步骤的完整描述如下。给定输入序列X，首先通过标准的attention层获得中间表示H。然后进入MLP块，MLP块分为冻结部分和可更新部分：冻结部分计算Ui = HWup⊤和Gi = HWgate⊤，通过激活函数和门控得到Zi = ϕ(Gi) ⊙ Ui。在chunk-wise更新中，将Z, V, O分割成k个非重叠块，大小为C，记为[i] = □iC+1:(i+1)C。对于每个块i ∈ [k]，执行两个顺序操作：应用操作：使用当前fast weights状态Wdown(i)处理块Z[i]，即O[i] = Z[i] (Wdown(i))⊤；更新操作：使用Z[i]作为keys，V[i]作为values更新fast weights，通过一个梯度下降步骤：Wdown(i+1) = Wdown(i) - η∇W L(Z[i] (Wdown(i))⊤, V[i])。在使用相似度度量的损失函数L(·, ·) = -⟨·, ·⟩F下，fast weights的梯度可以直接推导为Wdown(i) = Wdown(i-1) + ηV̂⊤[i] Z[i]。在文档边界处，fast weights重置为预训练状态。对于context parallelism实现，分为三个阶段：对所有块并行计算中间激活和fast weight更新增量；对增量进行单个前缀和计算每个块的聚合更新；并行计算每个块的有效fast weights和对应输出。
技术新颖性
技术新颖性体现在三个方面。首先是架构层面，首次提出将现有MLP块直接复用为fast weights的in-place设计，避免了引入新层或修改架构的需要，使方法能够无缝集成到预训练模型中。其次是算法层面，由于只适配MLP块并保持attention层完整，方法从per-token约束中解放出来，能够使用高效的chunk-wise更新策略，特别是能够使用大块大小512-1024，避免了独立TTT层需要小块来维持性能的限制。第三是目标函数层面，超越了之前工作的通用重构目标，引入了显式与Next-Token Prediction目标对齐的新目标，理论分析表明这个目标鼓励fast weights存储对自回归语言建模预测有用的信息，而不仅仅是记忆当前token的表示。这种组合使得In-Place TTT成为第一个实用、高效且与LLM生态系统完全兼容的TTT框架。
实验结果
核心发现来自三个实验设置。在作为预训练LLM的'drop-in'增强实验中，使用Qwen3-4B-Base模型进行持续训练，分为两个阶段：第一阶段约20B tokens，32k上下文长度；第二阶段约15B tokens，128k上下文长度。在RULER基准上的评估显示，In-Place TTT显著提升了模型的长上下文能力。具体而言，在短上下文4k-16k下，两个模型都很有竞争力，但随着序列长度增加，In-Place TTT的优势逐渐扩大。在64k和128k上下文长度下获得了实质性增益。关键是在外推到256k上下文时保持了这个优势，证明了优越的泛化能力。将In-Place TTT扩展到LLaMA-3.1-8B和Qwen3-14B-Base，按照相同的持续训练协议，两个模型在所有上下文长度上都显示出一致的RULER分数提升，增益在更长上下文时特别明显，例如LLaMA-3.1-8B在64k时提升+2.1，Qwen3-14B在64k时提升+2.7。当与YaRN结合用于位置外推时，提升得以保持，确认了方法与RoPE扩展技术的正交性。在从头预训练的比较分析中，在500M和1.5B参数规模上与多种高效注意力方法基准比较，包括滑动窗口attentionSWA、门控线性注意力GLA、DeltaNet和大块TTTLaCT。在Pile数据集上的滑动窗口困惑度评估显示，In-Place TTT在所有上下文长度2k-32k上都实现了比所有竞争基线更低的验证困惑度，性能稳步提升到完整的32k上下文。在4B参数规模上，比较了Full Attention和SWA的Transformer及其对应的In-Place TTT增强版本，训练了120B tokens，8k上下文长度。在常识推理基准上，4B参数的Full Attention和SWA Transformer在大多数任务上都得到一致改进。在长上下文评估上，In-Place TTT产生了优越的性能，例如Full Attention的Transformer在RULER-16k上从6.58提升到19.99，SWA模型在RULER-8k上从9.91提升到26.80。消融研究在1.7B参数模型上对关键设计选择的影响进行了调查。状态大小的影响显示出明确趋势，随着fast weights大小通过TTT启用层的数量控制的扩展，In-Place TTT的性能一致改善，证实更大的fast weights允许模型更有效地适应上下文信息。块大小的影响显示，C=512和C=1024都比其他选择竞争性地实现更好的性能，其中C=1024具有更好的效率。LM-aligned目标的深入分析显示，Conv1D和投影组件都是性能保证所必需的，其中Conv1D在长上下文中起关键作用，Wtarget在短上下文中至关重要。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| RULER Long-Context Benchmark (Qwen3-4B-Base, 64k context) | Average Accuracy (%) | 78.7 | 74.3 | +4.4 absolute percentage points |
| RULER Long-Context Benchmark (Qwen3-4B-Base, 128k context) | Average Accuracy (%) | 77.0 | 74.8 | +2.2 absolute percentage points |
| RULER Long-Context Benchmark (Qwen3-4B-Base, 256k extrapolation) | Average Accuracy (%) | 43.9 | 41.7 | +2.2 absolute percentage points |
| Sliding Window Perplexity on Pile (500M model, 32k context) | Perplexity | ~9.0 | LaCT ~9.5, SWA ~11.0 | ~0.5-2.0 lower perplexity compared to baselines |
| RULER-16k (4B model, Full Attention backbone) | Average Accuracy (%) | 19.99 | 6.58 | +13.41 absolute percentage points (203% relative improvement) |
| RULER-8k (4B model, SWA backbone) | Average Accuracy (%) | 26.80 | 9.91 | +16.89 absolute percentage points (170% relative improvement) |
| RULER-64k (LLaMA-3.1-8B) | Average Accuracy (%) | 83.7 | 81.6 | +2.1 absolute percentage points |
| RULER-64k (Qwen3-14B-Base) | Average Accuracy (%) | 70.6 | 67.9 | +2.7 absolute percentage points |
局限与改进
作者承认的局限性包括：核心原则与损失函数和优化器的具体选择是正交的，这些在更广泛的TTT文献中已经得到了广泛研究，本文没有探索这些方向。此外，方法虽然与YaRN等RoPE扩展技术正交，但对于更复杂的位置编码策略的兼容性需要进一步验证。从个人观察来看，方法的性能提升与上下文长度正相关，这意味着对于短上下文任务4k以下，增益可能不够显著，限制了其在这些场景下的实用性。虽然chunk-wise更新机制提高了效率，但与标准Transformer相比仍然引入了额外的计算开销，在资源受限的部署环境中可能成为瓶颈。fast weights的重置策略在文档边界处虽然避免了上下文泄漏，但也限制了跨文档的长期记忆能力，这在需要维护跨文档信息的应用中可能成为限制。此外，虽然方法被证明适用于4B-14B规模的模型，但对于更大规模的模型如70B+的扩展性和训练稳定性尚未得到验证，这在大规模商业应用中是一个重要考虑因素。
独立分析的弱点
独立分析的弱点包括：首先，方法对超参数块大小、学习率、状态大小相对敏感。消融研究显示块大小需要在512-1024范围内才能获得最优性能，这意味着在不同任务和模型规模上需要仔细调优，增加了部署的复杂度。其次，虽然方法在长上下文任务上表现出色，但在短上下文任务上的收益有限，这可能限制了其在通用对话系统等短上下文主导场景中的应用价值。第三，fast weights的更新规则虽然简单高效，但这种线性更新可能无法捕获更复杂的上下文依赖关系，特别是在需要长期推理或多跳推理的任务中。改进方向包括：设计自适应块大小机制，根据上下文长度和任务类型动态调整块大小，平衡效率和性能；探索非线性更新规则或多层fast weights架构，以增强对复杂上下文依赖的建模能力；研究跨文档记忆机制，例如引入持久fast weights状态，在文档之间选择性保留关键信息；开发针对短上下文任务的轻量级变体，减少在短上下文场景下的开销。
未来方向
作者提出的未来工作方向包括探索与损失函数和优化器的具体选择相关的更广泛设计空间，这些在更广泛的TTT文献中已经得到了广泛研究。基于成果可延伸的方向包括：将In-Place TTT与高效长上下文架构如线性attention、状态空间模型自然集成，因为这些架构也有MLP块，可以无缝适配；扩展到多模态模型，探索将MLP块复用为fast weights的设计在视觉-语言模型中的适用性；研究fast weights的更复杂更新机制，如引入动量、自适应学习率或二阶优化方法，以提高更新效率和质量；探索将fast weights应用于其他下游任务，如代码生成、数学推理等，验证方法的通用性；研究fast weights的解释性和可控性，使模型能够提供关于其适应过程的更透明解释，这对于可信赖AI系统至关重要；开发增量学习和持续学习的框架，使模型能够在部署后持续从新数据中学习，同时避免灾难性遗忘。
复现评估
复现评估显示，论文提供了详细的技术实现描述，包括context parallelism的伪代码算法和训练超参数表。作者提到代码将在https://github.com/ByteDance-Seed/In-Place-TTT开源，这将大大便利复现。数据方面，作者使用了自定义收集的数据集，包括英语和中文文本、高知识或推理密度数据、代码、数学数据和多语言文本。对于持续预训练，数据集设计用于增强长上下文建模：短文档部分遵循与预训练数据类似的分布，长文档部分结合自然数据书籍和仓库级代码与合成数据包括检索增强和长上下文QA风格构建。实验在Nvidia H800 GPU上进行，大规模预训练需要显著的计算资源120B tokens训练4B模型。虽然论文提供了详细的超参数设置，但复现完整实验需要访问类似规模的计算资源，这对于研究机构可能构成挑战。然而，方法的'drop-in'特性使得在较小规模上验证核心概念成为可能，降低了初步验证的门槛。
论文图表
