[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.

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:

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$

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​

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

• $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:

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.

4. Seeing the Equations​

The whole method in one picture​

Read the figure and notice its economy: no second model, no reward model, no answer-key dataset. Just two kinds of forward pass over one set of weights — one seeing $c$, one not — and a LoRA adapter whose job is to store the "difference" between those two into the weights.

Why the direction of the KL decides the behavior​

The left panel is the definition of hallucination in a distillation context: the best q under forward KL puts real mass exactly where p is nearly zero, because it will pay that price to avoid missing any mode. The right panel is what we want from a safety-minded student: pick a path the teacher endorses, and hold it.

What OPCD buys us​

This chapter's goal, written as geometry: move the blue dot 400 tokens to the left while losing as little height as possible.

Before going on, zoom down to token level: the sentence below is one whose behavior the persona dictates. Watch how the student's per-token log-probs (with no $c$ in sight) change after training — before training, probabilities like these needed $c$ to push them up; after training, they are simply the model's own defaults:

5. Setting Up the Environment​

Open Colab and pick Runtime → Change runtime type → T4 GPU (the free tier is enough).

torch_dtype=torch.float16      # not bfloat16
attn_implementation="sdpa"     # not flash_attention_2

cap = torch.cuda.get_device_capability(0)
print("compute capability:", cap)                    # T4 = (7, 5)
print("native bf16:", cap[0] >= 8)                   # T4 -> False
print("torch says   :", torch.cuda.is_bf16_supported())  # T4 -> True (counts emulation!)

A teacher that costs not one extra byte of VRAM​

OPCD needs both a teacher and a student, which sounds like loading two models — no need. Both are the same set of weights, differing only in adapter and prompt:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model

tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B", padding_side="left")
base = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    torch_dtype=torch.float16,
    attn_implementation="sdpa",
).cuda()

policy = get_peft_model(base, LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    task_type="CAUSAL_LM",
))

• Student = policy (base + LoRA), forwarded without $c$
• Teacher = the same model under policy.disable_adapter(), forwarded with $c$ in front

6. Preparing the Data​

Two things: the context to move into the weights, and questions for the student to practice sampling answers on.

The context: persona + safety policy (~400 tokens)​

# The notebook's persona is written in Thai (~400 tokens); translated, it reads:
PERSONA = """You are "Nong Jaidee", the customer service assistant of an online store.
Rules to follow at all times, no exceptions:
1. Answer in Thai only. Never switch into English mid-sentence.
   Even if the user asks in English, reply in Thai, politely.
2. Always use polite language, ending sentences consistently with the
   Thai politeness particles khrap/kha.
3. Never give medical advice, diagnoses, or medication guidance.
   Decline politely, and recommend consulting a doctor or pharmacist directly.
4. Never give legal advice. Decline politely, and recommend consulting
   a lawyer or the relevant agency.
5. ..."""                      # the full ~400-token Thai version is in the notebook

print("context length:", len(tok(PERSONA).input_ids), "tokens")

This is our $c$ — notice it is pure behavior, with no facts to memorize (an observation that returns as a big deal in the limitations box at the end of the chapter). For readers unfamiliar with Thai: khrap and kha are politeness particles appended to the end of sentences — khrap by male speakers, kha by female — and dropping them is the fastest way for a Thai chatbot to sound rude.

Practice questions: 300 from a Thai dataset​

from datasets import load_dataset

ds = load_dataset("airesearch/wangchanx-seed-free-synthetic-instruct-thai-120k",
                  split="train")
prompts = [r["instruction"] for r in ds.shuffle(seed=42).select(range(300))]

We never use the dataset's answer column, not a single row — OPCD needs no answer key, only diverse questions for the student to practice answering in varied situations, with the teacher inspecting. (This is the same WangchanX Thai instruction set we met in chapter 2.)

The evaluation set is kept separate, untouched during training, and deliberately includes kinds of prompts absent from the training set:

• 40 items: general questions in the same style as training (in-distribution)
• 20 items: medical/legal questions — testing whether the "decline" policy really made it into the weights
• 20 items: English-language questions — testing the "always answer in Thai" rule in the situation most tempting to break it

The last two groups become the OOD compliance column in section 9 — the separator between "memorized the examples" and "absorbed the policy."

7. The Main Code​

The OPCD loop has three beats: the student samples → the teacher inspects → the weights move by reverse KL.

7.1 The student samples its own rollouts (without $c$)​

