[LLM 6/10] Context Distillation: Moving the System Prompt into the Model's Weights

Every time a user messages your chatbot, you attach the same system prompt, hundreds of tokens long — every request, for the lifetime of the system, paid again and again with no end. This chapter moves that block of knowledge from the prompt into the model's weights, using a technique called Context Distillation in its on-policy version (OPCD). The most beautiful part: the teacher and the student are the exact same model — the only thing that differs is who gets to see the prompt.

Open in Colab06_context_distillation.ipynb

Thai LLM seriesPart 6 of 10

1. 1Continue Pretraining
2. 2SFT and LoRA
3. 3RLHF and PPO
4. 4DPO: Direct Preference Optimization
5. 5GRPO
6. 6Context Distillation(you are here)
7. 7Model Distillation
8. 8Guardrails
9. 9Benchmarking
10. 10Deployment

1. The Problem

A typical Thai customer-service assistant has a system prompt that goes something like this: define a persona, require answers in Thai at all times, be polite and end sentences with the proper particles, never give medical or legal advice. Written out properly, that's about 400 tokens — and it is sent with every request.

Run the numbers: a system serving 100,000 requests a day pays for the same identical block of text 40 million tokens a day, 1.2 billion a month — while the content never changes by a single character. And that's before two prices that never show up on the bill:

• Latency — the model must prefill 400 tokens before it can start thinking about the first token of every answer
• Context budget — every persona token is space taken away from conversation history and attached documents

Framed in this series' terms, knowledge has three places it can live, each with a different payment schedule:

Where the knowledge lives	When you pay	Best for
System prompt	Every request, forever	Behavior/policy that still changes often
RAG	Every request (retrieval + a long prompt)	Large volumes of facts that change often and need source citations
Model weights	Once, at training time	Behavior/policy that has settled

A system prompt that has settled but still rides along on every request is knowledge stored in the wrong place — it belongs in the last row of this table, not the first. This chapter is how you move it.

2. What We're Going to Do

Context distillation trains a student that does not see the context $c$ to behave like a teacher that does see $c$ — put differently, it moves the effect of $c$ from the prompt into the weights. The original offline idea goes back to Askell et al. (2021); the version we use in this chapter is OPCD (On-Policy Context Distillation) by Ye, Dong, Wu, Huang and Wei (2026, arXiv:2602.12275), which adds two key ingredients that section 3 will take apart one at a time:

1. The student samples its own answers (on-policy), without seeing $c$
2. On those answers, minimize the reverse KL against the teacher who sees $c$

The core idea of this chapter

A system prompt is knowledge stored in the wrong place — kept in the prompt, you pay every request, forever. OPCD moves it into the weights, and you pay once, at training time.

And in this chapter, teacher and student are the same set of weights — the teacher is the model with $c$ in front of it; the student is the same model without $c$. What the distance between the two measures is the "influence of $c$," pure and simple.

Let me pin one sentence here, because chapter 7 will discuss another "distillation" that people confuse with this one constantly:

Context distillation changes "what the model knows without being told" — model distillation changes "the size of the model."

In this chapter the model does not shrink by a single parameter. It simply stops needing the prompt. Chapter 7 is about compressing a large model into a small one — an entirely different axis.

3. The Equations

3.1 The OPCD objective

$$\mathcal{L}(\theta) = \mathbb{E}_{(x,c),\; y\sim\pi_\theta(\cdot|x)}\left[\frac{1}{|y|}\sum_{t=1}^{|y|} \mathbb{D}_{\text{KL}}\Big(\pi_\theta(\cdot \mid x, y_{<t}) \,\Big\|\, \pi_{\text{teacher}}(\cdot \mid c, x, y_{<t})\Big)\right]$$

where the KL at each token position is a sum across the whole vocabulary $\mathcal{V}$:

$$\mathbb{D}_{\text{KL}}\Big(\pi_\theta \,\Big\|\, \pi_{\text{teacher}}\Big) = \sum_{v\in\mathcal{V}} \pi_\theta(v \mid x, y_{<t})\,\log\frac{\pi_\theta(v \mid x, y_{<t})}{\pi_{\text{teacher}}(v \mid c, x, y_{<t})}$$

