Doc-to-LoRA: Learning to Instantly Internalize Contexts

In plain words

Putting a long document into a model's prompt makes every later answer slower and more memory-hungry, because the model keeps re-reading the document. A known fix is to teach the model the document by training its weights on it, but that training takes tens of seconds or several minutes per document.

Doc-to-LoRA trains a second, helper network once. You give it a document, and in one forward pass it writes a small set of weight tweaks (a LoRA adapter) for the model. The model then answers questions about the document without the document in its prompt. The helper learns this by imitating what the model would answer with the document in front of it, so it has absorbed a fast version of that training process.

This matters when the information changes often, because each new document would otherwise need its own slow training run. In the paper's experiments the update takes under a second.

TL;DR

• D2L is a hypernetwork that maps a document to a rank-8 LoRA on the down-projection of a frozen LLM, in one pass §3.2.
• It is meta-trained to reproduce context distillation, using self-generated questions and the base model's answers §3.1 app. B.1.
• Author-reported: it beats vanilla context distillation under small query budgets and updates in under a second, but it trails the oracle that trains per document tab. 1 tab. 4.
• Its inference memory is much lower than in-context reading on long inputs fig. 4 §5.1.2.
• The paper and the released code disagree on several details; see Paper vs code.

Problem

In-context learning is transient and memory-intensive: long prompts raise latency through attention cost and KV-cache growth, and answer quality tends to fall with longer context §1. Supervised fine-tuning removes the prompt but needs a task dataset and repeated runs whenever the information changes §1. Context distillation (CD) trains a model without the context to imitate its own outputs with the context present, which internalizes the information, but each training run is slow and memory-hungry §1 §2. The paper's goal is internalization that costs one forward pass, so it can be repeated for every new document §1.

Assumptions: the base model is fixed; the learned mapping must work for queries it has not seen; and the training signal is query-independent, meaning a set of queries and self-generated answers per context rather than a single query §2.

Method

Overview. A frozen target LLM first reads the context and exposes its per-layer token activations. A Perceiver-style hypernetwork turns those activations into LoRA matrices for the target model's down-projection layers. Long contexts are split into chunks. Each chunk yields its own adapter, and the adapters are concatenated along the rank dimension, so the output shape does not depend on the context length §3.2. The hypernetwork is trained across many contexts so that the resulting adapted model matches the output of the same model that sees the context directly §3.1.

Steps.

1. Feed the context through the frozen LLM and keep each layer's hidden states Z_l, with Z_0 the token embeddings §3.2.
2. For each layer l, cross-attend r learnable, input-independent latent queries to Z_{l-1} to get latent vectors. The paper describes a cross-attention stack with no self-attention, 8 blocks, and 309M trainable parameters §3.2 §5.
3. Map each latent to one row of A_l and one column of B_l. The paper describes two per-layer output heads; the released code uses one head that emits both (see Paper vs code) §3.2.
4. For a context split into K chunks, concatenate the chunks' A matrices stacked vertically and their B matrices side by side, giving rank r·K §3.2.
5. Apply the adapter to each layer as W' = W + α_l B_l A_l, where α_l is a learned per-layer scalar §3.2.
6. Train with the query-independent objective below, then use the adapter for any query about the context §3.1.

# Internalize a context into LoRA weights  [§3.2]
# Input:  context tokens (N,); frozen LLM; trained hypernetwork
# Output: per layer, LoRA factors A and B for the down-projection

def internalize(context_tokens, llm, hypernet):
    # Split the context into chunks of at most 8192 tokens  [§5]
    chunks = split_into_chunks(context_tokens, max_len=8192)
    # Run the frozen LLM once and keep every layer's hidden states  [§3.2]
    activations = llm.hidden_states(chunks)  # (K, L, N, D)
    # Cross-attend learned latent queries to each layer's states  [§3.2]
    latents = hypernet.perceiver(activations)  # (K, L, r, d_latent)
    # Map each latent to a row of A and a column of B  [§3.2]
    lora_a, lora_b = hypernet.head(latents)
    # Concatenate the chunks along the rank dimension  [§3.2]
    lora_a = concat_rank(lora_a)  # (L, r*K, d_in)
    lora_b = concat_rank(lora_b)  # (L, d_out, r*K)
    return lora_a, lora_b

Training objective. The student is the base model with the generated adapter and no context. The teacher is the same base model with the context in its prompt. The hypernetwork parameters phi are optimized to reduce the divergence between the two over a set of (context, queries, responses) triples §3.1:

\min_{\phi}\;\mathbb{E}_{(c,\mathcal{D}_c)}\,\mathbb{E}_{(x,y)\sim\mathcal{D}_c}\;\mathrm{KL}\big(p_\theta(y\mid x,c)\,\Vert\,p_{\theta+H_\phi(c)}(y\mid x)\big) \quad [\S 3.1]

The paper states the query-independent CD objective as eq. 1 and eq. 2; the meta objective is written in §3.1.

