🤖 machine learning

One Student, Many Teachers: Multi-Task On-Policy Distillation via Soft-Prompt Privileged Context

This paper introduces Multi-Task On-Policy Distillation via Soft-Prompt Privileged Context, a method that employs frozen backbone models with task-specific learnable soft prompts as teachers to enable a single student model to efficiently absorb knowledge from multiple tasks in parallel without suffering from catastrophic forgetting or requiring full weight fine-tuning.

Original authors: Yingzi Ma, Zichen Zhu, Ming Jiang, Chaowei Xiao

Published 2026-07-22

📖 5 min read🧠 Deep dive

CC0 1.0

Original authors: Yingzi Ma, Zichen Zhu, Ming Jiang, Chaowei Xiao

Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a brilliant but slightly clumsy student how to solve complex puzzles. In the world of artificial intelligence, this student is a "Large Language Model" (LLM)—a computer program that has read almost everything on the internet and can write stories, answer questions, and solve math problems. But sometimes, even these super-smart models get stuck or make silly mistakes. To fix them, scientists use a technique called "distillation." Think of this as a master chef (the teacher) guiding a junior cook (the student) through a recipe. The student doesn't just copy the final dish; they watch every chop, stir, and taste, learning the process of cooking, not just the result.

For a long time, the best way to teach these AI students was to have the teacher write out the perfect answer and show it to the student. But this had a hidden trap: the teacher would sometimes just guess the answer first and then make up a story to justify it, rather than actually figuring it out step-by-step. The student would learn this bad habit, too. Another method involved changing the teacher's entire brain (its internal weights) to make it smarter, but this often caused the teacher to forget how to do other things, like writing poetry or chatting, while focusing on the new task. The big question scientists have been asking is: How do we create a teacher that is smart enough to show the right path, but stable enough not to forget its old skills or cheat by looking at the answer key?

This paper introduces a clever new solution called PROMPTSD (Prompt-based On-Policy Soft-Distillation). The authors, a team from universities in the US and Singapore, propose a way to build a teacher that is almost identical to the student, except for a tiny, invisible "magic note" attached to the front of its instructions.

Here is how it works, using a simple analogy: Imagine the student and the teacher are twins who share the exact same brain and personality. Usually, to teach the student a new skill, you might give the teacher a whole new set of glasses (changing its brain) or a cheat sheet with the answers written on it (giving it the answer key). Both of these methods have problems: the new glasses make the teacher forget how to see other things, and the cheat sheet makes the teacher lazy, justifying the answer instead of solving the problem.

PROMPTSD does something different. It gives the teacher a soft prompt. Think of this as a tiny, invisible sticky note that says, "Remember, for this specific math problem, think about fractions first." This note is made of numbers, not words, and it's the only thing that changes between the teacher and the student. The teacher's brain remains frozen and exactly the same as the student's. Because the teacher isn't looking at the answer key, it has to actually think through the problem to generate the right path. Because its brain hasn't been rewritten, it doesn't forget how to do other things.

The researchers tested this idea on a model called Qwen3-1.7B-Base and another called Phi-4-mini-instruct. They taught these models four different skills: Science, using tools (like a calculator or search engine), Biology, and Math. They found that their "magic note" teacher was incredibly effective. In single tasks, the student learned just as well as if the teacher had been fully retrained, but the student learned much faster because it only had to adjust a tiny fraction of its parameters (about 0.05% of the total brain, compared to 100% for full retraining).

Even more impressively, they tried teaching the student four different skills at the same time. Usually, when you try to learn too many things at once, you end up forgetting the old ones (a problem called "catastrophic forgetting"). But because the PROMPTSD teacher only used these tiny, task-specific notes, the student could absorb knowledge from all four teachers simultaneously without getting confused or forgetting its general smarts. In fact, on the Qwen3-1.7B-Base model, this multi-task approach achieved an average score of 56.2, beating all other methods.