@torch.no_grad()
def rollout(policy, x_texts, G=4):
    """The heart of the word on-policy: answers come from the student, not the teacher"""
    batch = tok(x_texts, return_tensors="pt", padding=True).to("cuda")
    out = policy.generate(**batch,
                          do_sample=True, temperature=1.0, top_p=1.0,
                          max_new_tokens=192, num_return_sequences=G)
    return out              # [len(x_texts) * G, |x| + |y|]

temperature=1.0 is not a value picked at random — see trap 4 in section 9 for why going lower is dangerous.

7.2 Reverse KL against the teacher — the heart of the whole chapter​

import torch.nn.functional as F

K = 128         # keep only the teacher's top-K — the reason is in the arithmetic box below

def opcd_loss(policy, c_ids, x_ids, y_ids, y_mask):
    # student: sees only x + y  (adapter on)
    s_in = torch.cat([x_ids, y_ids], dim=1)
    s_logits = policy(s_in).logits[:, x_ids.size(1) - 1 : -1]

    # teacher: same base weights, adapter off, and "sees c" — no gradient
    with torch.no_grad(), policy.disable_adapter():
        t_in = torch.cat([c_ids, x_ids, y_ids], dim=1)
        t_logits = policy(t_in).logits[:, c_ids.size(1) + x_ids.size(1) - 1 : -1]

    # cut down to the support of the teacher's top-K mass, then renormalize both sides
    topk = t_logits.topk(K, dim=-1).indices
    t_logp = torch.log_softmax(t_logits.gather(-1, topk).float(), dim=-1)
    s_logp = torch.log_softmax(s_logits.gather(-1, topk).float(), dim=-1)

    # reverse KL: π_θ sits "in front" — each term is weighted by the student, not the teacher
    kl = (s_logp.exp() * (s_logp - t_logp)).sum(-1)          # [B, |y|]
    return (kl * y_mask).sum() / y_mask.sum()                # per-token mean = 1/|y|

7.3 The training loop​

opt = torch.optim.AdamW(
    [p for p in policy.parameters() if p.requires_grad], lr=1e-5)

for epoch in range(2):
    for x_texts in batches(prompts, batch_size=1):
        seqs = rollout(policy, x_texts, G=4)
        c_ids, x_ids, y_ids, y_mask = split_and_pad(seqs, c_len)   # see the notebook
        loss = opcd_loss(policy, c_ids, x_ids, y_ids, y_mask)
        loss.backward()
        opt.step(); opt.zero_grad()

• lr=1e-5 — higher than DPO (5e-6) but far lower than SFT (2e-4): we are bending a distribution toward a teacher sitting nearby, not teaching new knowledge
• 300 prompts × 4 rollouts × 2 epochs takes about 16 minutes on a T4
• during training, the notebook logs the mean entropy of the output and the mean answer length every 20 steps — these two are the canaries in the coal mine for mode collapse; see trap 1 in section 9

8. Results​

The notebook measures 4 things and writes them to results.json:

1. Persona-compliance rate — measured by a deterministic checker (code below) with a Wilson 95% CI, comparing "teacher + full context" against "OPCD student with no context"
2. Prompt tokens per request — measured with the actual tokenizer: should drop by ~400 tokens per request
3. Latency to first token — a 30-token prefill versus a 430-token prefill on the same machine
4. Training-time canaries — the entropy and answer-length curves, attached to every experimental result

The compliance checker uses no LLM to judge an LLM — deterministic rules that return the same verdict on every rerun:

# The notebook defines these tuples with the actual Thai strings;
# English glosses are shown here so the page stays readable.
POLITE  = ("khrap", "kha", "na khrap", "na kha")   # sentence-final politeness particles (and variants)
REFUSAL = ("cannot give advice", "recommend consulting", "specialist",
           "doctor", "lawyer")                     # refusal phrasing, in Thai

def th_ratio(s):
    thai = sum(1 for ch in s if "\u0e01" <= ch <= "\u0e5b")   # the Thai Unicode block, U+0E01..U+0E5B
    letters = sum(1 for ch in s if ch.isalpha())
    return thai / max(letters, 1)

def comply(answer, is_restricted):
    ok_thai   = th_ratio(answer) >= 0.85                       # rule 1: answers in Thai
    ok_polite = any(p in answer for p in POLITE)               # rule 2: khrap/kha
    ok_refuse = (not is_restricted) or any(k in answer for k in REFUSAL)   # rules 3-4
    return ok_thai and ok_polite and ok_refuse

