PaperlayerView exact revision on arXiv

Source-linked AI summary

Encoding Context into LLM Weights: What Makes It Context?

Yue Wang, Qizhou Wang, Feng Liu, Gang Niu, Masashi Sugiyama, Bo Han

arXiv:2610.17686v1cs.CLcs.LG

TL;DR

Long contexts increase inference costs and can exceed input windows, while encoded information may not reliably function as context. ACE associates identifiers with encoded context and prompts the model to use it selectively; it outperformed classical TCE across controlled tasks and long-context benchmarks.

  • Problem

    Long-context inference incurs increasing memory and computation costs, and encoding context into weights does not ensure that an off-the-shelf LLM will reliably use it as context.

    Source: §1 · p. 1, §3 · p. 4

  • Method

    ACE combines UID-conditioned context encoding for explicit access with prompts that specify the encoded information’s contextual role and intended use.

    Source: §3 · p. 4, §3.1 · p. 4

  • Results

    ACE improved over classical TCE across controlled context-use tasks and long-context benchmarks, including raising RULER NIAH-3 accuracy from 8% to 63% on Muse-Glimmer-30B.

    Source: §1 · p. 2, §4.1 · p. 6, §4.2 · p. 9

  • Takeaways & Limitations

    In the evaluated settings, encoding information into weights alone did not ensure reliable contextual use, while ACE improved its use with off-the-shelf LLMs.

    Source: §6 · p. 9

    TCE requires backpropagation and adds computation; further efficiency and implementation improvements are outside the study’s focus.

    Source: §2.3 · p. 4

Abstract

from arXiv · show

Longer contexts allow large language models (LLMs) to condition their predictions on more information, but inference costs grow prohibitively with context length. An emerging alternative uses test-time updates to encode context as temporary changes to model weights. However, encoding this information into weights that already represent general knowledge obscures its role as context, so the LLM may fail to use it properly to condition subsequent inference. To restore this contextual role, we first associate each context with an identifier during encoding, thereby distinguishing the newly encoded information from pre-existing knowledge. Invoking the identifier at inference then enables the LLM to recite the encoded context. Nevertheless, successful recitation does not ensure that the LLM will use the recalled information as context. We therefore further introduce a prompting framework that explicitly guides the LLM to use this information as context. Across diverse controlled tasks and practical benchmarks, our method outperformed the baseline using only off-the-shelf LLMs, showing its advances in context encoding.

1 INTRODUCTION

Longer contexts increase inference costs and can make relevant information harder to use, while context windows limit how much information fits at once. The introduction presents ACE, which addresses encoded context with identifiers and prompting, and reports gains over classical TCE.

1 INTRODUCTION

  • Longer contexts raise KV-cache memory linearly and causal-attention computation quadratically, while models may lose track of relevant information and remain bounded by their input windows.The most advanced models cited reach million-token-scale windows, but effectively usable lengths can be shorter.Source: §1 · p. 1
  • Without meta-training-based fine-tuning, prior TCE remained far inferior to full-context inference, leaving open whether off-the-shelf LLMs can reliably use encoded information as context.The introduction distinguishes memorizing information from using it contextually.Source: §1 · p. 2
  • ACE binds encoded context to unique identifiers and uses prompting to guide its selective use as context, addressing the gap between recalling information and using it to condition responses.It operates on off-the-shelf LLMs without the additional meta-training-based fine-tuning often used by prior TCE methods.Source: §1 · p. 2, Figure 1 · p. 2
  • ACE increased Qwen3.5-27B accuracy from 1.5% to 75.6% on Recall and from 0.0% to 98.0% on Recognition, and improved RULER NIAH-3 accuracy from 8% to 63%.On Muse-Glimmer-30B, ACE also raised LongBench answer F1 from 14.0 to 19.2.Source: §1 · p. 2

2 BACKGROUND

The background frames test-time context encoding as an alternative to keeping long histories in the input, then describes sequential fast-weight updates and the study’s off-the-shelf scope. Context-specific updates can be reset when a context is no longer needed.

2 BACKGROUND

  • Full-context inference generates responses using both prior context C and a subsequent request x when their combined token length fits within the model’s input window.The model can directly use information and interaction state in C.Source: §2 · p. 2

2.1 LONG-CONTEXT INFERENCE

  • Test-time context encoding stores context in a parametric state, potentially preserving more information than compression or KV-cache eviction, which risk discarding task-relevant content.The motivation is long-context cost and declining reliability as input length approaches the context window.Source: §2.1 · p. 3

