PaperlayerView exact revision on arXiv

Source-linked AI summary

Select-TTT: Query-Conditioned Selective Fast Weight Updates for In-Place Test-Time Training

Yue Wu, Yang Zhang, Weinan Gan, Chongming Gao, Xinsen Guo, Xiaoyuan Li, Fuli Feng

arXiv:2610.36794v1cs.LGcs.AIcs.CL

TL;DR

Long-context QA requires models to use sparse relevant evidence, but uniform In-Place TTT writes every chunk equally. Select-TTT uses query-conditioned scores from the frozen output head to rescale chunk writes, and consistently improves results across two benchmarks and two backbones.

  • Problem

    In realistic long-context QA, In-Place TTT’s uniform write schedule gives irrelevant chunks the same write strength as answer-bearing chunks, diluting query-relevant information.

    Source: §1 · pp. 1–2

  • Method

    Select-TTT estimates chunk relevance from query-token probability mass under the frozen output head and uses temperature-scaled softmax scores to rescale fast-weight writes.

    Source: §1 · p. 2

  • Results

    Select-TTT consistently outperforms uniform In-Place TTT and corresponding baselines on LongMemEval-S and LoCoMo with Llama-3.1-8B-Instruct and Qwen3-4B.

    Source: §1 · p. 2, §7 · p. 9

  • Takeaways & Limitations

    Query-conditioned fast-weight writes are a promising direction for practical test-time training on long-context question answering.

    Source: §7 · p. 9

    Evaluation covers two backbones up to 8B parameters and two English benchmarks; scaling to larger models, multilingual settings, and other modalities remains future work.

    Source: §8 · p. 9

Abstract

from arXiv · show

Test-time training (TTT) enables a deployed language model to internalize incoming context on-the-fly into a small set of fast-weight parameters, rather than merely treating it as transient context for attention. In-Place TTT recently made this idea practical for large language models by repurposing the down-projection of standard MLP blocks as fast weights. However, we find that its benefit largely disappears on realistic web long-context question answering. We attribute this to In-Place TTT's uniform write schedule: every chunk is written with equal strength, so irrelevant content dilutes query-relevant information in fast-weight memory. Building on this insight, we propose Select-TTT, a model-agnostic, query-conditioned rule for In-Place TTT that scales each chunk's write strength by its relevance to the user query. This relevance signal requires no supervision and no new parameters: it is read directly from the model's frozen output head, as the probability mass that a chunk's hidden states assign to the query tokens. A temperature-scaled softmax turns these scores into learning-rate multipliers, leaving the chunk-wise, context-parallel computation of In-Place TTT intact. Select-TTT consistently improves long-context question answering on two benchmarks and two backbones. These findings highlight the value of query-conditioned fast weight writes for practical test-time training on long-context question answering, which is also potentially promising for long-horizon agent tasks involving multiple user queries over evolving contexts.

1 Introduction

Long-context QA leaves models with sparse answer-bearing evidence while their deployed weights remain frozen. Select-TTT addresses the dilution caused by uniform chunk writes by conditioning write strength on query relevance, improving results across two benchmarks and backbones.

1 Introduction

  • In long-context QA, answer-bearing evidence occupies only a small part of the context, while deployed models’ weights remain frozen and cannot internalize incoming information.Models can reread context through attention, but may struggle to exploit it even when it fits within the context window.Source: §1 · p. 1
  • In-Place TTT makes test-time training practical by repurposing a gated MLP’s final projection as fast weights and organizing updates chunk-wise for context-parallel computation.It requires no architectural modification or training from scratch.Source: §1 · p. 1
  • Uniform In-Place TTT writes every chunk at the same strength, allowing irrelevant writes to dilute query-relevant information and yielding little or no improvement, sometimes below zero-shot performance.Oracle annotations show that amplifying answer-chunk writes improves accuracy, whereas amplifying irrelevant-chunk writes can push it below baseline.Source: §1 · pp. 1–2
  • Select-TTT scales each chunk’s fast-weight write by query relevance, estimated from query-token probability mass under the frozen output head and converted to learning-rate multipliers.Its temperature-scaled softmax produces positive, budget-preserving multipliers with an upper cap; without a query, the rule reduces to In-Place TTT.Source: §1 · p. 2
  • Across LongMemEval-S and LoCoMo with Llama-3.1-8B-Instruct and Qwen3-4B, Select-TTT consistently outperforms uniform In-Place TTT and corresponding baselines.The evaluation covers inference-time use and use during continual training on LongCite-45k.Source: §1 · p. 2

2 Preliminary

Test-time training uses fast weights as a dynamic memory, updating them while processing input and using the updated weights to incorporate context. In-Place TTT adapts this mechanism to standard gated MLP blocks with a next-token-prediction-aligned objective and chunk-wise updates.

