Must the teacher and student use the same tokenizer for knowledge distillation?

For per-token logit alignment, the safest approach is to share the tokenizer and vocabulary. If vocabularies differ, you need additional token alignment, common space mapping, or sequence-level distillation; you cannot directly compare logits at the same position between the two models.

Is caching only the Top-K Teacher Logprobs sufficient?

It depends on the divergence direction, the K value, and how tail probabilities are handled. When caching Top-K, you should also save the residual probability mass of the unselected tokens and include it as a Tail Bucket in normalization and loss computation; otherwise, the teacher distribution will be systematically distorted.

Why might a student model still degrade after offline distillation?

A common reason is that training only covers teacher-generated trajectories, while online inference uses the student's own generations, causing a distribution shift. You should add On-policy replay, allowing the student to receive teacher feedback on its own erroneous trajectories, and separately validate quality, calibration, length, and cost metrics.

How should I choose between Forward KL, Reverse KL, and JSD?

There is no single best answer for all tasks. Forward KL emphasizes covering the teacher distribution, Reverse KL focuses on the student's high-probability regions, and JSD offers a compromise. Compare quality, calibration, output diversity, and stability on the same calibration set, rather than just looking at training loss.

LLM Knowledge Distillation in Production: Controlling Student Model Degradation with Teacher Logprob Contracts, On-Policy Replay, and Top-K Tail Buckets

Executive Summary

Knowledge Distillation (KD) for large language models is often oversimplified as “generate answers with the big model, then SFT the small model.” However, when pushing distillation into production, the truly tricky issues are rarely in the training script itself. They lie in: how to version the teacher’s probability distribution, how to fix temperature and divergence, how to preserve tail probabilities after Top-K truncation, how to replay the student’s self-generated trajectories, and when the student model is ready to replace the teacher.

This article takes an engineering deployment perspective, systematically covering five core aspects of knowledge distillation—Teacher Logprob contracts, Top-K Tail Buckets, On-policy replay, Tokenizer alignment, and deployment gating—to help teams avoid silent degradation after distillation.

Background: Distillation is Not Just Another SFT

The goal of knowledge distillation is to have a smaller Student model learn the behavior distribution of a larger Teacher model, thereby reducing inference memory, latency, and per-request cost. The classic approach combines a hard-label loss with a divergence term between the teacher’s and student’s soft distributions:

L = (1 - α) · L_CE(student, label) + α · T² · KL(p_teacher^T || p_student^T)

Where T is the temperature and α controls the weight between the hard label and the distillation signal.

The formula itself is not complex. The production problems lie outside the formula:

Can old Logprobs be reused after the teacher model, tokenizer, chat template, or generation parameters change?
Do Top-K Logprobs, after being transmitted over a network or through object storage, still represent the original distribution?
If the student only imitates the teacher’s generated text, will it expose unseen erroneous trajectories when generating on its own?
The student’s overall score may be close to the teacher’s, but does it degrade on long texts, low-frequency domains, format adherence, or probability calibration?
Can the costs of teacher calls, cache construction, and repeated training be tracked and reproduced?

Therefore, a distillation system must first manage probability contracts and data lineage, and only then the training script.

Core Principles

1. Teacher Logprobs Must Be Versioned Artifacts

For each training record, you should not just save the teacher’s output text. It is recommended to record at least the following structured information:

{
  "record_id": "sha256(prompt_canonical_bytes)",
  "teacher_model_revision": "model-commit-or-checkpoint-hash",
  "tokenizer_revision": "tokenizer-artifact-hash",
  "chat_template_hash": "sha256(template)",
  "generation_config_hash": "sha256(config)",
  "temperature": 1.0,
  "loss_support": "teacher_top_k_plus_tail",
  "top_k": 32,
  "token_ids": [123, 456],
  "teacher_topk_token_ids": [[123, 789], [456, 321]],
  "teacher_topk_logprobs": [[-0.1, -2.3], [-0.2, -1.7]],
  "tail_logprob_mass": [-0.08, -0.12]
}

Key Principle: Whenever the model revision, tokenizer, template, temperature, or Top-K rules change, this batch of Logprobs should be treated as a new artifact and should not be silently overwritten. This is the data lineage foundation of the entire distillation system.

2. Top-K Truncation Must Preserve a Tail Bucket

Transmitting and storing full-vocabulary logits is expensive, so engineering often saves only the Top-K. However, if you simply discard the rest and re-normalize, the teacher distribution becomes artificially sharp.