• $c$ = the context we want moved into the weights (persona + safety policy, ~400 tokens)
• $x$ = the user's question; $y$ = the answer sampled by the student itself, without seeing $c$
• $\pi_\theta$ = the student (predicting from $x$ alone); $\pi_{\text{teacher}}$ = the teacher (the original weights, but seeing $c$ too)
• $\frac{1}{|y|}$ = per-token averaging so long answers don't carry extra weight (sound familiar? — length bias, from chapter 4)

Notice this is not cross-entropy against any "answer key" — the target is the teacher's entire row of probabilities at every token position. The student isn't learning "what is the next word"; it is learning "if $c$ were sitting up front, what would the probability of every word in the vocab look like?"

The equation carries two decision points that bear the whole method's weight. Take them one at a time.

3.2 Decision one — the KL must be reverse ($\pi_\theta$ in front)

KL is asymmetric, and its order is a choice of behavior:

• Forward KL $\mathbb{D}_{\text{KL}}(\pi_{\text{teacher}} \| \pi_\theta)$ blows up when the teacher has mass and the student has none → the student is forced to "cover" every mode of the teacher (mode-covering). With insufficient capacity, it spreads mass thin to blanket everything — including the valleys between modes where the teacher never goes — which in LLM terms is the answer that "blends two styles into something wrong," or a hallucination.
• Reverse KL $\mathbb{D}_{\text{KL}}(\pi_\theta \| \pi_{\text{teacher}})$ blows up when the student puts mass where the teacher has none → the student is forced to never do what the teacher doesn't do, then commit firmly to one of the teacher's modes (mode-seeking).

For this chapter's task — a persona and a safety policy — we want the latter without a moment's thought: a student that "matches the teacher one way, reliably" is worth more than a student that "hedges probability across every path the teacher might take, plus the paths the teacher forbids."

If this looks familiar — yes, the KL in the RLHF equations of chapters 3–4 also puts $\pi$ in front, for the same reason: we govern the behavior of the thing being trained, not of the reference.

3.3 Decision two — the rollouts must be the student's own (on-policy)

Look at $y\sim\pi_\theta(\cdot|x)$ in equation 3.1: the answers used for training are sampled from the student, not from the teacher.

The easier alternative would be to have the teacher (who sees $c$) write a batch of answers, then SFT the student on them — and that route has a structural problem named exposure bias: the student is taught only on text trajectories the teacher wrote, but at deployment it must continue from prefixes it itself wrote. One token off, and it lands in a state it was never taught, and the errors compound from there.

On-policy sampling deletes this problem by construction: the states the student meets during training are the same kind of states it will meet at inference, because it authored both. The teacher's one job is to "stand inspection" along the student's path — saying, at this point you've just walked to, here's where you should go next if $c$ were present. (This is the same reason chapter 5 had to sample its own answers instead of continuing with DPO.)

One line of honesty: when computing gradients we treat the sampled $y$ as a constant — no gradient flows back through the sampling. That is standard practice in on-policy distillation.

3.4 The baseline you have to beat: offline context distillation

What most blogs call "context distillation" is the offline version:

$$\mathcal{L}_{\text{offline}}(\theta) = -\mathbb{E}_{y\sim\pi_{\text{teacher}}(\cdot|c,x)}\left[\sum_{t=1}^{|y|}\log\pi_\theta(y_t \mid x, y_{<t})\right]$$

Read it straight: let the teacher who sees $c$ write answers, then SFT the student who doesn't see $c$ on them — plain cross-entropy on the teacher's text. No full-row KL, no on-policy.

This is no strawman. It is a genuinely strong baseline, and cheaper (it trains exactly like chapter 2). Section 9 pits OPCD against it fairly, on the same data. If OPCD's two ingredients (reverse KL + on-policy) are worth anything, it has to win exactly where the theory says it will: generalization to kinds of prompts it never saw.

The full lesson is in the course

This post is roughly the first 30% of the chapter. The rest — environment setup, data preparation, the main code, measured results and the wrap-up — is in the free LLM Finetuning course. Sign in with Google to read it.

Read the full lesson in the course →