🧪 Defensive Distillation

1. Knowledge Distillation Fundamentals

Knowledge distillation is a technique where a smaller "student" model learns from a larger "teacher" model by mimicking its soft predictions.

Mathematical Foundation

L_distill = α * L_soft + (1-α) * L_hard

Temperature Scaling

import torch
import torch.nn.functional as F

def temperature_scaling(logits, temperature):
    """
    Apply temperature scaling to logits
    """
    return logits / temperature

def distillation_loss(student_logits, teacher_logits, labels, temperature, alpha):
    """
    Compute distillation loss
    """
    # Soft targets from teacher
    soft_targets = F.softmax(temperature_scaling(teacher_logits, temperature), dim=1)
    
    # Soft predictions from student
    soft_preds = F.log_softmax(temperature_scaling(student_logits, temperature), dim=1)
    
    # Distillation loss (KL divergence)
    soft_loss = F.kl_div(soft_preds, soft_targets, reduction='batchmean')
    
    # Hard loss (cross-entropy)
    hard_loss = F.cross_entropy(student_logits, labels)
    
    return alpha * soft_loss + (1 - alpha) * hard_loss

2. Defensive Distillation

Adaptation of knowledge distillation specifically designed to improve model robustness against adversarial attacks.

Key Differences from Standard Distillation

• Higher temperatures: Use higher temperature values (T > 1) to smooth predictions
• Adversarial training: Train student model on adversarial examples
• Robust teacher: Use a teacher model trained with adversarial training
• Gradient masking: Reduce gradient information available to attackers

Defensive Distillation Implementation

class DefensiveDistillation:
    def __init__(self, teacher_model, student_model, temperature=10.0, alpha=0.7):
        self.teacher_model = teacher_model
        self.student_model = student_model
        self.temperature = temperature
        self.alpha = alpha
        
    def train_step(self, x, y, optimizer, adversarial_generator=None):
        """
        Single training step with defensive distillation
        """
        # Generate adversarial examples if provided
        if adversarial_generator:
            x_adv = adversarial_generator.generate(x, y)
            x_combined = torch.cat([x, x_adv], dim=0)
            y_combined = torch.cat([y, y], dim=0)
        else:
            x_combined, y_combined = x, y
        
        # Teacher predictions
        with torch.no_grad():
            teacher_logits = self.teacher_model(x_combined)
            teacher_soft = F.softmax(teacher_logits / self.temperature, dim=1)
        
        # Student predictions
        student_logits = self.student_model(x_combined)
        
        # Compute distillation loss
        loss = distillation_loss(
            student_logits, teacher_logits, y_combined, 
            self.temperature, self.alpha
        )
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        return loss.item()
    
    def evaluate_robustness(self, test_loader, attack_method, epsilon=0.3):
        """
        Evaluate model robustness against attacks
        """
        self.student_model.eval()
        
        clean_correct = 0
        adv_correct = 0
        total = 0
        
        for x, y in test_loader:
            # Clean accuracy
            with torch.no_grad():
                clean_pred = self.student_model(x).argmax(dim=1)
                clean_correct += (clean_pred == y).sum().item()
            
            # Adversarial accuracy
            x_adv = attack_method.generate(x, y)
            with torch.no_grad():
                adv_pred = self.student_model(x_adv).argmax(dim=1)
                adv_correct += (adv_pred == y).sum().item()
            
            total += x.size(0)
        
        return {
            'clean_accuracy': clean_correct / total,
            'adversarial_accuracy': adv_correct / total,
            'robustness_gap': (clean_correct - adv_correct) / total
        }

3. Advanced Distillation Techniques

Enhanced distillation methods that provide additional security benefits.

Multi-Teacher Distillation

class MultiTeacherDistillation:
    def __init__(self, teachers, student, temperature=10.0):
        self.teachers = teachers
        self.student = student
        self.temperature = temperature
        
    def ensemble_teacher_predictions(self, x):
        """
        Combine predictions from multiple teachers
        """
        teacher_logits = []
        for teacher in self.teachers:
            with torch.no_grad():
                logits = teacher(x)
                teacher_logits.append(logits)
        
        # Average teacher predictions
        ensemble_logits = torch.stack(teacher_logits).mean(dim=0)
        return ensemble_logits
    
    def compute_distillation_loss(self, x, y):
        """
        Compute loss using ensemble teacher predictions
        """
        # Ensemble teacher predictions
        teacher_logits = self.ensemble_teacher_predictions(x)
        
        # Student predictions
        student_logits = self.student(x)
        
        # Distillation loss
        soft_targets = F.softmax(teacher_logits / self.temperature, dim=1)
        soft_preds = F.log_softmax(student_logits / self.temperature, dim=1)
        
        kl_loss = F.kl_div(soft_preds, soft_targets, reduction='batchmean')
        ce_loss = F.cross_entropy(student_logits, y)
        
        return 0.7 * kl_loss + 0.3 * ce_loss

Progressive Distillation

class ProgressiveDistillation:
    def __init__(self, teacher, student, temperature_schedule):
        self.teacher = teacher
        self.student = student
        self.temperature_schedule = temperature_schedule
        
    def train_with_progressive_temperature(self, train_loader, epochs):
        """
        Train with gradually increasing temperature
        """
        optimizer = torch.optim.Adam(self.student.parameters())
        
        for epoch in range(epochs):
            # Get current temperature
            current_temp = self.temperature_schedule[epoch]
            
            epoch_loss = 0
            for x, y in train_loader:
                loss = self.train_step(x, y, optimizer, current_temp)
                epoch_loss += loss
            
            print(f"Epoch {epoch+1}/{epochs}, Temp: {current_temp:.1f}, Loss: {epoch_loss/len(train_loader):.4f}")
    
    def train_step(self, x, y, optimizer, temperature):
        """
        Single training step with specified temperature
        """
        # Teacher predictions
        with torch.no_grad():
            teacher_logits = self.teacher(x)
            soft_targets = F.softmax(teacher_logits / temperature, dim=1)
        
        # Student predictions
        student_logits = self.student(x)
        
        # Distillation loss
        soft_preds = F.log_softmax(student_logits / temperature, dim=1)
        kl_loss = F.kl_div(soft_preds, soft_targets, reduction='batchmean')
        
        # Standard loss
        ce_loss = F.cross_entropy(student_logits, y)
        
        total_loss = 0.7 * kl_loss + 0.3 * ce_loss
        
        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()
        
        return total_loss.item()