A safer approach is to merge the probability mass of the tokens not in the Top-K into a Tail Bucket:

p_tail = 1 - Σ p_top_k

When computing the divergence, align the sum of probabilities outside the student’s Top-K with p_tail. This reduces the transmission load without pretending the tail probabilities don’t exist.

Divergence Type	Characteristics	Sensitivity to Tail Probabilities
Forward KL	Covers the full support of the teacher (mean-seeking)	High—missing the tail causes bias
Reverse KL	Focuses on the student’s high-probability regions (mode-seeking)	Medium—tail has a relatively smaller impact
Generalized JSD	A compromise between Forward and Reverse	Medium—adjustable via the β parameter

Note that Forward KL, Reverse KL, and JSD have different requirements for the support set. When using an external Teacher Server, include loss_top_k, the divergence direction, and the Tail Bucket rules in the training configuration fingerprint, rather than relying on framework defaults.

3. On-Policy Replay Solves the Training-Inference Distribution Shift

Pure offline distillation typically has the student learn sequences generated by the teacher. Once online, the student generates based on its own history of tokens. If an early deviation occurs, subsequent states can fall into regions not covered by the training set. This is the classic Distribution Shift problem.

On-Policy Distillation works as follows:

Have the student generate a portion of the answers first;
Send the states actually visited by the student to the teacher;
Obtain the teacher’s probability feedback on these states;
Use the teacher’s distribution to correct the student’s own erroneous trajectories.

This is very different from just doing SFT on the teacher’s answers: the former teaches the student “how to recover after making a mistake,” while the latter mainly teaches “what the teacher would write on an ideal trajectory.”

4. Inconsistent Tokenizers Prevent Direct Per-Position Logit Alignment

If the Teacher and Student use different tokenizers, the same string can produce token sequences of different lengths and with different boundaries. In this case, teacher_logits[t] and student_logits[t] have no natural one-to-one correspondence.

Production solutions, ordered from lowest to highest complexity, are typically:

Solution	Complexity	Use Case
Shared Tokenizer and Vocabulary	Low	Most stable, suitable for direct Logit KD
Sequence-Level Distillation	Medium	Treats teacher text as training target, no per-token distribution alignment
Common Vocabulary or Token Mapping	High	Compute loss only on reliably aligned regions
Cross-Tokenizer Representation Alignment	Very High	Requires additional models or alignment algorithms, must be validated separately

⚠️ Do not attempt “quick compatibility” by trimming vocabulary size or hard-pairing by position. Such errors often won’t raise an error, but they render the loss meaningless.

Engineering Implementation

1. Build a Three-Stage Distillation Pipeline

It is recommended to split the process into three independently re-runnable stages:

Trajectory Builder: Generates teacher trajectories, student trajectories, or mixed trajectories;
Teacher Scorer: Outputs versioned Logprob artifacts;
Student Trainer: Reads fixed artifacts for training, without a direct dependency on the online teacher.

The advantage of this approach is that a teacher service failure does not directly interrupt student training, Logprobs can be audited and reused, and different divergence or temperature configurations can be compared independently.

If adopting fully On-policy training, you can keep an external Teacher Server, but you should still write requests and responses to a replayable log and ensure the teacher model revision remains constant within a single experiment.

2. Define Immutable Experiment Fingerprints

Each experiment should fix at least the following configuration:

teacher:
  model_revision: teacher-v3.2
  tokenizer_hash: sha256:...
  chat_template_hash: sha256:...
  server_image_digest: sha256:...

student:
  base_revision: student-v1.8
  tokenizer_hash: sha256:...

distillation:
  temperature: 1.0
  kd_ratio: 0.7
  divergence: generalized_jsd
  beta: 0.5
  on_policy_ratio: 0.6
  top_k: 32
  tail_bucket: true

Training results must be traceable back to the teacher artifact, dataset snapshot, loss configuration, and code version.

3. Implement Contract Tests for the Teacher Server

The Teacher Server needs more than just an availability check. You should also verify the following contracts:

The returned Token IDs correspond to the declared Tokenizer;
Logprobs are natural logarithms;
Top-K is sorted by probability in descending order;
The sum of Top-K probabilities and the Tail Mass is close to 1;
Whether Temperature is applied server-side or training-side;
Whether EOS, special tokens, and padding participate in the loss;
Whether the same request, under a fixed configuration, meets the allowed numerical tolerance.

4. Design Caching Strategy Around “Probability Artifacts,” Not “Text Results”

The Logprob Cache Key should include the following dimensions:

prompt_hash + teacher_model_revision + tokenizer_revision
+ chat_template_hash + generation_config_hash + scoring_position_hash

After a cache hit, you must also validate the Schema Version. Never use only the Prompt text as the Key, otherwise, after a teacher upgrade, the old distribution will continue to be read, making training results irreproducible.

5. Capability Gating Must Compare Student, Teacher, and Original Baseline

Maintain at least three comparison groups:

Original Student Base
Distilled Student
Teacher

The release gate is recommended to cover the following dimensions:

Dimension	Check Item
Domain Capability	Domain task accuracy or business success rate
Behavioral Compliance	Instruction following, format constraints, and refusal behavior
Length Robustness	Performance bucketed by length
Long-Tail Coverage	Low-frequency domains and hard examples
Probability Calibration	Probability calibration and confidence
Performance Metrics	P50/P95 latency, memory, and throughput
Cost	Per-request and per-unit effective output cost

The goal of distillation is not for the student to match the teacher on all metrics, but to reach an acceptable capability boundary under clear cost constraints.

Applicable Scenarios

When Knowledge Distillation is Suitable:

The large model’s performance is acceptable, but online latency or GPU costs are too high;
The business is concentrated in a limited domain, and the student doesn’t need to retain the teacher’s full general capabilities;
There is a stable teacher model and a sufficiently large, high-quality prompt distribution;
The model needs to be deployed on smaller GPUs, edge devices, or high-concurrency services;
You can continuously collect real student failure trajectories for iterative distillation.

When it is Not Suitable:

The teacher itself is unstable, and the data distribution changes frequently and drastically;
The Teacher and Student tokenizers cannot be reliably aligned;
The business requires the student to fully inherit the teacher’s long-tail general capabilities.

Common Misconceptions

Misconception 1: Teacher-Generated Text Equals Knowledge Distillation

Training only on teacher-generated text is closer to Sequence-level KD or SFT. It loses the teacher’s probabilistic structure for sub-optimal tokens and cannot directly express the teacher’s uncertainty.

Misconception 2: Higher Temperature Always Means Richer Teacher Information

Increasing the temperature flattens the distribution, but too high a temperature amplifies tail noise. The temperature must be selected using a calibration set and evaluated together with the divergence direction and Top-K size.

Misconception 3: Smaller Top-K Means Cheaper Training with No Quality Impact

A Top-K that is too small loses the teacher’s ranking information among similar candidates. Even if you only keep a small number of tokens, you should still save the Tail Bucket and perform ablation studies on the K value.

Misconception 4: Passing the Overall Average Score is Enough for Launch

Distillation degradation often concentrates on long texts, rare formats, a few domains, and confidence calibration. The average score can mask significant regressions in critical business subsets.

Misconception 5: Cached Data Can Still Be Used After a Teacher Upgrade

Any change in teacher revision, tokenizer, template, or inference parameters can alter the probability distribution. Old caches should only be retained as artifacts from previous experiments and should not be mixed in by default.

Launch Checklist

Artifacts and Contracts

Teacher, Student, Tokenizer, Template, and code versions all have immutable fingerprints
Logprob Schema, logarithm base, temperature location, and Top-K rules are fixed
Probability mass outside Top-K is preserved via a Tail Bucket
Teacher Server contract tests are integrated into CI

Training and Data

The On-policy to Off-policy ratio is clearly configured
Student trajectories are replayable, and failed samples have stable Record IDs
Loss Mask for Padding, Prompt, Assistant, and special tokens has been validated
Training checkpoints can restore trajectory cursors and Logprob artifact versions

Quality and Release

Regression testing by domain, length, and difficulty has passed
Differences in probability calibration between student and teacher have been evaluated
Latency, Throughput, memory, and cost meet compression targets
A fallback path to the teacher is maintained during the canary phase
Online failure trajectories can be fed back into the next round of On-policy distillation

FAQ

Should I choose an offline Logit Cache or an online Teacher Server?

When the teacher is stable, the dataset is fixed, and you want high training reproducibility, prioritize an offline cache. When you need fully On-policy training and student trajectories are constantly changing, you can use an external Teacher Server, but you must fix the teacher revision and log all requests, responses, and probability contracts.

Is a standard SFT Loss still needed after distillation?

Usually, yes. Hard labels or high-quality reference answers provide the task objective, while the teacher’s soft distribution provides relative preferences. The kd_ratio should be determined experimentally; pure KD may inherit teacher biases, while pure SFT cannot fully leverage the teacher’s probabilistic information.