Knowledge Distillation: Logits, Sequences, and Rationales

Knowledge distillation (KD) is often mischaracterized as simply "training a small model on a large model’s outputs." While this captures the high-level goal, it obscures the mechanistic reality: KD is a structured transfer of information density. A large teacher model encodes knowledge in three distinct layers: its confidence distributions (logits), its generative sequence patterns, and its internal reasoning paths. To optimize a student model effectively, you must choose which layer to target.

This article is Part 2 of the Small Models & Distillation series.

For ML engineers, the choice between logit matching, sequence distillation, and rationale distillation determines whether you are compressing a classifier, a language model, or a reasoning engine.

Logit Matching: Transferring Class Relationships

The foundational mechanism of knowledge distillation, introduced by Hinton et al. (2015), operates on the premise that hard labels (e.g., one-hot vectors) discard valuable information. When a teacher classifies an image of a dog, it might assign 0.6 probability to "dog," 0.3 to "wolf," and 0.1 to "cat." A hard label would only preserve the "dog" signal, treating "wolf" and "cat" as equally wrong.

Logit matching preserves these inter-class relationships by using soft targets. The teacher’s final layer logits $z$ are passed through a softmax function with a temperature parameter $T$:

$$q_i = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}$$

A higher $T$ smooths the distribution, revealing the teacher’s uncertainty and relative similarities between classes. The student network is then trained to minimize the Kullback-Leibler (KL) divergence between its own softened output $p$ and the teacher’s $q$:

$$L_{KD} = T^2 \cdot \text{KL}(q || p)$$

The $T^2$ scaling factor ensures gradient magnitudes remain stable as $T$ changes. This mechanism teaches the student that "wolf" and "cat" are structurally similar errors, improving generalization even when the student has fewer parameters.

Sequence Distillation: Next-Token Probability Transfer

In natural language processing, logit matching alone is insufficient because the output space is vast and sequential. Sequence-level knowledge distillation transfers the teacher’s next-token prediction distribution at every time step.

Consider a machine translation task where the teacher model $T$ generates a target sequence $y$ from a source $x$. The teacher computes the conditional probability:

$$P_T(y_t | x, y_{<t})$$

The student model $S$ is trained to match this distribution for every token $t$ in the sequence. Unlike logit matching, which operates on static classes, sequence distillation operates on dynamic, context-dependent distributions. This preserves the teacher’s fluency and grammatical structure.

Jiao et al. (2020) demonstrated this with TinyBERT, where the student was distilled from a larger BERT teacher using both input-output pairs and the teacher’s internal representations. The key mechanism here is that the student learns not just what to say, but how likely each continuation is given the preceding context. This is critical for tasks where rare tokens or specific phrasings carry high semantic weight.

Rationale Distillation: Transferring Reasoning Steps

The most recent evolution in KD is rationale distillation, which targets the teacher’s internal reasoning process rather than just its outputs. This is particularly relevant for complex reasoning tasks like mathematics or logical deduction.

In rationale distillation, the teacher does not just produce an answer $y$. It produces a chain-of-thought (CoT) rationale $r$:

$$y, r = T(x)$$

The student is then trained on the joint distribution:

$$P_S(y, r | x) = P_S(y | x, r) \cdot P_S(r | x)$$

This dual objective forces the student to learn two things: (1) how to generate the correct reasoning steps given the input, and (2) how to derive the final answer from those steps. Hsieh et al. (2023) demonstrated in "Distilling Step-by-Step" that this approach significantly outperforms standard KD in mathematical reasoning tasks because the student internalizes the process of solving, not just the mapping from input to output [1].

Wei et al. (2022) emphasized that chain-of-thought prompting elicits reasoning in large language models during inference [2]. While their primary contribution focused on inference-time prompting, subsequent works by Luo et al. (2023) have extended these principles to finetuning smaller models on rationales generated by larger models, enabling better zero-shot generalization to new reasoning tasks. The mechanism here is that the rationale acts as a latent variable that structures the student’s understanding of the problem space.

[1] Hsieh, C., et al. "Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes." arXiv preprint arXiv:2305.02301 (2023). [2] Wei, J., et al. "Chain-of-Thought Prompting Elicits Reasoning in Large Language Models." Advances in Neural Information Processing Systems 35 (2022).

Choosing the Right Distillation Strategy

Method	Mechanism	Best For
Logit Matching	KL divergence on softened class probabilities	Image classification, tabular data, static outputs
Sequence Distillation	Next-token probability matching	Translation, summarization, fluent text generation
Rationale Distillation	Joint prediction of CoT and final answer	Math, logic, multi-step reasoning, QA

Logit matching is computationally cheap and effective for static classification tasks. Sequence distillation is essential for preserving linguistic nuance in LLMs. Rationale distillation is necessary when the task requires explicit reasoning steps.

Conclusion

Knowledge distillation is not a monolithic technique. It is a spectrum of information transfer mechanisms. Logit matching transfers structural knowledge about class relationships. Sequence distillation transfers probabilistic fluency. Rationale distillation transfers procedural reasoning. Understanding these mechanisms allows ML engineers to select the appropriate KD strategy based on the student’s architecture and the task’s complexity.

As models grow larger, the cost of inference becomes prohibitive. Distillation offers a path to deploy capable models on edge devices or within cost-constrained APIs. However, the success of this deployment depends on matching the distillation mechanism to the nature of the knowledge being transferred.

Common Pitfalls

1. Incorrect Temperature Selection: Using a temperature $T$ that is too high can wash out useful signal, while $T$ too low provides no benefit over hard labels. Always validate $T$ on a held-out validation set to find the optimal smoothing level.
2. Overfitting on Soft Targets: If the student model is too large or trained for too many epochs on the teacher’s soft targets, it may memorize the teacher’s specific errors or idiosyncrasies rather than learning generalizable patterns. Use early stopping or regularization.
3. Mismatched Task Types: Applying logit matching to sequential tasks where context matters, or rationale distillation to simple classification tasks, leads to inefficient training and poor performance. Align the distillation method with the structural complexity of the target task.

Practical Takeaways

• Match Distillation Type to Task Complexity: Use logit matching for simple classification, sequence distillation for text generation, and rationale distillation for complex reasoning tasks.
• Use Temperature Scaling for Smoother Distributions: Tune the temperature parameter $T$ to balance between preserving hard label signals and capturing inter-class relationships.
• Validate Student Performance Independently: Do not assume that minimizing KD loss directly correlates with improved downstream task performance. Always evaluate the student on the actual target metric.

FAQ

What is the ideal temperature for knowledge distillation? There is no single ideal temperature; it depends on the dataset and model size. A common starting point is $T=2$ to $T=5$, but you should experiment with values that maximize the validation performance of the student model.

Can knowledge distillation be used for non-LLM models? Yes. Logit matching is widely used in computer vision (e.g., distilling from ResNet to MobileNet) and tabular data models. It is not limited to natural language processing.

Does rationale distillation require the student to generate the chain-of-thought? Not necessarily. In some implementations, the student only predicts the final answer, but the loss function is influenced by the teacher's rationale. However, training the student to generate the rationale itself often leads to better interpretability and robustness.