PaperlayerView exact revision on arXiv

Source-linked AI summary

In-Place Test-Time Training

Guhao Feng, Shengjie Luo, Kai Hua, Ge Zhang, Di He, Wenhao Huang, Tianle Cai

arXiv:2604.06169v1cs.LGcs.AIcs.CLstat.ML

TL;DR

LLMs cannot update their weights during inference, limiting adaptation to streaming, evolving contexts. In-Place TTT repurposes MLP projections for efficient, language-model-aligned test-time updates, improving long-context performance and outperforming strong TTT baselines.

  • Problem

    Static LLMs cannot update weights during inference, limiting adaptation to streaming inputs and long-horizon, evolving tasks.

  • Method

    In-Place TTT updates existing MLP final projection matrices with chunk-wise, next-token-prediction-aligned objectives, requiring no architectural modification or retraining from scratch.

  • Results

    Experiments show In-Place TTT improves Qwen3-4B-Base on contexts up to 128k and outperforms strong TTT-related baselines when pretrained from scratch.

  • Takeaways & Limitations

    In-Place TTT is a promising step toward scalable, dynamic, continual learning for LLMs.

  • Takeaways & Limitations

    The approach is framed around compatibility with pretrained checkpoints and scalable parallel execution, addressing costly retraining and sequential per-token updates in existing TTT mechanisms.

Abstract

from arXiv · show

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.

1 Introduction

In-Place TTT equips LLMs with dynamic test-time adaptation by updating existing MLP projection matrices, while addressing TTT’s architectural, computational, and objective-alignment barriers. Its chunk-wise, attention-complementary design supports scalable language-model adaptation and enables strong long-context performance.

  • Motivation: TTT addresses static-model limitations by dynamically updating fast weights to compress and internalize information from streaming input contexts.Unlike fixed deployment, fast weights provide an online evolving state through self-supervised updates.
  • Efficient adaptation: In-Place TTT uses scalable chunk-wise updates and complements attention, avoiding the small chunks required by standalone TTT layers and improving throughput.This directly targets TTT’s sequential-update bottleneck and limited accelerator parallelism.
  • Language-model objective: The framework replaces generic reconstruction-based fast-weight updates with a bespoke mechanism tailored to causal Next-Token Prediction in autoregressive language models.The objective is intended to address the mismatch between canonical TTT updates and language modeling.
  • Core framework: In-Place TTT repurposes each MLP block’s final projection matrix as fast weights updated during inference, avoiding specialized layers and costly pretraining from scratch.The framework is designed as an in-place enhancement for existing LLM architectures.
  • Empirical validation: Qwen3-4B-Base achieves superior performance on tasks with contexts up to 128k through relatively cheap continual training.The experiments use language-modeling tasks at various compute scales as a proxy for long-horizon, evolving tasks.

2 Preliminary: Test-Time Training

Test-Time Training dynamically adapts models at inference by rapidly updating fast weights that serve as contextual neural memory. Applying TTT effectively to LLMs requires architectural compatibility, computational efficiency, and a learning objective tailored to next-token prediction.

  • TTT mechanism: TTT dynamically adapts models at inference by rapidly updating fast weights that continuously store and retrieve contextual information.Fast weights W form a small neural network f_W(·) and differ from standard weights, which remain frozen after training.
  • TTT mechanism: TTT sequentially updates fast weights using key-value pairs, then applies the updated network to queries enriched by preceding contextual information.The update minimizes a loss such as Mean Squared Error, while the apply operation computes o_i = f_Wi(q_i).
  • Desiderata for TTT within the LLM ecosystem: Effective TTT for LLMs requires architectural compatibility, computational efficiency, and a learning objective tailored to language modeling.These desiderata motivate the framework introduced in the paper.
  • Desiderata for TTT within the LLM ecosystem: Architectural compatibility requires warm-starting from a pretrained checkpoint rather than replacing attention with standalone recurrent layers that require pretraining from scratch.This drop-in property is especially important for billion-parameter LLMs.
  • Desiderata for TTT within the LLM ecosystem: Computational efficiency requires moving beyond inherently sequential per-token updates to support parallel processing and scalable language modeling.The canonical update bottlenecks GPUs and TPUs, making fine-grained updates impractical for high-throughput language modeling.
  • Desiderata for TTT within the LLM ecosystem: A tailored learning objective is needed because reconstruction-based key-value association does not guarantee alignment with next-token prediction.In prevailing TTT objectives, the target value v_i is typically derived from the input token x_i itself.

