Knowledge Distillation: Foundations and a Reproducible Iris Example

Knowledge distillation transfers part of a larger model’s behavior to a smaller model. This article explains the basic workflow, shows the loss on a reproducible Iris example, and then covers the situations where the method helps or falls short.

Introduction

Knowledge distillation trains a smaller student model with signals from a larger teacher model. The student still needs inputs and training examples. Teacher outputs add supervision, often as a probability distribution over classes, alongside the ordinary labels. A distilled student may use less memory or compute, but its accuracy depends on the task, architectures, data, and training recipe.

The concept of knowledge distillation was introduced by Geoffrey Hinton and colleagues in 2015. It addressed the problem of using a large model’s predictions to train a smaller model when deployment resources are constrained. The resulting accuracy tradeoff depends on the teacher, student, data, and objective.

The Process of Knowledge Distillation

Knowledge distillation commonly follows two stages:

1. Training or selecting the teacher: A large neural network is trained or chosen for the task. It provides predictions on the same inputs used to train the student.

2. Training the student: The student receives the inputs, hard labels when available, and the teacher’s soft targets. Its objective combines ordinary supervised loss with a term that brings its temperature-scaled distribution closer to the teacher’s.

This is often called “dark knowledge”: the relative probabilities in the teacher’s output can expose similarities among classes that a one-hot label does not encode. Whether that extra signal helps depends on the task and the training recipe.

The following compact example adapts the Iris demonstration from the accompanying Iris notebook. It makes the split, seed, teacher inference mode, KL direction, and temperature scaling explicit.

A reproducible Iris example

This example uses the same train split for all models. The validation split is available for model selection, and the test split is touched only for the final report. The teacher is put in evaluation mode and its probabilities are computed without gradients.

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

SEED = 42
torch.manual_seed(SEED)
np.random.seed(SEED)

features, labels = load_iris(return_X_y=True)
x_train, x_test, y_train, y_test = train_test_split(
    features, labels, test_size=0.20, random_state=SEED, stratify=labels
)
x_train, x_valid, y_train, y_valid = train_test_split(
    x_train, y_train, test_size=0.25, random_state=SEED, stratify=y_train
)

scaler = StandardScaler().fit(x_train)
to_tensor = lambda values: torch.tensor(values, dtype=torch.float32)
train_x = to_tensor(scaler.transform(x_train))
valid_x = to_tensor(scaler.transform(x_valid))
test_x = to_tensor(scaler.transform(x_test))
train_y = torch.tensor(y_train, dtype=torch.long)
valid_y = torch.tensor(y_valid, dtype=torch.long)
test_y = torch.tensor(y_test, dtype=torch.long)


class MLP(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(4, hidden_size), nn.ReLU(), nn.Linear(hidden_size, 3)
        )

    def forward(self, x):
        return self.layers(x)


def accuracy(model, x, y):
    model.eval()
    with torch.no_grad():
        return (model(x).argmax(dim=-1) == y).float().mean().item()


def fit_hard(model, epochs=200):
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    for _ in range(epochs):
        model.train()
        optimizer.zero_grad()
        loss = F.cross_entropy(model(train_x), train_y)
        loss.backward()
        optimizer.step()


def fit_distilled(student, teacher, epochs=200, temperature=7.0, alpha=0.3):
    teacher.eval()
    optimizer = torch.optim.Adam(student.parameters(), lr=0.01)
    for _ in range(epochs):
        student.train()
        optimizer.zero_grad()
        student_logits = student(train_x)
        with torch.no_grad():
            teacher_logits = teacher(train_x)
            teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
        student_log_probs = F.log_softmax(
            student_logits / temperature, dim=-1
        )
        distillation_loss = F.kl_div(
            student_log_probs,
            teacher_probs,
            reduction="batchmean",
        )
        hard_target_loss = F.cross_entropy(student_logits, train_y)
        loss = (
            alpha * hard_target_loss
            + (1.0 - alpha) * (temperature**2) * distillation_loss
        )
        loss.backward()
        optimizer.step()


teacher = MLP(hidden_size=64)
baseline = MLP(hidden_size=16)
student = MLP(hidden_size=16)
fit_hard(teacher)
fit_hard(baseline)
fit_distilled(student, teacher)

