Distilling Knowledge: Making Large Models Smaller And Smarter

Large Language Models (LLMs), like GPT-4, have revolutionized AI, unlocking new possibilities, but they come with significant challenges. These models demand immense computational power and storage, making them costly and impractical for standard devices. Their complexity introduces latency, causing frustrating delays in real-time responses, and their overparameterization leads to inefficiencies, with many parameters adding little value. Accessibility is another concern, as only resource-rich organizations can afford to deploy them. Moreover, their high energy consumption raises serious environmental issues. While LLMs are undeniably powerful, addressing these challenges is essential for their broader adoption and sustainable use.

Why Do We Need Knowledge Distillation?

Knowledge Distillation (KD) offers a game-changing solution by transferring the knowledge of large models into smaller, more efficient ones. These compact models retain most of the original’s performance but are faster, lighter, and less resource-intensive. With KD, AI becomes more accessible and sustainable, enabling us to leverage the power of LLMs without their limitations. It’s a step toward making advanced AI both practical and scalable, paving the way for smarter, greener innovations.

What is Knowledge Distillation?

Knowledge Distillation (KD) is a machine learning technique where a large, powerful teacher model trains a smaller, more efficient student model by passing on its knowledge. Unlike traditional training that relies only on true labels, KD uses the teacher’s outputs — soft probabilities or logits — to guide the student. This approach helps the student model learn not just the correct answers but also the nuanced relationships between classes, enabling it to mimic the teacher’s behavior effectively. The result? A compact model that’s faster and lighter, yet still retains the teacher’s expertise.

Training student model using KD

How Does Knowledge Distillation Work?

Knowledge Distillation transfers knowledge from a large teacher model to a smaller student model. The key idea is to train the student model on the true labels of the dataset and the soft predictions (logits) generated by the teacher model. These logits carry rich information about the relationships between different classes that aren’t captured by hard labels.

The Process in Simple Terms:

1. Teacher Model Training: The teacher model is first trained on the dataset using traditional methods, achieving high accuracy by learning complex patterns and relationships in the data.
2. Generating Soft Labels: During inference, the teacher generates soft labels — probabilities for each class instead of a single “hard” label. For example, instead of predicting “cat” with 100% certainty, the teacher might output a distribution like cat: 0.7, dog: 0.2, rabbit: 0.1, reflecting its nuanced understanding.
3. Student Model Training: The student model is then trained using two losses:

• Cross-Entropy Loss (CE): This loss compares the student’s predicted probabilities to the true class labels, ensuring the student learns to classify correctly based on the actual data.
• Kullback-Leibler Divergence Loss (Distillation Loss): This loss measures the difference between the teacher’s softened logits (probabilities) and the student’s predictions. By scaling the logits using a temperature parameter, the teacher’s output becomes softer, allowing the student to mimic the teacher’s knowledge better while focusing on more nuanced information.

The Role of Temperature:

The temperature parameter (T) in KD smooths the logits from the teacher model, amplifying smaller probabilities to expose subtle relationships between classes. For example, increasing T may transform logits like [0.7, 0.2, 0.1] to [0.5, 0.3, 0.2], making it easier for the student to grasp these relationships. This approach was first explored in the original paper *Distilling the Knowledge in a Neural Network*, which demonstrated that even a smaller, simpler model could match the performance of larger models.

Intuition:

Imagine the teacher as a skilled mentor who doesn’t just provide correct answers but also explains why other options are less likely. This nuanced guidance helps the student develop a deeper understanding, enabling them to perform well without requiring the same level of complexity as the teacher. By blending direct supervision (true labels) with this informed guidance (teacher logits), the student learns to generalize effectively, achieving performance close to the teacher’s — while being faster, smaller, and more efficient.

Phew, enough theory! Time to roll up our sleeves and dive into the code 💻

*Install Dependencies

• Start by installing torch, torchvision for image datasets, models, and transformations.

!pip install -q torch torchvision

*Import Libraries & Setup Device

• Import the necessary libraries and set up the device (CPU/GPU) for training.

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

*Load CIFAR-10 Dataset

• We load the CIFAR-10 dataset with image transformations (e.g., normalization) for preprocessing and split it into training and testing datasets.

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=2)

*Define Teacher and Student Networks

• The teacher (DeepNN) is a larger, more complex network, while the student (LightNN) is lightweight, making it suitable for resource-constrained scenarios.