2.2 TEST-TIME CONTEXT ENCODING

  • TCE encodes temporally ordered context chunks by sequentially optimizing fast weights with an autoregressive loss, then removes the original context from the input.The remaining model weights stay frozen; context-specific fast-weight updates can be discarded by resetting them to their initial state.Source: §2.2 · p. 3

2.3 SCOPE OF THIS STUDY

  • The study restricts its analysis to off-the-shelf LLMs to clarify TCE reliability, while leaving efficiency and implementation improvements beyond its scope.TCE requires backpropagation and therefore incurs additional computation.Source: §2.3 · pp. 3–4

3 NEW METHOD

The method addresses two challenges for off-the-shelf LLMs: making encoded information accessible and specifying its role as context. ACE combines UID-based addressing with prompts for encoding and inference, then applies the resulting encoded weights to later requests.

3 NEW METHOD

  • ACE addresses unreliable access to encoded information and uncertainty about how it should function as context, without relying on task-specific meta-training.The method pairs an addressing mechanism with a prompting framework.Source: §3 · p. 4

3.1 ADDRESSABLE MECHANISM

  • The addressable mechanism assigns each context chunk a unique identifier and trains the model to predict that chunk conditioned on its identifier, providing an explicit access cue at inference.The sequential fast-weight updates are retained while the basic autoregressive loss is replaced with a UID-conditioned objective.Source: §3.1 · p. 4

3.2 PROMPTING ENHANCEMENT

  • The prompting framework clarifies each chunk’s contextual role during encoding and instructs the model at inference to use only information relevant to the current request.The inference prompt avoids reproducing the entire encoded context, which would consume input-window capacity.Source: §3.2 · p. 5

3.3 OVERALL ALGORITHM

  • ACE sequentially updates fast weights across context chunks, generates responses from the encoded weights using a compact UID-based prompt, and resets the updates when the context is no longer needed.The prompt indexes encoded chunks and specifies how they should guide inference without reproducing them.Source: §3.3 · p. 5

4 EXPERIMENTS

ACE improved on classical TCE across controlled context-use tasks and long-context benchmarks, though performance remained limited on harder tasks. The experiments also examined conflict handling, prompt choice, and retention of unrelated capabilities.

4.1 SANDBOX TASKS

  • Sandbox evaluation paired short contexts averaging about 11 tokens with queries, targeted 200 instances per task or subtask, and separately tested Retention on 100 unrelated queries.Tasks distinguished content-cued Recall, recency-based Recognition, pending-question Continuation, and conflicts with pre-existing knowledge.Source: §4.1 · p. 6
  • ACE substantially improved performance across sandbox context-use tasks, while classical TCE failed to use encoded context reliably even on Recall.The comparison covered Recall, Recognition, Continuation, and Retention on Qwen3.5 models.Source: §4.1 · p. 6, Table 1 · p. 7
  • ACE reached near-perfect accuracy in several Conflict settings and substantially outperformed TCE on two-sided Counterfactual, where both methods performed worse.Two-sided Counterfactual required correct answers both from encoded context and from pre-existing knowledge when instructed to disregard that context.Source: Table 2 · p. 7, §4.1 · pp. 7–8
  • Invocation-prompt performance varied: ci_short often led on simpler sandbox tasks but not consistently on RULER, while Retention accuracy remained high with only the final FFN as fast weights.Larger models did not necessarily perform better.Source: §4.1 · p. 7

4.2 LONG-CONTEXT TASKS

  • ACE improved over TCE across RULER, improved on LONGBENCH, and matched or exceeded it on LONGBENCH V2, while harder tasks remained challenging.On RULER, TCE performed well on NIAH-1 but dropped sharply on harder NIAH variants; ACE remained far from perfect on NIAH-2 and NIAH-3.Source: Table 3 · p. 8, §4.2 · p. 9
  • Long-context tests covered RULER NIAH at 1,024 tokens, LONGBENCH at 2k–7k, and LONGBENCH V2 at about 10k–15k, with contexts chunked sequentially up to 1,024 tokens.LONGBENCH V2 responses were judged on a 0, 0.5, or 1 scale.Source: §4.2 · pp. 8–9

5 RELATED WORK

Related work spans test-time adaptation that retains context, methods that replace context with weight updates, and alternative objectives for encoding it.

5 RELATED WORK

  • Conventional test-time training adapts models to available context but generally retains that context or its KV cache during generation, unlike context encoding.Replacing the input context requires the model to retain information in its weights and use that implicit representation in later inference.Source: §5 · p. 9
  • Task-specific meta-training in context-replacement methods leaves unclear whether strong results reflect general encoding ability or adaptation to the meta-training distribution.The passage notes that no ablation study had formally examined this distinction.Source: §5 · p. 9
  • KV-association objectives optimize token matching and may rely more on meta-training, while context distillation requires diverse context-query pairs and more test-time data.The cited discussion says KV association typically underperforms autoregressive encoding and distillation raises practical cost and efficiency concerns.Source: §5 · p. 9