Data. About 900 million tokens of FineWeb-Edu form one set of contexts. The passage QA sets PwC, SQuAD, ROPES and DROP are added; after filtering out passages over 10,000 characters, the paper reports about 3.2 million unique contexts for the combined corpus app. B.1. Each FineWeb-Edu context gets 10 questions from gemma-3-12b-it, written in two rounds of five app. B.1. Each question gets one response from gemma-2-2b-it, and the top-16 logits per generated token are stored as the teacher target app. B.1. Generated answers to the questions are discarded app. B.1.

Training schedule. Stage 1 trains for 80K gradient steps where every context is one chunk. Stage 2 trains for 20K steps with random chunking: 50% one chunk, 12% two, 37.5% three to eight with equal odds. The paper says the two-stage setup was needed because training was unstable with multiple LoRAs from the start app. B.2. Both stages pack contexts into 4K-token sequences with gradient accumulation, for more than 200K context tokens per update across 8 GPUs app. B.2.

Needle-in-a-haystack variant. Trained on 32 to 256 token contexts with the same chunk probabilities, using cross-entropy on ground-truth tokens instead of the distillation loss, and the 6th layer's activations of gemma-2-2b-it as the only input app. A §4.

# Meta-training step on one context  [§3.1]
# Input:  context c; queries x_i; teacher top-16 logprobs per answer token
# Output: scalar loss for the hypernetwork parameters phi

def training_loss(context, queries, teacher_topk, hypernet, base_llm):
    # Generate the adapter for this context in one pass  [§3.2]
    lora = hypernet.generate(context)
    # Student is the base model with the adapter, and no context  [§3.1]
    student_logits = base_llm.with_lora(lora).forward(queries)
    # Match the teacher's stored top-16 distribution at each token  [§3.1]
    loss = divergence(teacher_topk, student_logits)
    return loss

Hyperparameters

Results

Author-reported. Setup: the main base model is gemma-2-2b-it with the context available as the ICL reference; the metric is word-level ROUGE-L F1, reported relative to the base model with the full context §5 §5.1. Hardware is a single H200 GPU app. B.2. Baselines are CD oracle (trained on the target query), CD with 5, 25 or 100 generated queries, T2L, LLMLingua-2 and the base model without context §5.

Long-document QA, 2WikiMultihopQA tab. 1

Reading comprehension, SQuAD (100-sample subset), query budgets tab. 4

Reading comprehension, D2L relative performance (ROUGE-L F1; the text puts SQuAD at 82.5% of the in-context upper bound §5.1; see Paper vs code)

Zero-shot image classification through the VLM tab. 2: the text-only target classifies Imagenette images at 75.03% over 10 classes, against 10% chance. The paper states the target never saw images §5.2.

Long-document QA, other benchmarks tab. 9 tab. 10

Needle-in-a-haystack §4 fig. 2: the base gemma-2-2b-it with 8K context drops sharply beyond 8K tokens. D2L stays close to perfect up to 40 chunks of 1024 tokens (40K tokens), five times the number of chunks seen in training, and degrades gracefully after that. The paper reports the base model needs more than 12 GB of extra memory for a 128K-token haystack, while the internalized model stays below 50 MB §4. No table is given for this experiment.

Query internalization (swapped roles) tab. 3: the target sees the SQuAD passage and the adapter has absorbed the question. Recall is 0.587 against 0.185 without context, but precision falls to 0.044 because answers are verbose §6 tab. 3.

Ablations (rank, data, objective, knowledge interference)

Knowledge interference (SQuAD with the context replaced) tab. 8

Paper vs code

Official code: github.com/SakanaAI/doc-to-lora, commit baa85db4d5df9b29d618af432d6ebf28b3ad5a29, MIT. The paper links this repository §1; the package inside it is named ctx_to_lora.