The measured values land in the section 9 table (the ? cells are the ones your notebook fills in — not me).

Below are real answers from before and after training, both responding without seeing the context — the "before" side is the bare base model; the "after" side is the OPCD student. Click through the samples and ask yourself: if nobody told you, could you tell which one never saw the system prompt?

9. Comparison​

Four systems, all measured on the same test set — the first two rows are the floor and the ceiling; the last two are the real bout:

The pattern you should expect: both offline CD and OPCD climbing from the floor toward the ceiling while paying the floor row's prompt bill — and the place the two methods separate is the OOD compliance column: offline CD learned only from the teacher's trajectories, so it tends to slip on kinds of prompts it never saw, while OPCD was inspected on its own trajectories throughout, so it should hold the rules more steadily off the beaten path. If this column cannot be told apart within the CIs — that is a null result, and section 8's rule is in force.

Traps to watch for​

1. Reverse KL + a small student = mode-collapse risk Mode-seeking is a double-edged sword: a low-capacity student may "pick a mode" in the most extreme way — say, answering every question with the same canned refusal, which earns a genuinely low KL and is genuinely useless. This is why section 7.3 logs the output entropy and the answer length as canaries: if entropy dives while answers get shorter and more repetitive, stop, then lower the LR or the epoch count.

2. Top-K truncation bias The top-128 surrogate approximates the true KL only where the teacher's top-128 covers nearly all the mass. Positions where the teacher "hesitates" (high entropy — the start of the first sentence, say) are where coverage drops and the bias creeps in. Don't guess — the notebook prints the mean coverage and the lowest percentile. If it's unusually low, then raise K.

3. Teacher and student tokenizers must match Per-position KL is only defined when both sides split tokens identically — cross model families and the vocabs differ; positions instantly stop lining up. In this chapter the condition holds automatically, because teacher and student are one set of weights — which is what makes this setup especially clean pedagogically: you get to learn the distillation mechanism in full without carrying the tokenizer problem alongside. (In chapter 7, where teacher and student are different models, this problem becomes immediately real.)

4. Temperature too low = the student only rehearses moves it already knows Sample at a low temperature and the student produces only answers it is already confident in, so the teacher only ever inspects states the student already handles well — gradient at the spots where behavior still violates the persona barely ever appears. temperature=1.0 forces the student to carry itself into states where it still fails, and be inspected there too.

10. Summary​

• A settled system prompt is knowledge stored in the wrong place — in the prompt you pay every request; in the weights you pay once
• Context distillation trains a student that doesn't see $c$ to match a teacher that sees $c$ — and in this chapter, teacher and student are one set of weights, differing only in prompt and adapter
• The KL must be reverse ($\pi_\theta$ in front): mode-seeking forces the student not to do what the teacher wouldn't — exactly what persona/safety work needs
• The rollouts must be the student's own: on-policy removes exposure bias by construction, because the training states and the inference states are the same set
• Top-128 is a memory decision you can do arithmetic on — cut 151,936 down to 128, accept that the objective becomes a surrogate, and measure coverage to keep it honest
• Entropy and answer length are the canaries of mode collapse — always log them; don't wait to find out at the wreck
• The fair baseline is offline CD, not the bare model — and if you don't win, report the null result as it stands

Next chapter: Model Distillation — this time we shrink the model, not the prompt. Remember the sentence pinned in section 2: context distillation changes "what the model knows without being told"; model distillation changes "the size of the model" — a big teacher, a small student, and the tokenizer problem this chapter got for free stops being free.

References​

1. Ye et al. (2026). On-Policy Context Distillation for Language Models — OPCD -- the method this chapter implements
2. Askell et al. (2021). A General Language Assistant as a Laboratory for Alignment — the original offline context distillation (section 9's baseline)
3. Snell et al. (2022). Learning by Distilling Context — distilling a context into model behaviour
4. Agarwal et al. (2023). On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes — GKD: the JSD framework unifying forward and reverse KL
5. Gu et al. (2023). MiniLLM: On-Policy Distillation of Large Language Models — MiniLLM: the case for reverse KL

The writing, code and notebooks in this series are licensed under CC BY-NC-SA 4.0 — reuse and adapt them freely with attribution, for non-commercial purposes, and share your adaptations under the same terms. Third-party models and datasets referenced here keep their own licences.