2.1 Test-Time Training

  • Test-Time Training designates a small set of fast weights as dynamic memory that is updated during inference to store and retrieve contextual information.The weights parameterize a small network that processes the input sequence.Source: §2.1 · p. 2
  • TTT alternates a self-supervised update that encodes key-value information in fast weights with an apply operation that uses the updated network to process the query.The update uses a learning rate, and the output is enriched with information stored in the updated weights.Source: §2.1 · p. 2

2.2 In-Place TTT

  • In-Place TTT addresses the specialized-layer and sequential-update barriers by reusing the final projection of standard gated MLP blocks as fast weights, without architectural modification.Its chunk-wise, dense computation supports large chunks and context parallelism.Source: §2.1 · p. 2, §2.2 · pp. 2–3
  • In-Place TTT freezes the gated MLP’s gate and up-projections while updating its down-projection using chunk-wise updates aligned with next-token prediction.The value target is constructed from token embeddings with future-token information; a causal convolution and trainable projection can produce next-token targets.Source: §2.2 · p. 2

3 Motivation Experiment

In-Place TTT provides marginal gains on realistic long-context QA, motivating a test of whether answer-bearing and irrelevant chunks should receive different write strengths. The experiment finds opposing accuracy trends as those strengths change.

3 Motivation Experiment

  • On LongMemEval-S with Llama-3.1-8B-Instruct, In-Place TTT raised average accuracy only from 33.4% zero-shot to 33.6%.The authors characterize its benefit on realistic long-context user QA as marginal.Source: §3 · p. 3
  • The experiment used LongMemEval-S answer-bearing chunks to vary answer-chunk and irrelevant-chunk learning rates independently on Llama-3.1-8B-Instruct.The setup tests whether write strength should depend on chunk relevance to the query.Source: §3 · p. 3
  • 37.8% accuracy was reached by strengthening answer-chunk writes, up from 33.6% with uniform In-Place TTT; increasing irrelevant-chunk write strength instead lowered accuracy from 34.8% to 29.6%.The authors interpret these trends as evidence that irrelevant writes hurt performance while stronger answer-chunk writes help.Source: §3 · p. 3

4 Select-TTT: Query-Conditioned Selective In-Place TTT

Select-TTT uses query-conditioned chunk relevance to scale In-Place TTT write strengths, retaining its chunk-wise update path. It estimates relevance from query-token probabilities under the frozen output head and converts chunk scores into per-chunk learning-rate multipliers.

4.1 The Selective Learning-Rate Rule

  • A per-chunk multiplier λ_i rescales In-Place TTT’s baseline learning rate according to query relevance, weakening weakly related writes and strengthening highly relevant ones.The baseline learning rate sets overall write strength; the multiplier adjusts it for each chunk.Source: §4.1 · p. 3

4.2 The Query-conditioned chunk selection

  • The frozen output head maps chunk hidden states to vocabulary probabilities, using probability mass on query tokens as a model-internal relevance signal.This addresses the lack of direct query relevance in the chunks’ intermediate hidden states.Source: §4.2 · p. 3
  • Top-k mean pooling averages positions with the lowest query-token cross-entropy, then a temperature-scaled softmax over negative chunk scores assigns learning-rate multipliers.The pooling ratio controls how many positions are retained, while temperature controls how sharply chunk relevance differentiates the multipliers.Source: §4.2 · p. 4
  • Every chunk receives a positive multiplier, the average learning rate remains aligned with In-Place TTT’s global rate, and clipping limits any single chunk’s amplification.The smooth allocation avoids entirely ignoring potentially answer-bearing chunks, while the upper clip guards against overly aggressive writes.Source: §4.2 · p. 4

4.3 Efficient Chunk-Wise Scoring and Training

  • Select-TTT rescales each chunk update without changing In-Place TTT’s main chunk-wise pass, preserving its matmul-rich context-parallel computation and adding no parameters or optimizer states.The scoring pass uses hidden states already available at TTT layers; the output head is frozen and shared across layers.Source: §4.3 · p. 5
  • Blockwise log-sum-exp and direct gathering of question-token logits avoid materializing a full [T,C,V] logit tensor for relevance scoring.The gathered logits and blockwise normalizer provide the per-token cross-entropy terms.Source: §4.3 · p. 5

5 Experiments

The experiments test whether query-conditioned writes improve long-context QA at inference and during continual training, then examine design choices, relevance ranking, and computational overhead.

5.1 Experimental Setup

  • The study tests LongMemEval-S and LoCoMo with Llama-3.1-8B-Instruct and Qwen3-4B, using inference-only and continual-training protocols, with TTT in six layers per backbone.Source: §5.1 · pp. 5–6

