Context distillation

Overview

A frozen base model reads a context, such as a document, and answers queries with it in the prompt. The goal is a model that answers the same queries with no context in its prompt. The invariant all variants share: the adapted model without the context should match the same model with the context present 2602.15902 §3.1. The student is the base model with some update and no context. The teacher is the base model with the context in its prompt.

Context distillation (CD) gets the update by gradient training on the student's outputs against the teacher's. Each document needs its own run. The paper calls this slow and memory-hungry 2602.15902 §1 §2.

Two cards in this wiki use the technique, in different roles. One makes a per-document memory of a document. The other bakes a fixed helpfulness prompt into a model before RLHF 2204.05862 app. A.1. The variants below are compared on their own papers' numbers. Those numbers measure different things, so they are not directly comparable across papers.

Variants

Per-document CD with an oracle query

The training signal uses the target query, which the paper labels as an oracle. It is the strongest row in the paper's tables. On 2WikiMultihopQA it reaches 0.901 normalized performance, at 40.171 s latency and 7.820 GB peak update memory 2602.15902 tab. 1. The oracle uses information a deployed system would not have, so treat it as an upper reference, not a deployable option 2602.15902 §5.1.

CD with generated queries

Without the target query, the training set is built from queries generated from the context. The paper reports results for 5, 25 and 100 generated queries. On 2WikiMultihopQA, 5 queries give 0.704 at 72.537 s and 25 queries give 0.745 at 465.454 s 2602.15902 tab. 1. On SQuAD, 100 queries give 0.650 at 631.101 s 2602.15902 tab. 4. More generated queries cost more time and did not close the gap to the oracle in these runs.

Hypernetwork internalization (Doc-to-LoRA)

A hypernetwork reads the context once and writes a rank-8 LoRA adapter for the base model's down-projection layers 2602.15902 §3.2 §5. The paper calls the method Doc-to-LoRA (D2L) Doc-to-LoRA: Learning to Instantly Internalize Contexts. It is trained once to reproduce CD, using self-generated questions and the base model's answers on about 3.2 million contexts 2602.15902 app. B.1. Two modes are reported:

• Batched. 0.857 normalized performance at 0.209 s and 11.522 GB peak update memory on 2WikiMultihopQA 2602.15902 tab. 1.
• Iterative. 0.844 at 0.551 s and 3.791 GB on the same benchmark 2602.15902 tab. 1. It has the lowest update memory in that table.

The paper reports D2L at 0.814 on SQuAD, 0.655 on DROP and 0.906 on ROPES 2602.15902 tab. 2, with the context-free query internalization and image-to-text results reported separately 2602.15902 §5.2 §6.

Long-haystack memory

For needle-in-a-haystack retrieval, the paper trains the hypernetwork on shorter contexts. The internalized model stays close to perfect up to 40 chunks of 1024 tokens, five times the training chunk count, and the paper reports the base model needs more than 12 GB of extra memory at 128K tokens, against under 50 MB for the internalized model 2602.15902 §4. No table is given for this experiment.

Swapped roles: internalizing the query

The paper also reverses the roles so the adapter absorbs the question and the target sees the passage. Recall rises to 0.587 from 0.185 without context, but precision falls to 0.044 because the answers are verbose [arxiv:2602.15902 tab. 3].

Prompt internalization before RLHF

This is a fixed prompt, not a per-document context. The RLHF paper prepends a helpfulness prompt to general text, takes the teacher's top-50 log-probabilities at each following token, and trains a student without the prompt on a KL loss over those 50 tokens plus one "other" bucket 2204.05862 app. A.1. The training mix is 350M tokens, half pretraining text and half StackExchange, with batches of 32 sequences app. A.1. The resulting model is the starting policy for RLHF 2204.05862 §4.1. The paper gives no speed or memory figures for this step. Its HHH results compare against context-distilled LMs, whose values appear only in a figure 2204.05862 §3.4.1 fig. 5. Card: Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback.

When to use which

The choice depends on how often the information changes, how much latency and memory the caller can spend per update, and how much quality loss is acceptable. The rows below use only the papers' reported numbers.

The hypernetwork route carries a one-time training cost: 80K steps with one chunk per context, then 20K steps with random chunking [arxiv:2602.15902 app. B.2]. The paper does not report that cost in its results tables.

Implementation sketch

The core steps of the hypernetwork route. The base model and the hypernetwork are frozen or trained as described in the paper; the lines marked with a citation are the paper's own design choices.

# Build a per-document adapter in one pass  [arxiv:2602.15902 §3.2]
# Input:  context tokens; frozen base LLM; trained hypernetwork
# Output: LoRA factors A and B for each layer's down-projection

def internalize(context_tokens, base_llm, hypernet):
    # Split long contexts into chunks; each chunk gets its own adapter  [§5]
    chunks = split(context_tokens, max_len=8192)
    # Read the context once and keep each layer's hidden states  [§3.2]
    states = base_llm.hidden_states(chunks)
    # Hypernetwork writes A and B for every chunk and layer  [§3.2]
    a, b = hypernet(states)
    # Concatenate chunk adapters along the rank dimension  [§3.2]
    return concat_rank(a), concat_rank(b)

# Training step for the hypernetwork, on one context  [arxiv:2602.15902 §3.1]
# Student: base model plus the generated adapter, no context
# Teacher: base model with the context in the prompt
def training_loss(context, queries, teacher_outputs, hypernet, base_llm):
    # Generate the adapter for this context  [§3.2]
    a, b = hypernet(base_llm.hidden_states(context))
    # Match the teacher's output distribution on each query  [§3.1]
    student = base_llm.with_lora(a, b).forward(queries)
    return divergence(teacher_outputs, student)

The prompt-internalizing variant 2204.05862 app. A.1 uses the same teacher-student loop with a fixed prompt in place of a document. The card describes no hypernetwork or adapter for it. The teacher's top-50 log-probs and an "other" bucket form the target, not a full distribution.

Contradictions

Within the 2602.15902 paper, not between papers. The card for 2602.15902 flags these and does not resolve them.

• SQuAD D2L numbers. The text puts D2L at 82.5% of the in-context upper bound 2602.15902 §5.1. The main table gives 0.814 for the same configuration 2602.15902 tab. 2, and against the 0.8692 truncated- context reference 2602.15902 tab. 11, 0.814 is about 94%. The SQuAD table on a 100-sample subset gives 0.866 2602.15902 tab. 4. The settings differ, and the paper does not reconcile them.
• Hypernetwork versus CD on MultiFieldQA. D2L trails CD with 25 generated queries there (0.481 and 0.485 against 0.528), while the text claims an advantage only over CD with generated queries on 2WikiMultihopQA 2602.15902 tab. 9 §5.1.2.
• Paper and released code. The paper states a KL objective and rank r·K. The code uses a top-16 cross-entropy and a rank of (K+1)·r. The card records these as implementation differences 2602.15902 §3.1 §3.2. An agent that reimplements the method should take the paper's equations as the design and the code as the tested behaviour, and check which one it follows.

Between papers: none known. The two cards measure different things, per-document QA against HHH and crowd preference, so their numbers do not compete.

Papers

• Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback: uses context distillation to put a fixed helpfulness prompt into the starting policy for RLHF.
• Doc-to-LoRA: Learning to Instantly Internalize Contexts: a hypernetwork that writes a LoRA adapter from a document in one pass, trained to reproduce context distillation.