Model Compression: Quantization, Pruning, and Distillation

Module: Machine Learning | Difficulty: Advanced

Knowledge Distillation

where is temperature.

Quantization

Structured Pruning

Remove entire channels/filters:

Unstructured Pruning

Remove individual weights:

import torch
import torch.nn as nn

class DistillationLoss(nn.Module):
    def __init__(self, temperature=4.0, alpha=0.7):
        super().__init__()
        self.T = temperature; self.alpha = alpha
    def forward(self, student_out, teacher_out, labels):
        soft_loss = nn.functional.kl_div(
            nn.functional.log_softmax(student_out / self.T, dim=1),
            nn.functional.softmax(teacher_out / self.T, dim=1),
            reduction='batchmean') * (self.T ** 2)
        hard_loss = nn.functional.cross_entropy(student_out, labels)
        return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

Research Insight: The key insight for compression is that neural networks have significant redundancy. Pruning and quantization exploit this, while distillation transfers the teacher's dark knowledge (soft probabilities) to a smaller student.