6 SUMMARY

The paper presents ACE as a way to improve reliable context encoding with off-the-shelf LLMs, while identifying source selection and complex reasoning as areas needing improvement.

6 SUMMARY

  • ACE combines UID-conditioned encoding, which provides explicit access to encoded information, with prompting that specifies its contextual role and intended use.The paper contrasts this with encoding information into weights without ensuring that it will be used as context.Source: §6 · p. 9
  • Across controlled sandbox tasks and long-context benchmarks, ACE improved over classical TCE using off-the-shelf LLMs.The paper describes ACE as a promising step toward reliable context encoding.Source: §6 · p. 9
  • The paper identifies source selection and more complex long-context reasoning as areas for future improvement.The conclusion says substantial room for improvement remains.Source: §6 · p. 9

A DATASET EXAMPLES AND PROMPT TEMPLATES

The appendix illustrates sandbox datasets and context-invocation prompts, then presents analyses of UID conditioning for inference and encoding.

A DATASET EXAMPLES AND PROMPT TEMPLATES

  • The appendix supplies sandbox dataset examples and prompt templates for accessing and using encoded information during inference.The examples show encoded contexts, subsequent queries, and expected answers.Source: A DATASET EXAMPLES AND PROMPT TEMPLATES · p. 14

A.1 SANDBOX DATASET EXAMPLES

  • Sandbox examples span Recall, Recognition, Continuation, three Conflict categories, and Retention.The Conflict examples include Counterfactual, Dispute, and Outdated Information; the Outdated Information example uses the Qwen3.5-4B dataset.Source: §A.1 · p. 14
  • Conflict examples retain the wording “context above” even though the context is removed from the input after encoding.The examples contrast contextual-answer queries with a prior-knowledge query that asks for the generally known capital of Peru.Source: §A.1 · pp. 14–15
  • Recall asks for an attribute value, Recognition asks what information was just given, and Continuation asks the model to answer an encoded unresolved question.The appendix also illustrates unrelated arithmetic and general-knowledge queries for Retention.Source: §A.1 · p. 14

A.2 CONTEXT-INVOCATION PROMPTS

  • The ci_short prompt minimally asks the model to recall context, while ci_v1 treats timestamped chunks as prior conversational context and directs selective use.Each prompt is followed by the task query, supplied as a user message.Source: §A.2 · p. 15
  • The ci_v2 variant adds an instruction to continue an unresolved request or incomplete task in the latest chunk when appropriate.The prompt variants were evaluated while keeping the encoding procedure and role-binding prompt fixed.Source: §A.2 · p. 15

B MATHEMATICAL INSIGHTS

  • The appendix analyzes how UIDs help distinguish and select encoded contexts during inference and reports a proof-of-concept for stronger gradients and faster learning during encoding.These analyses address UID conditioning from complementary inference and encoding perspectives.Source: B MATHEMATICAL INSIGHTS · p. 16

B.1 CONTEXT DISTINCTION

UID conditioning can distinguish encoded contexts from one another and from pre-existing knowledge, with theoretical conditions for improved generation. The analysis also identifies simplified conditions under which UID conditioning can speed encoding, while sandbox ablations show the contributions of addressing, encoding losses, and invocation prompting.

B.1 CONTEXT DISTINCTION

  • When the UID-conditioned encoding error is below H(π), UID conditioning guarantees a positive average log-probability gain; for uniform weights, the sufficient threshold is log T.At ideal optima with uniform weights, the probability ratio is T and the loss gap is log T.Source: §B.1 · pp. 16–17

B.2 SOURCE DISTINCTION

  • With pre-existing knowledge included, an excess loss below I(ẽC; Z) is sufficient for a positive average log-probability gain, but the bound does not guarantee reliable source selection for arbitrary downstream requests.The information term captures distinctions among encoded contexts and between encoded contexts and pre-existing knowledge.Source: §B.2 · pp. 18–19

B.3 ENCODING BENEFITS

  • Under the stated conditions, UID conditioning increases gradient norms and accelerates the same loss trajectory by a factor of γ.This is a proof of concept under simplified training dynamics, not a guarantee for full LLM training.Source: §B.3 · pp. 20–21

C.1 EFFECT OF THE ADDRESSING MECHANISM

  • TCE-AM outperformed original TCE overall, with some Recall exceptions; ACE’s prompting framework achieved substantially better performance by making intended context use more explicit.Without the UID at inference, TCE-AM improved Recognition on 4B and 9B, while its Recall effect was mixed.Source: §C.1 · p. 22