The paper explicitly rules out a few common ideas. It shows that simply giving the teacher the answer (like a "Gold Answer" teacher) actually hurts the student because the teacher just rationalizes the answer instead of deriving it. It also suggests that making the teacher much bigger (like using a 4B or 8B parameter model) doesn't always help; sometimes a teacher that is the same size as the student but uses the right "thinking pattern" is better. Furthermore, they found that Reinforcement Learning (a method where the teacher learns by trial and error) is good at sharpening skills the student already has, but it's not great at teaching brand-new knowledge from scratch.

The authors are quite sure about these results because they ran the experiments multiple times and checked the numbers carefully. They measured exactly how much the student's thinking overlapped with the teacher's and found that their method created the perfect amount of overlap: enough to learn, but not so much that there was nothing new to discover. They also proved that their method avoids the "SFT trap," where retraining a model makes it forget its general abilities.

In the end, PROMPTSD suggests that the secret to teaching AI isn't necessarily making the teacher bigger or giving it the answers. Instead, it's about giving the teacher a tiny, flexible nudge—a soft prompt—that guides its thinking without changing who it is. This allows a single student to learn from many different teachers at once, becoming a master of many skills without losing its own identity. It's a bit like having a personal tutor who whispers the right strategy into your ear for every subject, without ever rewriting your brain or letting you peek at the test answers.

Technical Summary: PROMPTSD

Problem Statement

On-policy self-distillation (OPSD) aims to teach large language models (LLMs) new skills by using a teacher model that shares the student's backbone and supervises its own rollouts. However, existing teacher construction methods suffer from significant structural limitations:

Observation-based teachers (e.g., those injecting gold answers or demonstrations as input context) induce post-hoc rationalization. By conditioning on the gold answer ( $y_{gold}$ ), the teacher justifies a predetermined outcome rather than deriving a reasoning trace, a failure mode the student inherits during inference when the answer is unavailable.
Weight-modifying teachers (e.g., Full Fine-Tuning (FFT) or LoRA) accumulate distributional drift and "intruder dimensions." While they can inject new knowledge, they perturb the model's representational geometry, leading to catastrophic forgetting of general capabilities when applied sequentially across tasks.
RL-tuned teachers often fail to inject new capabilities, instead amplifying existing latent behaviors, and require task-specific verifiers that are not universally available.

