You need to replace an expensive frontier model with a small on-device model for a narrow workflow. How would you design the distillation data, choose between logits and teacher-generated traces, and prove that the smaller model has retained the behaviors that matter?
Define the workflow’s critical behaviors first, then build a stratified dataset of real, rare, adversarial, and policy-sensitive cases. Use logits when teacher and student share an aligned output space; use concise structured traces when they do not, and validate the quantized on-device model against human-labeled behavioral gates rather than teacher agreement alone.
How to think about it
I would define a behavior contract first, then build a stratified dataset from real traffic, rare failures, policy edge cases, and adversarial examples. I would use logits when the teacher and student have compatible token or class spaces, use concise structured traces when they do not, and prove retention with slice-based tests on the final quantized device model.
Why this is the real problem
Distillation means training a smaller student model to imitate a larger teacher model. The trap is treating imitation as the goal. It is not. The goal is to retain the behaviors that make the workflow safe and useful.
A teacher may produce beautiful prose while making the wrong decision. A student may match the teacher’s wording while failing on a single policy exception that costs money. So I would write down the contract before collecting examples:
- What inputs must be handled?
- What actions are allowed?
- When must the model abstain?
- Which errors are merely annoying, and which are dangerous?
- What latency, memory, and battery budget applies on the device?
For a support workflow, “sounds helpful” is not a behavior. “Never issue a refund without a valid order and an eligible return window” is.
The mechanism matters because the available teacher signal contains different kinds of information. A hard label says which answer won. A teacher’s probability distribution says how the alternatives compared. A generated trace can show the evidence and intermediate decision structure, but it can also contain a confident fiction. Each signal teaches something different.
Designing the distillation data
I would combine four sources.
First, representative production data. Sample by the actual traffic distribution, after removing or protecting personal information. If 80 percent of requests are simple password resets, that fact belongs in training. It should not, however, dominate the evaluation set.
Second, risk-weighted examples. Add rare but costly cases: missing fields, contradictory records, policy exceptions, prompt injection, negation, long conversations, typos, and multilingual input if the product supports it.
Third, contrastive pairs. Change one fact and require the decision to change when it should. For example, “delivered 12 days ago” and “delivered 42 days ago” should not receive the same return decision if policy allows returns for 30 days. These pairs teach the boundary instead of merely showing the center of the distribution.
Fourth, trusted labels. Teacher outputs are useful pseudo-labels, but they are not ground truth. I would have domain experts label a smaller, carefully sampled set, especially the high-risk slices. The student should learn the policy, not faithfully reproduce a teacher’s occasional mistake.
I would split by customer, conversation, and time where possible. Otherwise, nearly identical tickets can land in both training and test sets, producing a gratifying but meaningless score. Every record should retain provenance: human label, teacher output, policy version, and whether it came from production, synthesis, or an adversarial test.
Logits versus traces
Logits are the raw scores a model assigns before converting them into probabilities. If teacher and student use the same vocabulary, or the same small set of output classes, I can compare their distributions directly.
For an autoregressive language model, a common objective combines the ordinary hard-target loss with a softened teacher distribution:
L = alpha * L_hard + (1 - alpha) * T^2 * KL(p_teacher^T || p_student^T)
Here, T is a temperature that softens the distributions, and KL is Kullback–Leibler divergence, a measure of how different the student distribution is from the teacher’s. The hard loss protects the known answer. The softened loss preserves “dark knowledge”: the teacher may assign substantial probability to two plausible alternatives and very little to the rest.
I would choose the mixture weight and temperature on a validation set, not by folklore. Teacher probabilities can be badly calibrated, so copying them blindly can make the student confidently wrong.
Logits are attractive when:
- teacher and student have compatible tokenizers and output vocabularies;
- the teacher exposes the required scores;
- preserving uncertainty or rankings matters;
- the student is learning a classification or small structured decision space.
They are awkward when the teacher’s vocabulary differs from the student’s. Token-level comparison then has no clean meaning. Logits are also expensive to store. For 50,000 examples, 128 target tokens, and a 100,000-token vocabulary, storing one score per token would mean 640 billion scores. Even at two bytes per score, that is roughly 1.28 terabytes before metadata.
Teacher-generated traces are input-output examples with an explanation or intermediate structure produced by the teacher. For production, I prefer constrained traces such as evidence fields, rule identifiers, extracted entities, and the final action. They are easier to inspect and work across different tokenizers.
The important caveat: a trace is not necessarily the teacher’s genuine internal reasoning. It is generated text. It may rationalize a decision after the fact, invent evidence, or teach a shortcut. I would not force the on-device model to emit long free-form reasoning unless the workflow truly requires it. That increases latency, output surface area, and opportunities for fabricated explanations.
The usual answer is a hybrid:
| Signal | Best use | Main risk |
|---|---|---|
| Hard labels | Policy decisions and exact fields | Throws away uncertainty |
| Logits | Shared output space and ranking | Vocabulary mismatch and storage cost |
| Structured traces | Evidence extraction and decomposition | Teacher rationales can be wrong |
A concrete design
Suppose the workflow is a mobile returns assistant. The input contains a customer message, order metadata, and the relevant policy excerpt. The student must choose one of six actions and produce a short response.
I might begin with 60,000 historical examples: 42,000 for training, 9,000 for validation, and 9,000 for a locked test set. I would add 2,000 manually authored adversarial cases, kept entirely outside training. The test set would report both the natural traffic score and separate slices for missing order IDs, negation, expired windows, damaged goods, and prompt injection.
A teacher-generated record could look like this:
{
"evidence": [
"delivered 12 days ago",
"item is within the 30-day return window"
],
"policy_rule": "standard_return_window",
"action": "approve_return",
"reply": "Your return is approved. We will send the next steps."
}
The action and evidence fields receive hard supervised loss. If the teacher and student share an aligned output space, I would add logit distillation on the action distribution, where preserving the distinction between “approve,” “ask for more information,” and “拒否” is useful. If they do not share that space, I would train on the structured fields and final action instead.
I would audit a risk-weighted subset with policy experts. A student that gets 99 percent of easy tickets right but mishandles “item never arrived” is not ready to ship.
Proving that behavior survived
I would evaluate the final artifact, not just the floating-point checkpoint. That means the exact tokenizer, quantization scheme, runtime, and device hardware used in production.
The test plan would include:
- Outcome metrics: action accuracy, field-level exact match, valid JSON rate, and response quality judged against a rubric.
- Safety metrics: false approvals, false refusals, unsafe tool calls, and correct abstention when required information is missing.
- Slice metrics: every important case above, with minimum thresholds rather than only one overall average.
- Metamorphic tests: change one relevant fact and verify the expected behavior changes, or change irrelevant wording and verify the decision remains stable.
- Operational metrics: memory footprint, cold-start time, median and tail latency, battery impact, and crash rate on real hardware.
- Human review: inspect every disagreement on high-risk slices and a random sample of apparent agreements.
I would compare against human labels and the existing system, not only against the teacher. Teacher agreement proves imitation. It does not prove correctness.
For a safety-sensitive action, “zero failures” needs humility. If 3,000 representative cases produce no observed failures, a rough 95 percent upper bound under simple independent sampling is about 0.1 percent, not zero. Rare risks need targeted adversarial testing, monitoring, and usually an abstain-or-fallback path.
The senior-level trade-off
Distillation is not automatically the right tool. If the workflow has six actions and strict rules, a small classifier plus deterministic policy code may be more reliable, cheaper, and easier to audit than a generative student. If the policy changes weekly, distilling a large model may create stale behavior; retrieval or explicit rules may age better.
The failure mode I would watch for first is fluent regression: offline accuracy stays high, but users report that the model approves returns for “I did not receive the item.” That usually means common examples overwhelmed negation and boundary cases. I would fix it with contrastive data, a dedicated risk slice, and a hard gate on the dangerous action—not with another pile of generic synthetic conversations.
What they’ll ask next
What if the teacher does not expose logits?
Use curated teacher outputs, structured labels, and human-reviewed edge cases. If uncertainty matters, estimate it from repeated samples or the student’s calibrated confidence, but do not pretend text outputs are equivalent to true teacher logits.
What if the tokenizers differ?
Avoid token-level KL. Distill at the level of actions, extracted fields, spans, structured records, or sequence-level outputs. A shared tokenizer can make logits practical, but it is not worth redesigning the whole system merely to obtain them.
How would you roll it out?
Run the student in shadow mode first, compare decisions and latency, then canary by device cohort. Keep abstention and fallback behavior explicit, and monitor the same risk slices used before launch.
One line to say in the room
“I would distill the behaviors in the contract, not the teacher’s prose, and I would sign off only after the quantized device model passes human-labeled, adversarial, and risk-sliced regression gates.”