3 In-Place Test-Time Training

In-Place TTT equips LLMs with test-time adaptation by repurposing existing MLP blocks, updating only their final projection in place. It combines chunk-wise computation with an objective aligned to next-token prediction and theoretically shows predictive advantages over reconstruction targets.

  • In-Place Adaptation: In-Place TTT repurposes existing Transformer MLP blocks as fast weights, avoiding replacement of the attention mechanism and enabling a lightweight drop-in enhancement.Only existing components are reused, preserving the model’s architectural integrity while adding test-time training capability.
  • Chunk-Wise Updates: Chunk-wise updates process token blocks sequentially with apply and gradient-descent update operations, allowing large chunks to exploit GPU or TPU parallelism.The mechanism partitions activations into non-overlapping chunks and updates the fast-weight state using each chunk’s keys and values.
  • In-Place Adaptation: The framework freezes Wup and Wgate while updating only Wdown, the gated MLP’s final projection matrix, as adaptable fast weights.This in-place update transforms the MLP without introducing new architectural components.
  • LM-Aligned Objective: The LM-Aligned objective replaces same-token reconstruction with targets containing future-token information, explicitly matching the Next-Token Prediction goal.The target is defined as ˆV = Conv1D(X0)Wtarget, with future-token information controlled through Wtarget and convolution settings.
  • Theoretical Analysis: Theorem 1 shows that LM-Aligned updates increase the correct next-token logit in expectation while leaving other logits approximately unchanged, unlike reconstruction updates.The analysis is grounded in the canonical induction-head setting and concludes that reconstruction provides no comparable predictive benefit.

4 Experiments

Experiments show that In-Place TTT improves long-context performance as a drop-in enhancement for pretrained LLMs and remains effective when trained from scratch. Ablations further support scaling fast weights, chunk-size choices, the LM-aligned objective, and efficient implementation.

  • Drop-in enhancement: In-Place TTT consistently improves RULER performance over pretrained Qwen3-4B-Base as context length increases, with widening gains at 64k and 128k.Both models follow the same continual training curriculum; the stages use approximately 20B tokens at 32k context and 15B tokens at 128k context.
  • Drop-in enhancement: +2.1 at 64k for LLaMA-3.1-8B and +2.7 at 64k for Qwen3-14B-Base demonstrate that the drop-in enhancement generalizes across additional models.In-Place TTT improves RULER scores across all context lengths on both models, with particularly pronounced gains at longer contexts.
  • Pretraining from scratch: In-Place TTT achieves lower sliding-window perplexity than all competitive baselines for 500M and 1.5B models, improving steadily through the full 32k context.The comparisons include SWA, GLA, DeltaNet, and LaCT, with In-Place TTT and LaCT using an SWA backbone.
  • Pretraining from scratch: 19.99 is the RULER-16k score for Full Attention with In-Place TTT, up from 6.58; 26.80 is the RULER-8k score for SWA with In-Place TTT, up from 9.91.Across 4B models, In-Place TTT also consistently improves performance on most common-sense reasoning tasks.
  • Ablations: Performance improves as fast-weight state size increases, while chunk sizes C = 512 and C = 1024 outperform other choices and C = 1024 is more efficient.State size is controlled by the number of TTT-enabled layers, and chunk size trades off update granularity, parallelism, efficiency, and performance.
  • Ablations: Both Conv1D and Wtarget are necessary for the LM-aligned objective, with Conv1D essential for long context and Wtarget crucial for short context.Conv1D produces targets containing future-token information, while Wtarget performs a projection transformation; efficiency experiments verify the practical implementation’s overhead.

5 Related Work

