import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F import torchvision import torchvision.transforms as transforms import numpy as np import copy import matplotlib.pyplot as plt # ------------------------------- # 1. Model Definitions # ------------------------------- class SimpleCNN(nn.Module): """A small CNN that can serve as Ground or Bend model.""" def __init__(self, num_classes=10): super(SimpleCNN, self).__init__() self.conv1 = nn.Conv2d(3, 32, 3, padding=1) self.conv2 = nn.Conv2d(32, 64, 3, padding=1) self.pool = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(64 * 8 * 8, 128) self.fc2 = nn.Linear(128, num_classes) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) x = x.view(-1, 64 * 8 * 8) x = F.relu(self.fc1(x)) x = self.fc2(x) return x # ------------------------------- # 2. CCT Overfitting Detector # ------------------------------- class CCTOverfitDetector: """ Monitors divergence between Ground (stationary) and Bend (trainable) models. If entropy (divergence) > threshold on the same batch, overfitting is signaled. """ def __init__(self, ground_model, threshold=0.27, divergence='mse'): self.ground_model = ground_model self.ground_model.eval() # Frozen: stationary kernel self.threshold = threshold self.divergence_type = divergence self.entropy_history = [] def compute_entropy(self, bend_logits, ground_logits): """Measure divergence = entropy (CCT metric).""" if self.divergence_type == 'mse': return F.mse_loss(bend_logits, ground_logits).item() elif self.divergence_type == 'kl': # KL divergence requires probabilities p_bend = F.softmax(bend_logits, dim=1) p_ground = F.softmax(ground_logits, dim=1) return F.kl_div(p_bend.log(), p_ground, reduction='batchmean').item() else: raise ValueError("Unsupported divergence type") def check_batch(self, bend_model, batch_x, batch_y, loss_value): """ Evaluate entropy on the given batch. Returns (overfit_flag, entropy) """ with torch.no_grad(): ground_logits = self.ground_model(batch_x) bend_logits = bend_model(batch_x) entropy = self.compute_entropy(bend_logits, ground_logits) self.entropy_history.append(entropy) overfit = entropy > self.threshold if overfit: print(f"⚠️ CCT Overfit detected! Entropy={entropy:.4f} > {self.threshold}") return overfit, entropy # ------------------------------- # 3. Specialized Training Method # ------------------------------- def train_cct(model, ground_model, trainloader, testloader, epochs=50, lr=0.001, device='cuda'): """ Specialized training with CCT overfitting detection. - model: Bend model (trainable) - ground_model: Stationary kernel (frozen) """ model.to(device) ground_model.to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=lr) detector = CCTOverfitDetector(ground_model, threshold=0.27, divergence='mse') train_losses = [] test_accs = [] overfit_epoch = None warmup_epochs = 2 # Allow model to learn before checking entropy for epoch in range(epochs): model.train() running_loss = 0.0 epoch_entropy = 0.0 batch_count = 0 overfit = False # Initialize for warmup epochs for i, (inputs, labels) in enumerate(trainloader): 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() # CCT check on the SAME batch (after gradient step) # Only check after warmup to avoid false positives on unseen samples if epoch >= warmup_epochs: with torch.no_grad(): overfit, entropy = detector.check_batch(model, inputs, labels, loss.item()) epoch_entropy += entropy batch_count += 1 # Log overfitting but continue training (overload to bend model) if overfit: overfit_epoch = epoch avg_loss = running_loss / len(trainloader) avg_entropy = epoch_entropy / batch_count if batch_count > 0 else 0.0 train_losses.append(avg_loss) # Evaluate test accuracy model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in testloader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs, 1) total += labels.size(0) correct += (predicted == labels).sum().item() acc = 100 * correct / total test_accs.append(acc) print(f"Epoch {epoch+1:3d} | Loss: {avg_loss:.4f} | Entropy: {avg_entropy:.4f} | Test Acc: {acc:.2f}%") return model, train_losses, test_accs, overfit_epoch, detector.entropy_history # ------------------------------- # 4. Main: CIFAR-10 Experiment # ------------------------------- def main(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # Data preparation transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) trainset = torchvision.datasets.CIFAR10(root='../data', train=True, download=True, transform=transform) testset = torchvision.datasets.CIFAR10(root='../data', train=False, download=True, transform=transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=2) testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False, num_workers=2) # Create Ground Model (Stationary Kernel) - frozen ground_model = SimpleCNN(num_classes=10) # Optional: pretrain ground model slightly or keep random? # For CCT, ground should represent an invariant baseline. # Here we use a randomly initialized but frozen model. # In practice, you might want a simpler or pre-trained model. # Create Bend Model (trainable) - same architecture, but could be larger bend_model = SimpleCNN(num_classes=10) # Train with CCT overfitting detection trained_model, losses, accs, stop_epoch, entropy_hist = train_cct( model=bend_model, ground_model=ground_model, trainloader=trainloader, testloader=testloader, epochs=50, lr=0.001, device=device ) # Plot results plt.figure(figsize=(12,4)) plt.subplot(1,2,1) plt.plot(losses, label='Train Loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.title('Training Loss') plt.legend() plt.subplot(1,2,2) plt.plot(accs, label='Test Accuracy') plt.xlabel('Epoch') plt.ylabel('Accuracy (%)') plt.title('Test Accuracy') plt.legend() plt.tight_layout() plt.savefig('cct_cifar10_results.png') plt.show() print(f"\n=== Final ===") if stop_epoch is not None: print(f"Training stopped at epoch {stop_epoch+1} due to CCT overfit detection.") else: print("Training completed all epochs without crossing entropy threshold.") print(f"Final test accuracy: {accs[-1]:.2f}%") print(f"Entropy threshold: 0.27") if __name__ == "__main__": main()