""" CCT Batch Inference for CIFAR-10 ================================ Implements a shared-manifold model where a single global state evolves over time steps, processing the entire batch in parallel. Each sample only pays a small probe cost after the shared evolution. Key ideas from the World-Compute-AI framework: - Stationary component: shared hidden state S(t) - Probability component: per-sample probes that collapse S(t) into logits - Entropy collapse: regularization that encourages confident predictions - Amortized compute: O(C_shared + batch_size * C_probe) Train and test on CIFAR-10. """ import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.utils.data import DataLoader import torchvision import torchvision.transforms as transforms import argparse import time import numpy as np from tqdm import tqdm # ------------------------------- # 1. Shared ODE Cell (Stationary Kernel) # ------------------------------- class SharedODECell(nn.Module): """ Evolves the shared state S(t) -> S(t+1) using aggregated batch information. This cell replaces several transformer layers in the original theory. """ def __init__(self, state_dim, feature_dim, hidden_dim=256): super().__init__() # Aggregation network: takes current state + batch-aggregated features self.agg_net = nn.Sequential( nn.Linear(state_dim + feature_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, state_dim) ) # Residual connection self.res_scale = nn.Parameter(torch.tensor(0.1)) def forward(self, S, batch_features): """ S: current shared state (state_dim,) batch_features: aggregated features from all samples (feature_dim,) Returns updated shared state S_new. """ inp = torch.cat([S, batch_features], dim=-1) delta = self.agg_net(inp) S_new = S + self.res_scale * delta return S_new # ------------------------------- # 2. Full CCT Batch Inference Model # ------------------------------- class CCTBatchModel(nn.Module): """ Processes a batch of images using a shared state that evolves over T steps. After evolution, per-sample lightweight probes produce logits. """ def __init__(self, state_dim=128, feature_dim=64, probe_dim=32, num_steps=3, num_classes=10): super().__init__() self.state_dim = state_dim self.num_steps = num_steps # Image encoder: extract per-sample features (low-level) self.image_encoder = nn.Sequential( nn.Conv2d(3, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, feature_dim, 3, padding=1), nn.AdaptiveAvgPool2d(1) ) # output: (batch, feature_dim, 1, 1) -> squeeze to (batch, feature_dim) # Shared ODE cell (stationary kernel) self.ode_cell = SharedODECell(state_dim, feature_dim) # Lightweight probe network (per-sample collapse) self.probe = nn.Sequential( nn.Linear(state_dim + feature_dim, probe_dim), nn.ReLU(), nn.Linear(probe_dim, num_classes) ) # Initialize shared state self.register_buffer("initial_state", torch.randn(state_dim) * 0.02) def forward(self, x, return_entropy=False): """ x: input images (batch, 3, 32, 32) return_entropy: if True, also return average prediction entropy for regularization """ batch_size = x.size(0) # 1. Encode each image to feature vector features = self.image_encoder(x).squeeze(-1).squeeze(-1) # (batch, feature_dim) # 2. Compute batch-aggregated features (e.g., mean, but can be any permutation-invariant agg) batch_agg = features.mean(dim=0) # (feature_dim,) # 3. Evolve shared state over multiple steps (shared ODE solver) S = self.initial_state.clone() # start from initial manifold for _ in range(self.num_steps): S = self.ode_cell(S, batch_agg) # 4. Per-sample probe: combine shared state with individual features # This is the lightweight "collapse" step. S_expanded = S.unsqueeze(0).expand(batch_size, -1) # (batch, state_dim) probe_input = torch.cat([S_expanded, features], dim=-1) # (batch, state_dim+feature_dim) logits = self.probe(probe_input) if return_entropy: # Compute average entropy of softmax predictions (encourages collapse) probs = F.softmax(logits, dim=-1) entropy = -(probs * torch.log(probs + 1e-8)).sum(dim=-1).mean() return logits, entropy else: return logits # ------------------------------- # 3. Specialized Training Method (with Entropy Collapse Regularization) # ------------------------------- def train_one_epoch(model, dataloader, optimizer, criterion, device, entropy_weight=0.01): model.train() running_loss = 0.0 correct = 0 total = 0 for images, labels in tqdm(dataloader, desc="Training", leave=False): images, labels = images.to(device), labels.to(device) optimizer.zero_grad() logits, entropy = model(images, return_entropy=True) loss_ce = criterion(logits, labels) # Entropy regularization: encourage low-entropy (collapsed) predictions loss = loss_ce + entropy_weight * entropy loss.backward() optimizer.step() running_loss += loss.item() _, predicted = logits.max(1) total += labels.size(0) correct += predicted.eq(labels).sum().item() epoch_loss = running_loss / len(dataloader) epoch_acc = 100. * correct / total return epoch_loss, epoch_acc def evaluate(model, dataloader, criterion, device): model.eval() running_loss = 0.0 correct = 0 total = 0 with torch.no_grad(): for images, labels in tqdm(dataloader, desc="Evaluating", leave=False): images, labels = images.to(device), labels.to(device) logits = model(images, return_entropy=False) loss = criterion(logits, labels) running_loss += loss.item() _, predicted = logits.max(1) total += labels.size(0) correct += predicted.eq(labels).sum().item() epoch_loss = running_loss / len(dataloader) epoch_acc = 100. * correct / total return epoch_loss, epoch_acc # ------------------------------- # 4. Main: Train and Test on CIFAR-10 # ------------------------------- def main(): parser = argparse.ArgumentParser(description="CCT Batch Inference on CIFAR-10") parser.add_argument("--batch_size", type=int, default=128, help="Batch size for training") parser.add_argument("--epochs", type=int, default=30, help="Number of training epochs") parser.add_argument("--lr", type=float, default=0.001, help="Learning rate") parser.add_argument("--state_dim", type=int, default=128, help="Dimension of shared state") parser.add_argument("--feature_dim", type=int, default=64, help="Image feature dimension") parser.add_argument("--num_steps", type=int, default=3, help="Number of ODE steps (shared evolution)") parser.add_argument("--entropy_weight", type=float, default=0.01, help="Weight for entropy collapse regularization") parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu") args = parser.parse_args() device = torch.device(args.device) print(f"Using device: {device}") # Data preparation transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) transform_test = 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_train) testset = torchvision.datasets.CIFAR10(root='../data', train=False, download=True, transform=transform_test) trainloader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=2) testloader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, num_workers=2) # Model, optimizer, loss model = CCTBatchModel( state_dim=args.state_dim, feature_dim=args.feature_dim, num_steps=args.num_steps, num_classes=10 ).to(device) optimizer = optim.Adam(model.parameters(), lr=args.lr) criterion = nn.CrossEntropyLoss() print(f"Model has {sum(p.numel() for p in model.parameters()):,} parameters") print(f"Shared ODE steps: {args.num_steps}, Entropy weight: {args.entropy_weight}") best_acc = 0.0 for epoch in range(1, args.epochs + 1): start_time = time.time() train_loss, train_acc = train_one_epoch(model, trainloader, optimizer, criterion, device, args.entropy_weight) test_loss, test_acc = evaluate(model, testloader, criterion, device) epoch_time = time.time() - start_time print(f"Epoch {epoch:2d}/{args.epochs} | Time {epoch_time:.1f}s | " f"Train Loss: {train_loss:.4f} Acc: {train_acc:.2f}% | " f"Test Loss: {test_loss:.4f} Acc: {test_acc:.2f}%") if test_acc > best_acc: best_acc = test_acc torch.save(model.state_dict(), "cct_batch_best.pth") print(f" -> New best model saved (acc={best_acc:.2f}%)") print(f"\nTraining finished. Best test accuracy: {best_acc:.2f}%") # Optional: measure inference scaling (amortized cost) print("\n--- Inference Scaling Demonstration ---") model.eval() batch_sizes = [1, 16, 64, 128, 256] with torch.no_grad(): for bs in batch_sizes: dummy = torch.randn(bs, 3, 32, 32).to(device) # Warm-up for _ in range(5): _ = model(dummy) torch.cuda.synchronize() if device.type == 'cuda' else None start = time.time() for _ in range(20): _ = model(dummy) torch.cuda.synchronize() if device.type == 'cuda' else None elapsed = (time.time() - start) / 20 print(f"Batch size {bs:3d}: {elapsed*1000:.2f} ms per batch, " f"{elapsed/bs*1000:.3f} ms per sample (amortized)") if __name__ == "__main__": main()