for name, model in [
    ("teacher", teacher),
    ("baseline", baseline),
    ("distilled student", student),
]:
    print(
        f"{name}: validation={accuracy(model, valid_x, valid_y):.4f}, "
        f"test={accuracy(model, test_x, test_y):.4f}"
    )

A single run of this script is one toy measurement, not a general benchmark. For example, one CPU run with PyTorch 2.14.0 and scikit-learn 1.9.0 printed:

teacher: validation=0.9333, test=0.9333
baseline: validation=0.9333, test=0.9333
distilled student: validation=0.9000, test=0.9000

The exact accuracies can vary with library versions and initialization. The important details are that the student sees the training inputs and hard labels, the teacher supplies additional probabilities, and the final test set is kept out of optimization.

Key Concepts of Knowledge Distillation

The following concepts explain what the objective transfers and what it leaves to the student.

Teacher and Student Models

The teacher-student model configuration is central to knowledge distillation. The teacher model is typically a large, pre-trained neural network that has been trained on a comprehensive dataset. It possesses a high capacity to learn and generalize from data, capturing intricate patterns and representations. However, due to its size and complexity, deploying it on devices with limited computational resources may not be feasible.

The student model is a smaller network. It still trains on the input examples and their hard labels when those labels exist; teacher outputs provide an additional target. The student can therefore learn a task-specific approximation while using less compute or memory at inference time, but its accuracy must be measured for the task and setup at hand.

Soft Targets and Temperature

A common distillation target is a soft probability distribution over classes rather than a one-hot label. It can expose relative teacher preference between classes, while the student still receives the input that produced the distribution.

The softness of these targets is controlled by a temperature applied to the teacher and student logits before softmax. A higher temperature produces a softer distribution, while a lower temperature makes it sharper. The example multiplies the soft loss by T**2, the conventional scaling used with this temperature transformation.

Loss Function

The loss function used in knowledge distillation typically combines two components:

1. Distillation loss: This measures the difference between the teacher’s soft targets and the student’s temperature-scaled distribution. With PyTorch’s KLDivLoss, the student log-probabilities are the input and the teacher probabilities are the target.

2. Student Loss: This is the standard cross-entropy loss between the student’s predictions and the true labels.

The total loss is a weighted sum of these two components.

loss = alpha * hard_target_loss + (1.0 - alpha) * T**2 * distillation_loss

Where alpha is a hyperparameter that balances the importance of the two loss terms.

Feature-Based Distillation

In addition to distilling knowledge through soft targets, some approaches focus on transferring intermediate representations or features from the teacher to the student. This can be particularly useful when the architectures of the teacher and student models differ significantly. Feature-based distillation aims to align the intermediate activations or attention maps of the two models, encouraging the student to learn similar internal representations as the teacher.

Applications of Knowledge Distillation

Knowledge distillation has been studied across several domains. The benefit depends on the task and on how the teacher outputs are generated:

Image Classification

In computer vision, researchers have applied knowledge distillation to tasks such as object detection, image recognition, and semantic segmentation. Whether a smaller network preserves accuracy depends on the dataset, architecture, and distillation recipe.

Natural Language Processing (NLP)

Large language models can be distilled for tasks such as text classification, translation, and question answering. The DistilBERT paper reported 97% of BERT’s GLUE score with a model 40% smaller and 60% faster in its stated evaluation setup; those figures describe that benchmark and implementation, not every distilled language model.

Speech Recognition

In speech recognition systems, researchers have used knowledge distillation to study lower latency and compute. A smaller acoustic model can be useful for real-time applications when its measured quality remains sufficient for the workload.

Edge Computing

Knowledge distillation can make a task model small enough for edge devices such as smartphones, IoT devices, and embedded systems. A local deployment may reduce network dependence or latency and keep inputs on the device, but those benefits depend on the resulting model’s measured quality and the hardware budget.

Transfer Learning

Knowledge distillation can transfer task-specific behavior across architectures and model sizes. When labeled data is limited, teacher predictions may add supervision, but the student still depends on the available input examples and the quality of the teacher signal.

Ensemble Compression