• Loss. The paper states a KL divergence over the teacher's full distribution §3.1. The code computes minus the sum of the teacher's stored top-16 probabilities times the student's log-probabilities over the full vocabulary. That is cross-entropy restricted to the stored tokens, not the full KL, and the teacher's entropy term is absent trainer.py.
• Teacher responses. The paper says the teacher response is sampled app. B.1. The data paths in the main config say temp_0.0 self_gen_lv1_closed_qa_1_l2l.yaml.
• Blocks. The paper uses 8 cross-attention blocks §5. The main config sets 9 self_gen_lv1_closed_qa_1_l2l.yaml; the main scripts do not override it 1-train.sh. The NIAH script sets 8 1-train.sh.
• Latent queries. The paper uses r learnable latent queries, with r the LoRA rank §3.2. The code's Perceiver uses a fixed n_latent_queries (8 in the main config, 208 in the NIAH script), then one decoder block that produces num_modules × r output queries aggregator.py self_gen_lv1_closed_qa_1_l2l.yaml 1-train.sh.
• Output heads. The paper describes two output heads, one for A and one for B §3.2. The code uses one head producing d_in + d_out values per rank, split into A and B hypernet.py hypernet.py. The code also L2-normalizes each latent before the head and runs one residual MLP before it, which the paper does not mention hypernet.py hypernet.py.
• Per-layer scalers. The paper has one learned scalar α_l per layer §3.2. The code has separate per-layer scalers for A and B, initialized to one and zero. Zero-initialized B means the generated update starts at zero hypernet.py.
• Chunk rank. The paper says concatenated chunks give rank r·K §3.2. The code adds a learned bias adapter of rank r on every combined adapter, so the rank is (K+1)·r lora_merger.py lora_merger.py. The flag is on by default configs.py.
• Regularization. The paper states no penalty on the generated weights. The code adds an L1 penalty on the generated A and B, with coefficient 0.1 in both main stages and 1.5 in the NIAH run trainer.py 1-train.sh 1-train.sh.
• LoRA scaling. The paper gives no scaling factor. The code uses lora_alpha = r^(3/2) × 2, which is 45.25 for r = 8 model_loading.py. The CD baseline in evaluation uses 16 eval_utils.py.
• Context encoder. The paper says the frozen target LLM produces the activations §3.2. The code runs a copy of the base model with its last decoder block removed ctx_encoder.py, quantized to 4-bit NF4 during training model_loading.py 1-train.sh.
• Stage 2 packing and chunks. The paper says batches pack into 4K tokens app. B.2. The stage 2 script packs 1024 input and 2048 context tokens and caps training chunks at 512 tokens 2-train-chunk.sh 2-train-chunk.sh. Stage 1 has no cap; a value of -1 defaults to the model's maximum length train.py.
• Learning rate. The paper gives no learning rate for the main stages app. B.2. The code uses the default 4e-5 in stage 1 and 2e-5 with 2000 warmup steps in stage 2 2-train-chunk.sh.
• NIAH loss. The NIAH appendix uses cross-entropy with use_kl_loss off app. A 1-train.sh. The main text describes only the KL objective §3.1.

No discrepancies were found in the query-generation counts (10 per context, two rounds of five) or the top-16 stored logprobs gen_data.sh self_generate_qa.py.

Pitfalls

• Adapter rank is K+1 times r, not K times r, because of the learned bias adapter. Any memory estimate based on the paper's r·K will be low lora_merger.py.
• The KL in the paper is a top-16 cross-entropy in code. Results depend on the stored support. The data paths indicate temperature 0.0 for the teacher, while the paper says the response is sampled.
• Batched mode uses more update memory than the oracle. On 2WikiMultihopQA, batched D2L peaks at 11.522 GB against 7.820 GB for the oracle; on QASPER at 31.153 GB against 7.665 GB tab. 1 tab. 10. Iterative mode is the low-memory option, at 3.791 GB and 4.781 GB respectively.
• D2L does not beat every CD setting. On MultiFieldQA, D2L (0.481 batched, 0.485 iterative) trails CD with 25 generated queries (0.528) tab. 9. The text claims an advantage only over CD with generated queries on 2WikiMultihopQA §5.1.2.
• The SQuAD numbers do not reconcile. The text says 82.5% relative to the in-context upper bound §5.1. The tables list 0.814 for the same configuration tab. 2 tab. 5 tab. 7. Against the 0.8692 reference tab. 11, 0.814 is about 94%. The card reports both and does not resolve the gap.
• Knowledge interference is large. When the context is an unrelated assistant prompt, D2L falls to 0.096 against a base of 0.201 tab. 8. The paper hypothesizes that the training data only pairs queries with their own context, and suggests adding irrelevant queries as a remedy; it does not test that app. C.4.
• Verbose answers. On the SQuAD swap test, precision is 0.044 against recall of 0.587 tab. 3. ROUGE-L F1 penalizes verbosity, and the paper uses recall for the Mistral and Qwen models for that reason app. E.
• Mistral results are weak. The paper says every internalization method, including D2L, struggles to store long-context information in Mistral-7B-Instruct-v0.2. It gives no explanation app. E.
• Long-context test sets are small. At most 500 samples per dataset are evaluated app. B.2. Latency has a standard deviation, but accuracy has no error bars in the tables.
• NEFTune is on by default. The trainer default noise alpha is 5.0 and the main scripts do not override it; the NIAH script sets it to 0 configs.py 1-train.sh. The paper does not mention it. I did not trace whether it reaches the hypernetwork path, so its effect is unverified.
• Hardware. All latency numbers come from one H200 GPU app. B.2. Latency on other hardware is not reported.

Lineage

• Builds on: context distillation, as in the prior CD line cited in §1 and §2; hypernetworks, Ha et al. (arXiv:1609.09106, cited in the paper's bibliography); LoRA, Hu et al. §3.2; and the Perceiver, Jaegle et al. §3.2. None of these has a card in this wiki yet.
• Compared against: Text-to-LoRA (T2L), a hypernetwork trained with supervised fine-tuning on ground-truth tokens §5; Generative Adapter, a hypernetwork trained with next-token prediction §7 app. E; Cartridges, which learns prefix KV caches (arXiv:2506.06266) §7; and LLMLingua-2, a prompt compressor §5.
• Superseded by: not known from this wiki.