The core challenge is identifying a teacher construction that satisfies three desiderata simultaneously: (D1) Thinking-pattern consistency (sharing the student's weights to ensure the learning signal is absorbable), (D2) Capability gap absorption (encoding new task knowledge), and (D3) No post-hoc rationalization (not conditioning on the gold answer).

Methodology: PROMPTSD

The authors propose PROMPTSD (Prompt-based On-policy Distillation), a framework that constructs teachers via soft-prompt tuning rather than weight modification or input-side context injection.

1. Soft-Prompt Teacher Construction

For each task $k$ , a set of learnable continuous embeddings $P_k$ (a soft prompt) is attached to the input embedding layer of a frozen student backbone $\pi_\theta$ .

Training: The prompt $P_k$ is optimized via standard supervised fine-tuning (SFT) on $(x, y_{gold})$ pairs to minimize next-token cross-entropy.
Result: This yields a task-specific teacher $\pi_{\theta+P_k}$ that shares all transformer weights with the student but differs only in the input-side conditioning. The prompt acts as a parametric form of observation, encoding task knowledge without altering the model's internal geometry ( $\Delta W = 0$ ).

2. On-Policy Distillation

The student $\pi_\theta$ is updated via on-policy reverse-KL divergence minimization:

The student generates rollouts $Y_s$ given prompt $x$ .
The teacher $\pi_{\theta+P_k}$ (with the frozen prompt $P_k$ ) provides per-token supervision distributions on the same states visited by the student.
Crucially, the teacher does not see the gold answer during distillation; it only sees the soft prompt $P_k$ and the input $x$ . This eliminates post-hoc rationalization.

3. Parallel Multi-Task Extension

PROMPTSD extends to multi-task settings by merging datasets from $K$ tasks into a single corpus.

Routing: Each example in the merged corpus retains a task tag $k$ . During training, the example is routed to its corresponding soft-prompt teacher $\pi_{\theta+P_k}$ .
Parallelism: All $K$ soft prompts are trained and available simultaneously. The student learns from all teachers in parallel, avoiding the sequential forgetting associated with continual learning.
Inference: All soft prompts are discarded. The deployed model is the base backbone $\pi_\theta$ alone, incurring no parameter or latency overhead.

Key Contributions

Novel Teacher Architecture: The paper introduces the first on-policy distillation teacher constructed via prompt tuning. This design uniquely satisfies the three desiderata for effective OPD: it maintains thinking-pattern consistency (shared weights), enables capability absorption (learnable prompt), and prevents post-hoc rationalization (no gold answer in context).
Parallel Multi-Task Distillation: The framework demonstrates how to route a single student to multiple task-specific soft-prompt teachers simultaneously, achieving superior multi-task performance without catastrophic forgetting.
Empirical Validation of Principles: Through extensive experiments, the authors establish two design principles for OPD:
- The teacher family must match the task type (SFT-based teachers for capability injection; RL-based for elicitation).
- Teacher scale is less critical than whether the teacher carries new knowledge.

Experimental Results

Experiments were conducted on Qwen3-1.7B-Base and Phi-4-mini-instruct across four tasks: Science, Tool Use, Biology, and Math.

Single-Task Performance: The single-task variant (OPD with a PT teacher) matches or exceeds full fine-tuning (FFT) and LoRA baselines while training orders of magnitude fewer parameters. Notably, on the Biology task, the student achieved 62.0% accuracy, surpassing its own teacher (58.0%) and significantly outperforming LoRA-based students.
Multi-Task Performance: The multi-task PROMPTSD achieved the best overall average accuracy (56.2% on Qwen3-1.7B-Base), outperforming the strongest single-task baseline by +2.3 points and cross-size distillation baselines by +7.7 points.
General Capability Retention: Unlike sequential SFT, which degrades general benchmarks (MMLU-Pro, HellaSwag, TruthfulQA) due to catastrophic forgetting, PROMPTSD preserved general capabilities while improving task-specific accuracy.
Ablation Studies:
- Prompt Length: Longer prompts (e.g., $L=400$ ) generally improved performance on complex tasks like Math and Tool Use, though diminishing returns were observed.
- Divergence: Reverse-KL (mode-seeking) was found to be superior for Math reasoning, while other divergences performed similarly on knowledge-heavy tasks.
- Teacher Contribution: Leave-one-teacher-out analysis revealed that for Math, the gain came primarily from cross-task exposure (elicitation) rather than the specific Math teacher, whereas for Science and Biology, the dedicated teachers were essential for capability injection.

Significance and Claims

The paper claims that PROMPTSD resolves the trade-offs inherent in existing teacher designs. By using a soft prompt as a "parametric privileged context," it allows the student to absorb new task knowledge without the off-manifold drift caused by weight updates (FFT/LoRA) or the reasoning shortcuts caused by gold-answer conditioning (Observation-based).

The authors position PROMPTSD as a solution to the "SFT trap" in continual learning, demonstrating that a single student can absorb knowledge from multiple specialized teachers in parallel while retaining its general capabilities. The work suggests that for capability injection tasks, a prompt-tuned teacher is a more effective distillation source than a full fine-tuned teacher, as it avoids the catastrophic forgetting associated with modifying the base model's weights.

The paper remains modest regarding its scope, acknowledging that experiments were limited to models up to 1.7B parameters and did not cover open-ended generation, multilingual, or agentic settings. It notes that the operational definitions of the three desiderata are empirically validated rather than formally proven.