5.2 Inference-Time Select-TTT as a Drop-in Enhancement

  • Select-TTT leads both benchmarks and backbones: LongMemEval-S average gains over IP-TTT are +3.8 (Llama) and +3.4 (Qwen), while LoCoMo overall scores reach 40.42 and 26.19.Source: Table 1 · p. 6, Table 2 · p. 6, §5.2 · p. 6
  • LongMemEval-S gains are larger on Knowledge Update, Multi-session, and Temporal Reasoning than on single-session subtasks.Source: §5.2 · p. 6

5.3 Continual Pre-Training with Select-TTT

  • At 5k steps, Select-TTT finishes at 37.4%, ahead of IP-TTT (35.6%) and CPT (35.0%); both TTT methods beat CPT from step 2k onward.Source: Figure 3 · p. 7, §5.3 · p. 7

5.4 Ablation Studies On the Impact of Select-TTT Design Choices

  • Removing both the softmax and upper clip causes a performance drop, supporting both safeguards in the multiplier design.Source: Table 3 · p. 7, §5.4 · p. 8

5.5 Analysis: Which Relevance Signal Picks the Relevant Chunks?

  • Query-CE ranks answer-bearing chunks most effectively across tested chunk sizes and k values, outperforming cosine similarity and a pretrained cross-encoder on Hit@k and Rec@k.Source: Table 4 · p. 7, §5.5 · p. 8

5.4 Ablation Studies On the Impact of Select-TTT Design Choices

  • Temperature 𝜏=0.1 performs best, top-𝑘 mean beats Max and Mean pooling, and the non-monotonic pooling-ratio results motivate 𝜌=0.10 as a compromise.Source: §5.4 · p. 7, Figure 4 · p. 8

5.6 Overhead analysis

  • Select-TTT’s lightweight scoring mechanism adds negligible overhead to IP-TTT in prefill throughput and peak GPU memory measurements on Qwen3-4B.Source: Figure 5 · p. 8, §5.6 · p. 8

6 Related Work

Prior work addresses long-context use through efficient processing and memory, while test-time training adapts model parameters as context arrives. The approaches differ in what they change and retain.

Test-Time Training and fast weights

  • Test-time training research has focused on update objectives, optimizers, and parallelization, while retaining fixed write schedules that ignore task relevance.Every token or chunk is written with the same strength regardless of relevance.Source: §6 · p. 8

Efficient long-context modeling

  • Efficient long-context methods make long inputs processable, but benchmark studies find that models exploit less context than their advertised windows allow.Test-time adaptation is described as orthogonal: it changes what the model retains and can be combined with efficient processing methods.Source: §6 · p. 8

Memory and gating in sequence models

  • External retrieval and memory modules store information outside model parameters, while gates in linear-recurrent models control retention in a parametric state.External systems typically decide what to read at answer time through similarity retrieval.Source: §6 · p. 9

7 Conclusion

Select-TTT conditions fast-weight writes on the user query using scores from the frozen output head, while preserving In-Place TTT’s chunk-wise, context-parallel computation. It consistently improves over In-Place TTT and zero-shot baselines across two backbones and benchmarks; the authors view query-aware writes as a step toward test-time memory for evolving contexts.

7 Conclusion

  • Across two backbones and two benchmarks, Select-TTT consistently improves over In-Place TTT and zero-shot baselines in inference-time and continual-training settings.Its selective learning-rate rule scores chunks in the frozen output head’s logit space and rescales their fast-weight writes.Source: §7 · p. 9
  • Select-TTT preserves In-Place TTT’s chunk-wise, context-parallel computation and adds only a lightweight scoring pass.The authors present it as a model-internal way to select hidden states when conventional text-based retrieval signals are unavailable or poorly matched.Source: §7 · p. 9
  • The authors expect query-aware test-time adaptation to extend to agentic web applications with long, evolving contexts.They also frame the method as a step toward query-aware test-time memory.Source: §7 · p. 9

8 Limitations

The paper identifies limits in Select-TTT’s relevance estimate and evaluation scope, while its algorithm retains context-parallel processing through weighted, causal update aggregation.

8 Limitations

  • Query-CE is a heuristic relevance estimate, and applying the frozen LM head to states from multiple TTT layers may affect query-token score calibration.The authors note that learned relevance could improve the approach but may add computation, parameters, and training overhead.Source: §8 · p. 9
  • Evaluation covers two English benchmarks and backbones up to 8B parameters; larger models, multilingual settings, and other modalities remain future work.The method also uses fixed chunk sizes to preserve context parallelism; semantically adaptive chunking could improve boundaries and relevance estimates.Source: §8 · p. 9