Ensemble methods combine predictions from multiple models and can be expensive to serve. Ensemble distillation trains one student to approximate the ensemble on a chosen task; the approximation and its resource cost need to be measured.

Benefits of Knowledge Distillation

Knowledge distillation can offer practical benefits, especially when a smaller deployment model meets the task’s quality target:

Model Efficiency

The student has fewer parameters or a simpler architecture than the teacher in many distillation setups. That can reduce memory and compute per inference, which may make deployment on phones, embedded devices, or other constrained hardware easier. The quality tradeoff must be measured rather than assumed.

Maintained Performance

Some students retain a useful fraction of teacher performance after distillation. Soft targets can add information beyond a one-hot label, but they do not guarantee comparable accuracy or better generalization.

Reduced Training Time

Training a smaller model can reduce the cost of each student update and later evaluation, although generating teacher targets adds its own compute. The total training time depends on whether teacher predictions are cached, the dataset size, and the chosen architectures.

Ease of Deployment

When the student meets the required quality and latency target, its smaller footprint can simplify deployment. The deployment benefit is an engineering outcome to measure, including memory, latency, and hardware support.

Enhanced Generalization

Distillation can affect generalization because the student learns from teacher predictions as well as hard labels. It may help on some held-out tasks and hurt on others, so the relevant test set and baseline should accompany any claim.

Scalability and Accessibility

The lower resource requirement can make a model feasible in settings where the teacher is too expensive or large to serve. This is a deployment constraint, not a guarantee that the student will reproduce all of the teacher’s capabilities.

Performance Improvement

In some task-specific experiments, a student can outperform its teacher on the measured task. Such a result needs the task, model, data, and evaluation protocol because it is not a general consequence of distillation.

Challenges in Knowledge Distillation

While knowledge distillation offers significant benefits, it also presents several challenges that can impact its effectiveness:

Technical Complexities

The process of knowledge distillation involves several technical complexities. Training both a teacher and a student model requires more steps than training a single model, which can increase the overall computational burden. This complexity can make knowledge distillation less suitable for resource-constrained applications where computational resources are limited.

Difficulty in Multi-Task Learning

Knowledge distillation can be challenging when applied to multi-task learning scenarios. The student model may struggle to learn multiple tasks simultaneously, especially if the tasks require different types of knowledge or skills. This limitation can restrict the applicability of distillation techniques in environments where multi-task learning is essential.

Limitations Imposed by the Teacher Model

The student model is inherently limited by the capabilities of the teacher model. If the teacher model has biases or was trained on biased data, these biases may be inherited by the student model during the distillation process. Additionally, if the teacher model lacks certain information or capabilities, the student model will also lack them, potentially limiting its performance.

Loss of Information

During the distillation process, there is a potential for loss of minor details and nuances that the larger teacher model can interpret. While distilled models aim to emulate the performance of their larger counterparts, they may not capture all the subtleties present in the teacher’s predictions. This loss of information can affect the student’s ability to generalize effectively across different tasks or datasets.

Computational Overhead

Training both a teacher and a student model adds to computational overhead. The need for extensive hyperparameter tuning, such as adjusting the temperature parameter in soft label production, further complicates this process. Finding the optimal balance for these parameters can require significant experimentation and computational resources.

Sensitivity to Noisy Labels

Knowledge distillation can be sensitive to noisy labels in training data. If the teacher model’s predictions are based on noisy or unreliable data, these inaccuracies may be transmitted to the student model, affecting its performance. Ensuring high-quality training data is therefore essential to minimize this risk.

Limited Applicability on Proprietary Models

Knowledge distillation may have limited applicability when dealing with existing proprietary models. These models may not be easily accessible for modification or adaptation into a distillation framework. This restriction can hinder efforts to apply knowledge distillation techniques to certain commercial or closed-source systems.

Conclusion

Knowledge distillation gives a smaller model an additional training signal from a larger model. The student still needs the task inputs and a careful evaluation split. Soft targets, temperature scaling, and the balance with hard-label loss are choices to validate for each task. When the resulting student meets its quality and resource targets, it can be easier to serve; when it does not, the teacher’s probabilities do not make that outcome automatic.