Related work spans Test-Time Training, efficient long-context architectures, and memory augmentation. These directions address dynamic adaptation, reduced attention complexity, and enhanced reasoning through explicit memory modules.

  • Test-Time Training: Test-Time Training adapts models during inference by updating a small subset of parameters called fast weights, with applications across vision, language, video, and audio.Research has largely focused on two avenues, but the supplied passage truncates before specifying them.
  • Efficient Long-Context Architectures: Efficient long-context architectures mitigate the quadratic complexity of standard attention through sparse attention or linear-time approximations and recurrent or gated formulations.Sparse methods restrict token interactions using fixed patterns such as sliding or strided windows.
  • Memory Design and Augmentation: Memory augmentation enhances reasoning and contextual understanding through explicit modules that store persistent task-agnostic knowledge or transient context-dependent information.The passage distinguishes external memory banks for persistent knowledge from mechanisms capturing information in the immediate context.

6 Conclusion · Appendix

The paper introduces In-Place Test-Time Training as a practical, scalable framework for dynamic adaptation in LLMs. Its in-place mechanism, chunk-wise updates, and language-modeling-aligned objective support both drop-in enhancement and continual learning.

  • 6 Conclusion: In-Place Test-Time Training resolves critical barriers that previously limited Test-Time Training for LLMs.
  • 6 Conclusion: The framework repurposes existing MLP blocks to provide an in-place adaptation mechanism.
  • 6 Conclusion: A chunk-wise update rule improves the framework’s efficiency and scalability.
  • 6 Conclusion: A theoretically grounded objective aligns the test-time updates with language modeling.
  • 6 Conclusion: Extensive experiments show that the approach functions as a powerful drop-in enhancement for pretrained LLMs and outperforms strong baselines.
  • 6 Conclusion: The framework provides a scalable solution for on-the-fly adaptation in LLMs.
  • Appendix: The work represents a promising step toward more dynamic continual learning for LLMs.

A Proof of theorem 1

Theorem 1 shows that the NTP-aligned target changes the correct-token logit while leaving other logits nearly unchanged under the stated assumptions. The proof contrasts this with the reconstruction target, whose expected effect on the correct answer is negligible.

  • NTP-aligned target: Theorem 1 bounds the expected change in non-correct-token logits under the NTP-aligned target by |E [∆ℓn[w]]| ≤λlr · ϵ · calign, ∀w̸ = v∗.The bound applies for learning rate λlr > 0 under the specified setting and assumptions.
  • Proof mechanism: The proof reduces the expected update to the unique key position t∗, where Assumption 2 gives E[z⊤t∗zn] = calign and zero expected information gain elsewhere.This isolates the key-value pair (k∗, v∗) at position t∗.
  • NTP-aligned target: For the correct token w = v∗, Assumption 1’s non-trivial embedding magnitude yields the lower bound in equation (4).The assumption is stated as ∥ev∗∥2 ≥c2 norm.
  • Reconstruction target: For the reconstruction target, the expected effect on the correct answer v∗ is negligible, including when the induction task has distinct key k∗ and value v∗.The proof invokes near-orthogonality of distinct token embeddings to establish the corresponding bound.
  • Conclusion: Together, the two cases establish Theorem 1 and provide a theoretical basis for the superiority of the NTP-aligned objective in in-context learning.The proof concludes after analyzing both target choices.

B Context Parallel Algorithm for In-Place TTT

The context-parallel In-Place TTT algorithm partitions sequences into chunks, computes update deltas in parallel, aggregates them associatively, and applies prefix updates when producing outputs. It resets the fast weights at document boundaries.

  • Context Parallel Algorithm: Sequence chunks are partitioned for context parallelism, and each chunk independently computes its attention representations and update quantities.The attention block requires no architectural changes.
  • Context Parallel Algorithm: The algorithm computes NTP-aligned targets with causal padding and uses them to form each chunk’s fast-weight update gradient.The update gradient is computed from the target and intermediate chunk quantities.
  • Context Parallel Algorithm: Chunk update deltas are aggregated associatively before outputs are computed in parallel.This separates parallel delta computation from the subsequent update-and-output stage.
  • Context Parallel Algorithm: Each chunk uses an effective weight incorporating updates from preceding chunks, preserving the sequence order during parallel execution.The effective weight for chunk i uses updates from chunks < i.
  • Context Parallel Algorithm: At document boundaries, the algorithm resets the fast weights to their initial values.This prevents fast-weight state from carrying across documents.