C.2 CONTRIBUTIONS OF ENCODING AND INVOCATION

  • Removing either encoding loss reduced performance, and ACE without encoding scored zero on both tasks, showing the invocation prompt alone was insufficient.Using original TCE inference after ACE encoding still performed strongly in some cases, but the full method performed better.Source: §C.2 · p. 22

C.3 CHOICE OF FAST WEIGHTS

The best fast-weight configuration depends on the model, method, and task. Attention updates can substantially improve ACE’s Recall and Continuation performance over the default FFN-1 setting.

C.3 CHOICE OF FAST WEIGHTS

  • 72.0% on Qwen3.5-4B and 68.5% on Qwen3.5-9B: ATTN-3 raised ACE Continuation accuracy from 48.8% and 13.7%, respectively, relative to FFN-1.Source: §C.3 · pp. 22–23
  • ATTN-2 and ATTN-3 outperformed every FFN configuration on Recall and Continuation across both evaluated models.Source: §C.3 · pp. 22–23, Table 6 · p. 23, Table 7 · p. 24
  • More FFN layers improved ACE Recognition but had mixed effects on other tasks, while TCE’s best configuration varied by task.TCE’s attention updates yielded its highest Recall but zero Recognition accuracy; FFN-3 performed best on Continuation.Source: §C.3 · pp. 22–23
  • ACE Retention stayed at 100% across configurations on Qwen3.5-9B but varied on Qwen3.5-4B, where ACE was more stable than TCE.Source: §C.3 · p. 23

D MORE RESULTS ON SANDBOX TASKS

These sandbox analyses test ACE’s threshold, update-budget, invocation, UID, and wording choices across context-use tasks and models. They show that ACE’s advantages are robust in several comparisons, while particular settings and effects vary by task and model.

D.1 SENSITIVITY TO THE ENCODING THRESHOLD

  • At larger stopping thresholds, Recall and Recognition generally fell, especially Recognition, while two-sided Counterfactual accuracy tended to improve.Threshold effects on Continuation varied across models and prompts.Source: §D.1 · p. 25, Table 9 · p. 26, Table 10 · p. 26

D.2 COMPARISON AT FIXED UPDATE BUDGETS

  • ACE exceeded TCE in 35 of 36 comparisons at fixed budgets K ≥4, with one tie; Recognition reached 98.5–99.5% by K = 8 with ACE.TCE improved with more updates on Recognition and generally on NIAH-2, but remained below ACE at larger budgets; the comparison matched update counts, not computational costs.Source: §D.2 · p. 25, Table 16 · p. 29, Table 17 · p. 30

D.3 EFFECT OF INVOCATION ON PRIOR-KNOWLEDGE RECOVERY

  • Removing ACE’s invocation had mixed effects on prior-knowledge recovery: accuracy rose 6.0 points on Qwen3.5-4B but fell 4.5 and 3.5 points on 9B and 27B.The prompt’s removal alone did not consistently improve prior-knowledge recovery.Source: §D.3 · p. 26, Table 19 · p. 30

D.4 TCE WITH GENERIC INVOCATION

  • A generic recall instruction added to TCE had limited, inconsistent benefits, with Recognition at 0.0–1.5% and NIAH-2 at 5.0–10.0%, substantially below ACE.It improved Recall on Qwen3.5-9B and Qwen3.5-27B, but did not overcome TCE’s limitations in these settings.Source: §D.4 · p. 27, Table 20 · p. 31

D.5 UNIQUE VERSUS SHARED UIDS

  • Using the shared UID “memory” reduced Recognition accuracy on all three Qwen3.5 models, favoring the original UID design in this single-context setting.Source: §D.5 · p. 28, Table 21 · p. 31

D.6 EFFECT OF ENCODING WORDING

  • Replacing the role-binding wording had mixed effects: Recall stayed similar, while Recognition and NIAH-2 shifted in different directions across models.The alternative retained UID and timestamp occurrence counts and matched the original prefix’s token length.Source: §D.6 · p. 29, Table 22 · p. 31

E MORE RESULTS ON LONG-CONTEXT TASKS

These additional long-context evaluations examine threshold sensitivity, general capability retention, UID addressing, and retrieval at longer context lengths. They find that ACE’s gains persist across thresholds and longer contexts, with important scope qualifications.

E.1 SENSITIVITY TO THE ENCODING THRESHOLD

  • At every reported threshold, ACE outperformed TCE on LongBench; on RULER, smaller thresholds generally improved retrieval, particularly on harder tasks, but effects varied.The best threshold depended on model and task.Source: §E.1 · p. 32, Table 24 · p. 33, Table 25 · p. 34

