Deep learning / Autoencoders and generative
Efficient Knowledge Distillation Techniques
Contents
- Knowledge Distillation Fundamentals
- Teacher Assistant Knowledge Distillation
- Basic Student-Teacher Architecture
- Temperature Scaling Implementation
- Training Loop with Progressive Knowledge Transfer
- Teacher Assistant Selection Strategy
- Multi-teacher Knowledge Distillation
- Online Knowledge Distillation
- Feature-based Knowledge Distillation
- Attention-based Knowledge Distillation
- Knowledge Distillation with Data Augmentation
- Progressive Layer-wise Knowledge Distillation
- Real-world Implementation: Image Classification
- Real-world Implementation: Performance Metrics and Visualization
- Additional Resources
Knowledge Distillation Fundamentals
Knowledge distillation is a model compression technique where a smaller student model learns to mimic the behavior of a larger teacher model. The process involves training the student model to match the soft probability distributions produced by the teacher model rather than hard class labels.
import torch
import torch.nn as nn
import torch.nn.functional as F
class DistillationLoss(nn.Module):
def __init__(self, temperature=2.0):
super().__init__()
self.temperature = temperature
def forward(self, student_logits, teacher_logits, labels):
# Soften probability distributions
soft_targets = F.softmax(teacher_logits / self.temperature, dim=1)
soft_prob = F.log_softmax(student_logits / self.temperature, dim=1)
# Calculate distillation loss
distillation_loss = F.kl_div(soft_prob, soft_targets, reduction='batchmean')
# Calculate standard cross entropy with true labels
student_loss = F.cross_entropy(student_logits, labels)
# Combine losses (usually with a weighted sum)
total_loss = (0.5 * student_loss) + (0.5 * distillation_loss)
return total_loss
# Example usage
student_outputs = torch.randn(32, 10) # batch_size=32, num_classes=10
teacher_outputs = torch.randn(32, 10)
labels = torch.randint(0, 10, (32,))
criterion = DistillationLoss()
loss = criterion(student_outputs, teacher_outputs, labels)
Teacher Assistant Knowledge Distillation
The Teacher Assistant (TA) approach introduces an intermediate model between the teacher and student. This model bridges the capacity gap, making knowledge transfer more effective. The TA model is trained first from the teacher, then transfers knowledge to the student.
class TeacherAssistantDistillation:
def __init__(self, teacher, assistant, student, temperature=2.0):
self.teacher = teacher
self.assistant = assistant
self.student = student
self.temperature = temperature
def train_assistant(self, dataloader, optimizer, epochs):
criterion = DistillationLoss(self.temperature)
self.teacher.eval()
self.assistant.train()
for epoch in range(epochs):
for inputs, labels in dataloader:
with torch.no_grad():
teacher_outputs = self.teacher(inputs)
assistant_outputs = self.assistant(inputs)
loss = criterion(assistant_outputs, teacher_outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
def train_student(self, dataloader, optimizer, epochs):
criterion = DistillationLoss(self.temperature)
self.assistant.eval()
self.student.train()
for epoch in range(epochs):
for inputs, labels in dataloader:
with torch.no_grad():
assistant_outputs = self.assistant(inputs)
student_outputs = self.student(inputs)
loss = criterion(student_outputs, assistant_outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
Basic Student-Teacher Architecture
The fundamental architecture consists of a teacher network that is typically larger and pre-trained, and a student network that is smaller but architecturally similar. Both networks process the same input but produce different output distributions.
class TeacherNetwork(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.classifier = nn.Linear(128 * 8 * 8, 10)
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), -1)
return self.classifier(x)
class StudentNetwork(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.classifier = nn.Linear(64 * 8 * 8, 10)
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), -1)
return self.classifier(x)
Temperature Scaling Implementation
Temperature scaling is crucial in knowledge distillation as it softens probability distributions, making it easier for the student to learn from the teacher's knowledge. Higher temperatures produce softer probability distributions, revealing more information about the teacher's learned patterns.
class TemperatureScaling:
def __init__(self, temperature=2.0):
self.temperature = temperature
def scale_logits(self, logits):
"""Scale logits using temperature and return softened probabilities"""
scaled_logits = logits / self.temperature
return torch.softmax(scaled_logits, dim=1)
def compare_distributions(self, original_logits):
"""Compare original and scaled probability distributions"""
original_probs = torch.softmax(original_logits, dim=1)
scaled_probs = self.scale_logits(original_logits)
print(f"Original probabilities:\n{original_probs[0]}")
print(f"Scaled probabilities (T={self.temperature}):\n{scaled_probs[0]}")
# Example usage
scaler = TemperatureScaling(temperature=3.0)
sample_logits = torch.tensor([[2.0, 1.0, 0.1, 0.01]])
scaler.compare_distributions(sample_logits)
Training Loop with Progressive Knowledge Transfer
This implementation showcases a progressive knowledge transfer approach where the student model gradually learns from both hard labels and teacher's soft targets. The balance between these two sources is controlled by a dynamically adjusted alpha parameter.
class ProgressiveDistillation:
def __init__(self, teacher, student, temperature=2.0, max_epochs=100):
self.teacher = teacher
self.student = student
self.temperature = temperature
self.max_epochs = max_epochs
def calculate_alpha(self, current_epoch):
"""Dynamic weighting between soft and hard targets"""
return min(1.0, current_epoch / (self.max_epochs * 0.3))
def train_epoch(self, dataloader, optimizer, epoch):
total_loss = 0
alpha = self.calculate_alpha(epoch)
for inputs, labels in dataloader:
# Get teacher's predictions
with torch.no_grad():
teacher_logits = self.teacher(inputs)
teacher_probs = torch.softmax(teacher_logits / self.temperature, dim=1)
# Get student's predictions
student_logits = self.student(inputs)
student_probs = torch.log_softmax(student_logits / self.temperature, dim=1)
# Calculate losses
distillation_loss = F.kl_div(student_probs, teacher_probs, reduction='batchmean')
student_loss = F.cross_entropy(student_logits, labels)
# Combine losses with dynamic alpha
loss = (alpha * distillation_loss) + ((1 - alpha) * student_loss)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
return total_loss / len(dataloader)
Teacher Assistant Selection Strategy
The selection of an appropriate teacher assistant model is crucial for effective knowledge transfer. This implementation provides a strategy to automatically select the best TA architecture based on the capacity gap between teacher and student.
class TASelector:
def __init__(self, teacher_params, student_params):
self.teacher_size = teacher_params
self.student_size = student_params
def calculate_optimal_ta_size(self):
"""Calculate optimal TA model size using geometric mean"""
return int(np.sqrt(self.teacher_size * self.student_size))
def generate_ta_architecture(self):
optimal_size = self.calculate_optimal_ta_size()
class TeacherAssistant(nn.Module):
def __init__(self, hidden_size):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, hidden_size, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(hidden_size, hidden_size*2, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.classifier = nn.Linear(hidden_size*2 * 8 * 8, 10)
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), -1)
return self.classifier(x)
return TeacherAssistant(optimal_size)
# Example usage
selector = TASelector(teacher_params=128, student_params=32)
ta_model = selector.generate_ta_architecture()
Multi-teacher Knowledge Distillation
Multi-teacher knowledge distillation leverages knowledge from multiple expert models simultaneously. This approach combines diverse knowledge sources to create a more robust student model, particularly useful when different teachers excel at different aspects of the task.
class MultiTeacherDistillation:
def __init__(self, teachers, student, temperature=2.0):
self.teachers = teachers # List of teacher models
self.student = student
self.temperature = temperature
self.teacher_weights = nn.Parameter(torch.ones(len(teachers)) / len(teachers))
def ensemble_teacher_predictions(self, inputs):
teacher_outputs = []
for teacher in self.teachers:
with torch.no_grad():
logits = teacher(inputs)
probs = F.softmax(logits / self.temperature, dim=1)
teacher_outputs.append(probs)
# Weighted average of teacher predictions
weighted_probs = sum(w * p for w, p in zip(
F.softmax(self.teacher_weights, dim=0),
teacher_outputs
))
return weighted_probs
def train_step(self, inputs, labels, optimizer):
# Get ensemble teacher predictions
teacher_ensemble_probs = self.ensemble_teacher_predictions(inputs)
# Student forward pass
student_logits = self.student(inputs)
student_probs = F.log_softmax(student_logits / self.temperature, dim=1)
# Calculate losses
distillation_loss = F.kl_div(student_probs, teacher_ensemble_probs)
student_loss = F.cross_entropy(student_logits, labels)
# Combined loss
total_loss = 0.7 * distillation_loss + 0.3 * student_loss
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
return total_loss.item()
Online Knowledge Distillation
Online knowledge distillation enables simultaneous training of both teacher and student networks. This approach is particularly efficient as it eliminates the need for pre-training the teacher model and allows for dynamic knowledge transfer during training.
class OnlineDistillation:
def __init__(self, teacher_model, student_model, temperature=2.0):
self.teacher = teacher_model
self.student = student_model
self.temperature = temperature
def mutual_learning_step(self, inputs, labels, teacher_opt, student_opt):
# Teacher forward pass
teacher_logits = self.teacher(inputs)
teacher_probs = F.softmax(teacher_logits / self.temperature, dim=1)
# Student forward pass
student_logits = self.student(inputs)
student_probs = F.softmax(student_logits / self.temperature, dim=1)
# Teacher learning from ground truth and student
teacher_loss = (
F.cross_entropy(teacher_logits, labels) +
F.kl_div(F.log_softmax(teacher_logits / self.temperature, dim=1),
student_probs.detach(), reduction='batchmean')
)
# Student learning from ground truth and teacher
student_loss = (
F.cross_entropy(student_logits, labels) +
F.kl_div(F.log_softmax(student_logits / self.temperature, dim=1),
teacher_probs.detach(), reduction='batchmean')
)
# Update both networks
teacher_opt.zero_grad()
teacher_loss.backward()
teacher_opt.step()
student_opt.zero_grad()
student_loss.backward()
student_opt.step()
return teacher_loss.item(), student_loss.item()
Feature-based Knowledge Distillation
This implementation focuses on transferring knowledge through intermediate feature representations rather than just final outputs. This approach helps the student learn rich internal representations from the teacher's feature maps.
class FeatureDistillation(nn.Module):
def __init__(self, teacher_channels, student_channels):
super().__init__()
self.transform = nn.Sequential(
nn.Conv2d(student_channels, teacher_channels, 1),
nn.BatchNorm2d(teacher_channels),
nn.ReLU()
)
def forward(self, teacher_features, student_features):
# Transform student features to match teacher dimensions
adapted_student = self.transform(student_features)
# Calculate feature similarity loss
similarity_loss = F.mse_loss(
self._normalize_features(adapted_student),
self._normalize_features(teacher_features)
)
return similarity_loss
def _normalize_features(self, features):
return F.normalize(features.view(features.size(0), -1), dim=1)
# Usage in training loop
feature_distiller = FeatureDistillation(
teacher_channels=256,
student_channels=128
)
def train_with_features(inputs, labels, teacher, student, optimizer):
teacher.eval()
student.train()
with torch.no_grad():
t_features, t_output = teacher(inputs, return_features=True)
s_features, s_output = student(inputs, return_features=True)
# Combine feature and output distillation
feature_loss = feature_distiller(t_features, s_features)
output_loss = F.kl_div(
F.log_softmax(s_output / 2.0, dim=1),
F.softmax(t_output / 2.0, dim=1)
)
total_loss = feature_loss + output_loss
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
Attention-based Knowledge Distillation
Attention-based knowledge distillation transfers the teacher's attention patterns to the student model. This approach helps the student focus on important regions of the input data, similar to how the teacher model processes information.
class AttentionDistillation(nn.Module):
def __init__(self):
super().__init__()
def compute_attention_map(self, features):
"""Compute spatial attention maps from feature maps"""
batch_size, channels, height, width = features.size()
attention = features.pow(2).mean(1, keepdim=True)
attention = F.normalize(attention.view(batch_size, -1), dim=1)
return attention.view(batch_size, 1, height, width)
def attention_loss(self, teacher_features, student_features):
"""Calculate attention transfer loss"""
teacher_attention = self.compute_attention_map(teacher_features)
student_attention = self.compute_attention_map(student_features)
return F.mse_loss(
student_attention,
teacher_attention,
reduction='mean'
)
class AttentionDistillationTrainer:
def __init__(self, teacher, student, temperature=2.0):
self.teacher = teacher
self.student = student
self.temperature = temperature
self.attention_distiller = AttentionDistillation()
def training_step(self, inputs, labels, optimizer):
# Get teacher's features and predictions
with torch.no_grad():
t_features, t_logits = self.teacher(inputs, return_features=True)
# Get student's features and predictions
s_features, s_logits = self.student(inputs, return_features=True)
# Calculate attention distillation loss
attention_loss = self.attention_distiller.attention_loss(
t_features, s_features
)
# Calculate standard distillation loss
distill_loss = F.kl_div(
F.log_softmax(s_logits / self.temperature, dim=1),
F.softmax(t_logits / self.temperature, dim=1),
reduction='batchmean'
)
# Combine losses
total_loss = 0.7 * distill_loss + 0.3 * attention_loss
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
return total_loss.item()
Knowledge Distillation with Data Augmentation
This implementation combines knowledge distillation with advanced data augmentation techniques to enhance the student model's learning process and improve generalization capabilities.
import torchvision.transforms as transforms
from torch.utils.data import Dataset
class AugmentedDistillationDataset(Dataset):
def __init__(self, base_dataset, num_augmentations=2):
self.dataset = base_dataset
self.num_augmentations = num_augmentations
self.augment = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(15),
transforms.ColorJitter(
brightness=0.2,
contrast=0.2,
saturation=0.2
),
transforms.RandomResizedCrop(
size=32,
scale=(0.8, 1.0),
ratio=(0.9, 1.1)
)
])
def __getitem__(self, idx):
image, label = self.dataset[idx]
augmented_images = [
self.augment(image) for _ in range(self.num_augmentations)
]
return augmented_images, label
def __len__(self):
return len(self.dataset)
class AugmentedDistillationTrainer:
def __init__(self, teacher, student, temperature=2.0):
self.teacher = teacher
self.student = student
self.temperature = temperature
def train_step(self, augmented_inputs, labels, optimizer):
teacher_outputs = []
student_outputs = []
# Process all augmented versions
for aug_input in augmented_inputs:
with torch.no_grad():
teacher_logits = self.teacher(aug_input)
teacher_outputs.append(
F.softmax(teacher_logits / self.temperature, dim=1)
)
student_logits = self.student(aug_input)
student_outputs.append(
F.log_softmax(student_logits / self.temperature, dim=1)
)
# Calculate consistency loss across augmentations
distill_losses = [
F.kl_div(s_out, t_out, reduction='batchmean')
for s_out, t_out in zip(student_outputs, teacher_outputs)
]
# Average losses
avg_distill_loss = sum(distill_losses) / len(distill_losses)
optimizer.zero_grad()
avg_distill_loss.backward()
optimizer.step()
return avg_distill_loss.item()
Progressive Layer-wise Knowledge Distillation
This advanced implementation focuses on transferring knowledge layer by layer progressively, ensuring better feature alignment between teacher and student networks while maintaining computational efficiency.
class LayerwiseDistillation:
def __init__(self, teacher_layers, student_layers):
self.layer_pairs = list(zip(teacher_layers, student_layers))
self.adaptation_layers = nn.ModuleList([
nn.Sequential(
nn.Conv2d(s.out_channels, t.out_channels, 1),
nn.BatchNorm2d(t.out_channels)
) for t, s in self.layer_pairs
])
def layer_loss(self, teacher_feat, student_feat, adapter):
adapted_student = adapter(student_feat)
return F.mse_loss(
F.normalize(adapted_student, dim=1),
F.normalize(teacher_feat, dim=1)
)
def progressive_training_step(self, inputs, current_layer_idx, optimizer):
teacher_features = []
student_features = []
# Forward pass through teacher layers
x_teacher = inputs
for i, (t_layer, _) in enumerate(self.layer_pairs):
x_teacher = t_layer(x_teacher)
if i <= current_layer_idx:
teacher_features.append(x_teacher.detach())
# Forward pass through student layers
x_student = inputs
for i, (_, s_layer) in enumerate(self.layer_pairs):
x_student = s_layer(x_student)
if i <= current_layer_idx:
student_features.append(x_student)
# Calculate layer-wise losses
layer_losses = []
for i in range(current_layer_idx + 1):
layer_loss = self.layer_loss(
teacher_features[i],
student_features[i],
self.adaptation_layers[i]
)
layer_losses.append(layer_loss)
total_loss = sum(layer_losses)
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
return total_loss.item()
Real-world Implementation: Image Classification
Complete implementation of knowledge distillation for a practical image classification task, demonstrating the entire pipeline from data preprocessing to model evaluation.
class ImageClassificationDistillation:
def __init__(self, num_classes, pretrained_teacher=True):
# Initialize teacher (ResNet50) and student (ResNet18)
self.teacher = models.resnet50(pretrained=pretrained_teacher)
self.student = models.resnet18(pretrained=False)
# Modify final layers for num_classes
self.teacher.fc = nn.Linear(2048, num_classes)
self.student.fc = nn.Linear(512, num_classes)
self.temperature = 2.0
self.alpha = 0.5 # Weight for distillation loss
def preprocess_data(self, batch_size=32):
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
])
# Load your dataset here
trainset = datasets.ImageFolder('path/to/train', transform=transform)
trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True)
return trainloader
def train_epoch(self, dataloader, optimizer, device):
self.teacher.eval()
self.student.train()
running_loss = 0.0
correct = 0
total = 0
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
# Teacher predictions
with torch.no_grad():
teacher_logits = self.teacher(inputs)
# Student predictions
student_logits = self.student(inputs)
# Calculate losses
distillation_loss = F.kl_div(
F.log_softmax(student_logits / self.temperature, dim=1),
F.softmax(teacher_logits / self.temperature, dim=1),
reduction='batchmean'
) * (self.temperature ** 2)
student_loss = F.cross_entropy(student_logits, labels)
# Combined loss
loss = (self.alpha * distillation_loss +
(1 - self.alpha) * student_loss)
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
# Calculate accuracy
_, predicted = student_logits.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
epoch_loss = running_loss / len(dataloader)
accuracy = 100. * correct / total
return epoch_loss, accuracy
Real-world Implementation: Performance Metrics and Visualization
This implementation focuses on comprehensive evaluation metrics and visualization tools to analyze the effectiveness of knowledge distillation in practical scenarios.
class DistillationMetrics:
def __init__(self, teacher, student):
self.teacher = teacher
self.student = student
self.metrics_history = {
'teacher_accuracy': [],
'student_accuracy': [],
'compression_ratio': self._calculate_compression_ratio(),
'inference_speedup': None
}
def _calculate_compression_ratio(self):
teacher_params = sum(p.numel() for p in self.teacher.parameters())
student_params = sum(p.numel() for p in self.student.parameters())
return teacher_params / student_params
def measure_inference_time(self, sample_input, num_runs=100):
self.teacher.eval()
self.student.eval()
# Measure teacher inference time
torch.cuda.synchronize()
teacher_start = time.time()
for _ in range(num_runs):
with torch.no_grad():
_ = self.teacher(sample_input)
torch.cuda.synchronize()
teacher_time = (time.time() - teacher_start) / num_runs
# Measure student inference time
torch.cuda.synchronize()
student_start = time.time()
for _ in range(num_runs):
with torch.no_grad():
_ = self.student(sample_input)
torch.cuda.synchronize()
student_time = (time.time() - student_start) / num_runs
self.metrics_history['inference_speedup'] = teacher_time / student_time
return teacher_time, student_time
def plot_learning_curves(self):
plt.figure(figsize=(12, 6))
plt.plot(self.metrics_history['teacher_accuracy'],
label='Teacher Accuracy')
plt.plot(self.metrics_history['student_accuracy'],
label='Student Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.title('Knowledge Distillation Learning Curves')
plt.legend()
plt.grid(True)
return plt.gcf()
def generate_performance_report(self):
report = {
'compression_ratio': self.metrics_history['compression_ratio'],
'inference_speedup': self.metrics_history['inference_speedup'],
'final_accuracy_gap': (self.metrics_history['teacher_accuracy'][-1] -
self.metrics_history['student_accuracy'][-1])
}
return report
Additional Resources
- Understanding Knowledge Distillation: A Survey
- Teacher Assistant Knowledge Distillation
- Feature-based Knowledge Distillation for Image Classification
- Progressive Layer-wise Distillation
- Attention-based Knowledge Distillation
- Multi-teacher Knowledge Distillation
- Online Knowledge Distillation
For implementation details and practical examples, search for:
- Pytorch knowledge distillation implementations
- Hugging Face transformers distillation examples
- TensorFlow model compression techniques