C Experiment Details

The appendix specifies the experimental settings, datasets, model configurations, and training hyperparameters used in Section 4, covering three primary experiment sets to support reproducibility.

  • The appendix documents the settings, datasets, model configurations, and training hyperparameters underlying the Section 4 results.These details are presented to make the findings reproducible.
  • The experiments comprise continual pre-training of Qwen3-4B-Base, from-scratch pre-training at 500M, 1.5B, and 4B scales, and targeted ablations.These are identified as the three primary sets of experiments.

C.1 Details of Datasets

The study uses self-collected datasets for large-scale pretraining, continual pretraining, and ablation studies. These datasets combine broad linguistic, knowledge, reasoning, programming, mathematical, multilingual, and long-context data.

  • The experiments use self-collected datasets for large-scale pretraining, continual pretraining, and ablation studies.
  • From-scratch pretraining combines English and Chinese text with knowledge- and reasoning-dense data, code, mathematics, and multilingual text.The mixture is designed to balance linguistic diversity, knowledge and reasoning content, programming material, and mathematical reasoning.
  • Continual pretraining combines similarly distributed short documents with natural and synthetic long-document data for challenging long-context scenarios.Long-document data includes books, repository-level code, retrieval-augmented constructions, and long-context question-answering styles.

C.2 Details of Training and Evaluation

Training and evaluation use Nvidia H800 GPUs, with hyperparameters specified in Tables 4–7 and benchmarks evaluated through lm-evaluation-harness and opencompass. Long-context Qwen3-4B evaluation additionally applies clipped fast-weight updates and measures throughput and memory from 8k to 128k tokens.

  • Training Details: All models are trained on Nvidia H800 GPUs, with hyperparameters provided in Tables 4 through 6.Table 7 additionally specifies continual pre-training hyperparameters for LLaMA-3.1-8B and Qwen3-14B-Base.
  • Evaluation Details: Common-sense reasoning uses lm-evaluation-harness, while long-context benchmarks use opencompass, with evaluations conducted on Nvidia H800 GPUs.
  • Inference Stability: For continual-pretrained Qwen3-4B, updates exceeding threshold τ are rescaled by Frobenius norm before application to stabilize long-context inference.The threshold for all reported Qwen3-4B evaluations is τ = 1e-5.
  • Efficiency Evaluation: 8k to 128k token sequences are used to measure In-Place TTT prefill throughput and peak memory with batch size 1 on Nvidia H800 GPUs.The experiments use continual-pretrained Qwen3-4B-Base checkpoints; sliding-window settings manually change attention to a 1024-token window.

C.3 Details of Model Configuration

The experiments use decoder-only Transformer models with standard components, while varying attention mechanisms and integrating In-Place TTT into selected MLP layers. For continual pre-training, the modules are initialized to preserve the original model behavior before gradually learning NTP-aligned targets.

  • Architectural configurations: All models trained from scratch use decoder-only Transformers with SwiGLU activations and Rotary Position Embeddings (RoPE).Their key architectural parameters are summarized in Table 8.
  • Attention mechanisms: Attention varies by experiment: 500M and 1.5B models use sliding-window attention, while 4B experiments compare full and sliding-window variants.The 4B and 1.7B backbones match Qwen3-4B-Base and Qwen3-1.7B-Base, respectively.
  • Continual pre-training models: Continual pre-training starts from Qwen3-4B-Base, LLaMA-3.1-8B, and Qwen3-14B-Base, inheriting their architectures without modification.These models are publicly available pre-trained models.
  • TTT integration: In-Place TTT is integrated into MLP blocks every sixth layer, with ablation studies varying this frequency.Training hyperparameters are provided separately for Qwen3-4B-Base, LLaMA-3.1-8B, and Qwen3-14B-Base.
  • Initialization of In-Place TTT Modules: The Conv1D operator and projection matrix Wtarget are initialized so the TTT update ∆Wdown is negligible, preserving pre-trained behavior at initialization.The depth [i]Z[i] ≈0, so effective Wdown remains identical to its pre-trained value; training then enables gradual emergence of NTP-aligned targets.
Loading 2604.06169v1…