A Algorithmic Details

  • The context-parallel algorithm computes chunk relevance scores and standard updates, converts scores to learning-rate multipliers, and aggregates weighted updates with a causal prefix sum.Each chunk’s output applies only preceding updates.Source: A Algorithmic Details · p. 10

B.1 Backbone Architectures

The experiments use two decoder-only Transformer backbones, and Select-TTT’s module configuration and initialization are specified alongside separate evaluation and continual-training datasets.

B.1 Backbone Architectures

  • Both decoder-only backbones use SwiGLU gated MLPs and Rotary Position Embeddings; Qwen3-4B additionally uses YaRN, while Llama-3.1-8B-Instruct uses its native context window.Table 5 summarizes their architectures.Source: §B.1 · p. 10, Table 5 · p. 10

B.2 TTT Module Configuration

  • The TTT module is inserted into the gated MLP of every sixth layer, where the down-projection is treated as the fast weight.Table 6 lists the TTT module and Select-TTT scoring configuration.Source: §B.2 · p. 10, Table 6 · p. 11
  • At continual-training initialization, zero-initialized convolution and a sparse diagonal target projection make the fast-weight update negligible, preserving the pretrained down-projection initially.The update mechanism emerges as training progresses.Source: §B.2 · p. 11

B.3 Dataset Statistics

  • Inference-time evaluation and continual pre-training use separate datasets, and the continual-training corpus is filtered by language and context length.Source: §B.3 · p. 11

B.3.1 Inference Benchmarks.

The appendix specifies the long-term conversational-memory benchmarks and filtered continual-pretraining data, then gives training and query-evaluation procedures. These details define the settings used across the experiments.

B.3.1 Inference Benchmarks.

  • LongMemEval-S covers 500 questions over timestamped histories of about 40 sessions and 115K tokens; LoCoMo covers 1,986 questions across ten dialogues of 19–32 sessions.LongMemEval-S uses six question types and official per-subtask accuracy; LoCoMo uses five query types and reports overall score by query type.Source: §B.3.1 · p. 11, Table 7 · p. 11

B.3.2 CPT Dataset.

  • Continual pretraining uses the English LongCite-45k subset with context lengths from 8K to 65K tokens.The source corpus has approximately 44,600 long-context QA instances with sentence-level citation annotations and Chinese and English examples.Source: §B.3.2 · p. 11, Table 8 · p. 11

B.4 Training Hyperparameters

  • Within each backbone, IP-TTT and Select-TTT share the optimizer, global batch size, TTT learning rate, and chunk size; their per-chunk multiplier differs.Experiments use Nvidia H200 GPUs.Source: §B.4 · p. 11, Table 9 · p. 12

B.5 Query Extraction and Evaluation

  • Questions are tokenized with each backbone’s tokenizer; LongMemEval-S is scored by official per-subtask accuracy and LoCoMo by query-type score.The Section 5.5 relevance-signal analysis uses LongMemEval-S ground-truth answer-chunk annotations.Source: §B.5 · p. 12

C.1 Full Numerical Results of Continual-Training

The appendix provides per-subtask LongMemEval-S accuracy results for continual-training methods across checkpoints and both backbones.

C.1 Full Numerical Results of Continual-Training

  • Tables 10 and 11 report per-subtask accuracy for CPT, IP-TTT, and Select-TTT at each checkpoint from 1k to 5k steps on both backbones.All reported values are percentages.Source: §C.1 · p. 12
  • Table 10 reports continual-training results on LongMemEval-S with Qwen3-4B, with values expressed as percentages.The supplied passage identifies the table but does not state its cell values.Source: Table 10 · p. 12
  • Table 11 reports LongMemEval-S continual-training results with Llama-3.1-8B-Instruct; bold average values mark the best result at each training step.Avg. is the question-weighted overall accuracy, and values are percentages.Source: Table 11 · p. 12

C.2 Full Numerical Results of the Ablation Study

The appendix reports LongMemEval-S ablations of temperature, pooling ratio, and pooling strategy for both backbones. It also describes planned resource release and the authors’ use of language models.

C.2 Full Numerical Results of the Ablation Study

D Online Resources

  • The authors state that source code, training and evaluation configurations, and Select-TTT-enhanced checkpoints will be released upon publication.The release is planned for the project repository.Source: D Online Resources · p. 13

E Usage of LLMs

  • Language-model assistance was limited to grammar and phrasing suggestions; the authors prepared and verified the technical content, experiments, and final text.The authors state that they take full responsibility for the paper.Source: E Usage of LLMs · p. 13

Related papers

Loading 2610.36794v1…