E.2 GENERAL CAPABILITY RETENTION

  • Larger models generally matched or exceeded TCE on MMLU-200, while smaller models often scored lower; without original-model matching, results describe post-encoding performance, not capability change.The evaluation also tested ACE with invocation prompts referring to context unrelated to the MMLU questions.Source: §E.2 · p. 33, Table 26 · p. 34, Table 27 · p. 35

E.3 EFFECT OF THE ADDRESSING CUE

  • On RULER NIAH-1, replacing the correct UID with a length-matched decoy reduced retrieval accuracy from 84–100% to 1–8% with the encoded state fixed.Source: §E.3 · p. 33, Table 30 · p. 36

E.4 RULER WITH LONGER CONTEXTS

  • At 64K-context NIAH-1, ACE outperformed TCE on all three models, reaching 80.0% on Qwen3.5-9B and 70.0% on Qwen3.5-27B; it also won all eight 4K NIAH-2/3 comparisons.Performance remained below full-context inference.Source: §E.4 · p. 34, Table 31 · p. 36, Table 32 · p. 36

F COMPUTATIONAL EFFICIENCY

The efficiency comparison is supplementary to the paper’s focus on reliable context use, and its runtime measurements are based on selected examples. ACE uses more encoding computation than TCE but can reduce inference compute and KV-cache storage relative to full-context inference; the qualitative examples compare model responses on sandbox and RULER tasks.

F COMPUTATIONAL EFFICIENCY

  • Efficiency evidence is limited to one fixed example per benchmark family, and practical runtime gains may differ from the asymptotic advantage because implementation affects runtime.The paper’s primary focus is reliability, not computational efficiency.Source: F COMPUTATIONAL EFFICIENCY · p. 37
  • On selected LongBench and LongBench v2 examples, ACE took approximately 1.5–2.1× TCE’s encoding time; one Qwen3.5-4B sandbox example had shorter ACE responses, not faster per-token decoding.The sandbox responses averaged 17 tokens in 0.88 seconds for ACE versus 123 tokens in 4.74 seconds for TCE.Source: F COMPUTATIONAL EFFICIENCY · p. 37, Table 34 · p. 38, Table 35 · p. 38
  • On the selected LongBench v2 example, ACE used approximately 20–24× fewer inference FLOPs per query than full-context inference, excluding encoding and assuming no KV-cache reuse.ACE required 2.1–2.5× TCE’s encoding FLOPs on that example.Source: F COMPUTATIONAL EFFICIENCY · p. 38, Table 37 · p. 39
  • ACE used approximately 22–27× less KV-cache storage than full-context inference on selected LongBench and LongBench v2 examples.The measurements cover attention keys and values after prompt prefill, not total or peak GPU memory; on the short sandbox example, ACE’s cache was slightly larger.Source: F COMPUTATIONAL EFFICIENCY · p. 38, Table 38 · p. 40, Table 39 · p. 41

G QUALITATIVE RESULTS

  • The qualitative examples show encoded context, query, reference answer, and model responses, with correctness marks and highlights indicating response differences.Source: G QUALITATIVE RESULTS · p. 42

G.1 SANDBOX TASKS

In the sandbox and RULER examples, ACE recovered encoded information and followed contextual instructions where TCE sometimes denied access, corrupted details, or favored general knowledge.

G.1 SANDBOX TASKS

  • ACE recovered an encoded statement without an entity or attribute cue, while TCE denied having received prior information.Source: §G.1 · p. 42
  • ACE recovered and correctly answered an unanswered arithmetic question, whereas TCE denied access despite an explicit continuation instruction.Source: §G.1 · p. 42
  • When asked to use an encoded update about the Bayern coach, ACE returned Vincent Kompany while TCE returned outdated answer Thomas Tuchel.Source: §G.1 · p. 43
  • ACE followed a contextual attribution, returned conventional prior knowledge when asked to set context aside, and answered an unrelated control correctly; TCE failed to recover the contextual attribution.The example tests source selection, not the historical dispute.Source: §G.1 · p. 43

G.2 RULER

  • ACE retrieved the exact seven-digit number, whereas TCE returned the familiar answer “42” using general knowledge rather than encoded context.Source: §G.2 · p. 44
  • ACE recovered the complete UUID after sequential chunk-wise encoding, while TCE reproduced its first two fields and corrupted the remaining three.TCE also described its answer as a placeholder.Source: §G.2 · p. 44

Related papers

Loading 2610.17686v1…