# Teacher Model (DeepNN)
class DeepNN(nn.Module):
    def __init__(self, num_classes=10):
        super(DeepNN, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(2048, 512),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

# Student Model (LightNN)
class LightNN(nn.Module):
    def __init__(self, num_classes=10):
        super(LightNN, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(16, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(1024, 256),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(256, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

*Define Training and Testing Functions

• It focus on training and testing both the Teacher (DeepNN) and Student (LightNN) models using only Cross-Entropy (CE) loss. The Teacher model is trained first to minimize the CE loss, and then the Student model is trained with the same loss, serving as a baseline. In real-world scenarios, the Teacher model is typically pre-trained, and Knowledge Distillation (KD) is used to transfer its knowledge to a lighter version of the Teacher model, creating a more efficient Student model that is ideal for resource-constrained environments.

# Training Function
def train(model, train_loader, epochs, learning_rate, device):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    model.train()
    for epoch in range(epochs):
        running_loss = 0.0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_loader)}")

# Testing Function
def test(model, test_loader, device):
    model.to(device)
    model.eval()

    correct, total = 0, 0
    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f"Test Accuracy: {accuracy:.2f}%")
    return accuracy

*Train Teacher and Student Models

• The teacher (DeepNN) is trained first, followed by the student (LightNN) using only cross-entropy loss for comparison.

torch.manual_seed(42)
teacher = DeepNN(num_classes=10).to(device)
train(teacher, train_loader, epochs=10, learning_rate=0.001, device=device)
test_accuracy_teacher = test(teacher, test_loader, device)

student = LightNN(num_classes=10).to(device)
train(student, train_loader, epochs=10, learning_rate=0.001, device=device)
test_accuracy_student = test(student, test_loader, device)

*Train Student Model with Knowledge Distillation

• The student model is trained with KD, using two losses: Kullback-Leibler divergence loss and Cross Entropy loss . The teacher’s soft logits are passed through a temperature-scaled softmax to create soft targets, which the student learns to match. Simultaneously, the student’s predictions are compared to the true labels using CE loss. Both losses are combined to update the student model, allowing it to benefit from the teacher’s knowledge while still learning from the actual data, helping it improve performance with fewer parameters.

def train_kd(teacher, student, train_loader, epochs, learning_rate, T, soft_target_loss_weight, ce_loss_weight, device):
    ce_loss = nn.CrossEntropyLoss()
    optimizer = optim.Adam(student.parameters(), lr=learning_rate)

    teacher.eval()
    student.train()

    for epoch in range(epochs):
        running_loss = 0.0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()

            # Get teacher predictions (soft targets)
            with torch.no_grad():
                teacher_logits = teacher(inputs)

            # Get student predictions
            student_logits = student(inputs)

            # Compute distillation loss
            soft_targets = nn.functional.softmax(teacher_logits / T, dim=1)
            soft_prob = nn.functional.log_softmax(student_logits / T, dim=1)
            soft_targets_loss = torch.sum(soft_targets * (soft_targets.log() - soft_prob)) / soft_prob.size()[0] * (T**2)

            # Compute cross-entropy loss
            label_loss = ce_loss(student_logits, labels)

            # Combine losses
            loss = soft_target_loss_weight * soft_targets_loss + ce_loss_weight * label_loss

            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_loader)}")

*Apply KD and Compare Results

• Finally, the student model is trained with KD and its accuracy is compared with and without the teacher’s help.

train_Vkd(teacher=teacher, student=student, train_loader=train_loader, epochs=10, learning_rate=0.001, T=2, soft_target_loss_weight=0.25, ce_loss_weight=0.75, device=device)
test_accuracy_student_kd = test(student, test_loader, device)

print(f"Teacher accuracy: {test_accuracy_teacher:.2f}%")
print(f"Student accuracy without KD: {test_accuracy_student:.2f}%")
print(f"Student accuracy with KD: {test_accuracy_student_kd:.2f}%")

The slight accuracy improvement with Knowledge Distillation (KD) on CIFAR-10 (70.63% vs. 70.22%) is due to the relatively simple nature of the dataset and the small model size. In cases where the student model already performs well, the gains from KD are often marginal. Additionally, CIFAR-10’s simplicity means the student can already capture most features without the need for extra knowledge transfer. However, on more complex datasets (e.g., ImageNet) or with larger, deeper models, KD can provide substantial improvements as the teacher model’s knowledge helps the student learn more complex features, resulting in better generalization and performance.

In real-world scenarios, using only KD loss without cross-entropy is generally not ideal. While KD loss helps the student model learn from the teacher’s logits, cross-entropy loss ensures the student also learns from the true labels, improving generalization. Combining both losses allows the student to benefit from the teacher’s guidance while also leveraging the actual data, leading to better performance. While it is possible to train with only KD loss, particularly when the teacher model has learned rich representations and no ground truth is available (as in unsupervised distillation), this approach requires a very strong teacher capable of providing meaningful soft labels. However, relying solely on the teacher’s predictions, especially in the presence of noisy data or errors, is not ideal. In most cases, particularly for general-purpose tasks like classification, combining KD loss with cross-entropy loss offers a more effective solution.