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 copy import matplotlib.pyplot as plt # ------------------------------- # Simple CNN Model # ------------------------------- class SimpleCNN(nn.Module): 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 # ------------------------------- # CCT Relaxation & Consolidation Trainer # ------------------------------- def train_cct_relax_consolidate(bend_model, ground_model, trainloader, testloader, epochs=50, lr=0.001, device='cuda', entropy_threshold=0.27, relax_threshold=0.15, ground_update_decay=0.99, # how fast ground absorbs bend reset_on_overfit=True, entropy_penalty_weight=0.5): """ - bend_model: fast learner (trainable) - ground_model: slow anchor (initially copy of bend) - ground_update_decay: high = slow consolidation """ bend_model.to(device) ground_model.to(device) ground_model.eval() optimizer = optim.Adam(bend_model.parameters(), lr=lr) criterion_ce = nn.CrossEntropyLoss() entropy_history = [] test_acc_history = [] reset_epochs = [] consolidate_epochs = [] for epoch in range(epochs): bend_model.train() running_ce_loss = 0.0 running_entropy = 0.0 batch_count = 0 for inputs, labels in trainloader: inputs, labels = inputs.to(device), labels.to(device) outputs = bend_model(inputs) loss_ce = criterion_ce(outputs, labels) # Entropy = divergence from ground model with torch.no_grad(): ground_out = ground_model(inputs) entropy = F.mse_loss(outputs, ground_out).item() # Total loss = cross‑entropy + small entropy penalty # (keeps bend from drifting too fast) loss = loss_ce + entropy_penalty_weight * entropy optimizer.zero_grad() loss.backward() optimizer.step() running_ce_loss += loss_ce.item() running_entropy += entropy batch_count += 1 avg_entropy = running_entropy / batch_count entropy_history.append(avg_entropy) # ---- CCT Logic ---- if avg_entropy > entropy_threshold: # Overfitting → reset bend to ground if reset_on_overfit: bend_model.load_state_dict(copy.deepcopy(ground_model.state_dict())) print(f"⚠️ Epoch {epoch+1}: Entropy {avg_entropy:.3f} > {entropy_threshold} → RESET bend to ground") reset_epochs.append(epoch+1) # Optionally reinitialize optimizer state? elif avg_entropy < relax_threshold: # Relaxation achieved → consolidate bend into ground # Slow update: ground = decay*ground + (1-decay)*bend with torch.no_grad(): for g_param, b_param in zip(ground_model.parameters(), bend_model.parameters()): g_param.data.mul_(ground_update_decay).add_(b_param.data, alpha=1 - ground_update_decay) print(f"✓ Epoch {epoch+1}: Entropy {avg_entropy:.3f} < {relax_threshold} → CONSOLIDATE ground ← bend") consolidate_epochs.append(epoch+1) # Optional: also reset bend to ground after consolidation? # If you reset, bend starts fresh from the improved ground. # Uncomment next line to reset bend after consolidation: # bend_model.load_state_dict(copy.deepcopy(ground_model.state_dict())) else: print(f" Epoch {epoch+1}: Entropy {avg_entropy:.3f} (normal range)") # Evaluate test accuracy bend_model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in testloader: inputs, labels = inputs.to(device), labels.to(device) outputs = bend_model(inputs) _, predicted = torch.max(outputs, 1) total += labels.size(0) correct += (predicted == labels).sum().item() acc = 100 * correct / total test_acc_history.append(acc) print(f" Test Acc: {acc:.2f}%") return bend_model, ground_model, entropy_history, test_acc_history, reset_epochs, consolidate_epochs # ------------------------------- # Main # ------------------------------- def main(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using {device}") # Data 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) # Initialize models bend = SimpleCNN(num_classes=10) ground = SimpleCNN(num_classes=10) ground.load_state_dict(bend.state_dict()) # start identical # Train with CCT relaxation & consolidation bend_final, ground_final, entropy_hist, acc_hist, reset_ep, cons_ep = train_cct_relax_consolidate( bend_model=bend, ground_model=ground, trainloader=trainloader, testloader=testloader, epochs=50, lr=0.001, device=device, entropy_threshold=0.27, relax_threshold=0.15, ground_update_decay=0.99, reset_on_overfit=True, entropy_penalty_weight=0.5 ) # Plot plt.figure(figsize=(12,4)) plt.subplot(1,2,1) plt.plot(entropy_hist, label='Entropy (bend vs ground)') plt.axhline(y=0.27, color='r', linestyle='--', label='Overfit threshold') plt.axhline(y=0.15, color='g', linestyle='--', label='Relax threshold') for ep in reset_ep: plt.axvline(x=ep-1, color='r', alpha=0.3, linestyle=':', label='Reset' if ep==reset_ep[0] else '') for ep in cons_ep: plt.axvline(x=ep-1, color='g', alpha=0.3, linestyle=':', label='Consolidate' if ep==cons_ep[0] else '') plt.xlabel('Epoch') plt.ylabel('Entropy') plt.title('CCT Entropy – Resets prevent overfitting') plt.legend() plt.subplot(1,2,2) plt.plot(acc_hist, label='Test Accuracy') plt.xlabel('Epoch') plt.ylabel('Accuracy (%)') plt.title('No overfitting – steady improvement') plt.legend() plt.tight_layout() plt.savefig('cct_relax_consolidate.png') plt.show() print("\n=== Final ===") print(f"Resets at epochs: {reset_ep}") print(f"Consolidations at epochs: {cons_ep}") print(f"Final test accuracy: {acc_hist[-1]:.2f}%") print("The bend model never overfits because it is reset whenever entropy > 0.27.") print("Useful learning is consolidated into the ground model during low‑entropy relaxation periods.") if __name__ == "__main__": main()