import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim import torchvision import torchvision.transforms as transforms import numpy as np from torch.utils.data import DataLoader from sklearn.metrics import confusion_matrix torch.backends.cudnn.benchmark = True torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True EPOCHS = 100 # raise/lower as you like; equity weighting makes long runs safer # ------------------------------- # 1. 3x3 Algebraic Cell with Spectral Regularization (vectorized) # ------------------------------- class AlgebraicCell3x3(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.weight = nn.Parameter(torch.randn(out_channels, in_channels, 3, 3) * 0.1) self.bias = nn.Parameter(torch.zeros(out_channels)) def forward(self, x): return F.conv2d(x, self.weight, bias=self.bias, padding=1) def spectral_regularization(self): mats = self.weight.mean(dim=1) sym = mats + mats.transpose(-1, -2) eigvals = torch.linalg.eigvalsh(sym) psd_pen = F.relu(-eigvals[:, 0]).sum() * 0.1 gaps = eigvals[:, 1:] - eigvals[:, :-1] collision_pen = torch.exp(-gaps.var(dim=1, unbiased=True)).sum() * 0.01 return psd_pen + collision_pen # ------------------------------- # 2. Invariant Cell Network # ------------------------------- class InvariantCellNet(nn.Module): def __init__(self, num_classes=10): super().__init__() self.conv1 = AlgebraicCell3x3(3, 64) self.conv2 = AlgebraicCell3x3(64, 128) self.conv3 = AlgebraicCell3x3(128, 256) self.pool = nn.MaxPool2d(2) self.drop = nn.Dropout(0.3) self.fc1 = nn.Linear(256 * 4 * 4, 512) self.fc2 = nn.Linear(512, num_classes) def forward(self, x): x = F.relu(self.conv1(x)); x = self.pool(x) x = F.relu(self.conv2(x)); x = self.pool(x) x = F.relu(self.conv3(x)); x = self.pool(x) x = x.flatten(1) x = F.relu(self.fc1(x)) x = self.drop(x) return self.fc2(x) def spectral_reg_loss(self): return (self.conv1.spectral_regularization() + self.conv2.spectral_regularization() + self.conv3.spectral_regularization()) # ------------------------------- # 3. Class Equity Monitor — the non-sacrificial strategy # ------------------------------- class ClassEquityMonitor: """Online detector AND corrector of class sacrifice. Repeating the same batch ~100 times reveals which classes the optimizer abandons: their loss stalls or rises while the favored classes' loss keeps dropping. But that signal is already present, for free, in the stream of normal batches — you just have to accumulate it per class. This monitor keeps an exponential moving average (EMA) of per-class cross-entropy across ordinary training batches. The EMA over ~100 recent batches is statistically the same evidence the 100-iteration inner loop would surface, at zero extra forward/backward cost. Correction ("altruism"): each batch is reweighted GroupDRO-style — classes with high EMA loss receive exponentially more gradient budget, classes that are comfortably learned donate theirs. The mean weight is always 1, so total gradient magnitude (and the effective learning rate) is unchanged; the budget is only redistributed, never inflated. Detection: a class whose EMA loss sits warn_z standard deviations above the mean is flagged as being sacrificed. """ def __init__(self, num_classes, device, ema=0.99, eta=1.0, warn_z=1.5): self.num_classes = num_classes self.ema = ema # 0.99 ≈ memory of the last ~100 batches self.eta = eta # >0 sharpens reallocation toward lagging classes self.warn_z = warn_z self.loss_ema = torch.zeros(num_classes, device=device) self.seen = torch.zeros(num_classes, device=device) @torch.no_grad() def update(self, per_sample_loss, targets): sums = torch.zeros(self.num_classes, device=per_sample_loss.device) cnts = torch.zeros(self.num_classes, device=per_sample_loss.device) sums.scatter_add_(0, targets, per_sample_loss) cnts.scatter_add_(0, targets, torch.ones_like(per_sample_loss)) present = cnts > 0 batch_mean = sums[present] / cnts[present] first_time = self.seen[present] == 0 new_vals = self.ema * self.loss_ema[present] + (1 - self.ema) * batch_mean self.loss_ema[present] = torch.where(first_time, batch_mean, new_vals) self.seen += cnts @torch.no_grad() def weights(self): """Per-class loss weights, mean exactly num_classes/num_classes = 1.""" w = torch.softmax(self.eta * self.loss_ema, dim=0) * self.num_classes return w @torch.no_grad() def sacrificed_classes(self): if (self.seen > 0).sum() < self.num_classes: return torch.tensor([], dtype=torch.long) mu = self.loss_ema.mean() sd = self.loss_ema.std().clamp_min(1e-8) z = (self.loss_ema - mu) / sd return torch.nonzero(z > self.warn_z).flatten().cpu() def report(self, class_names=None): flagged = self.sacrificed_classes() ema = self.loss_ema.cpu().numpy() w = self.weights().cpu().numpy() print(f" Per-class EMA loss: {np.array2string(ema, precision=3)}") print(f" Equity weights: {np.array2string(w, precision=3)}") if len(flagged) > 0: names = [class_names[i] if class_names else str(int(i)) for i in flagged] print(f" !! SACRIFICE WARNING: classes {names} are being abandoned " f"(loss >{self.warn_z} sd above mean). Gradient budget reallocated.") else: print(" No class sacrifice detected.") # ------------------------------- # 4. Diagnostics # ------------------------------- def detect_eigenvalue_collision(model): model.eval() with torch.no_grad(): # works whether or not the model was wrapped by torch.compile net = getattr(model, '_orig_mod', model) w = net.conv1.weight[0].mean(dim=0).cpu().numpy() eigvals = np.linalg.eigvalsh(w + w.T) gap_var = np.var(np.diff(eigvals)) collision = gap_var < 0.01 print(f"Eigenvalue collision score (low var = collapse): {gap_var:.4f} -> " f"{'COLLAPSED' if collision else 'OK'}") return collision def simpson_paradox_check(model, test_loader, device): model.eval() all_preds, all_labels = [], [] with torch.no_grad(): for images, labels in test_loader: images = images.to(device, non_blocking=True) with torch.autocast(device_type=device.type, enabled=(device.type == 'cuda')): outputs = model(images) all_preds.append(outputs.argmax(1).cpu()) all_labels.append(labels) all_preds = torch.cat(all_preds).numpy() all_labels = torch.cat(all_labels).numpy() cm = confusion_matrix(all_labels, all_preds) per_class_acc = cm.diagonal() / cm.sum(axis=1) overall_acc = np.mean(per_class_acc) low_classes = np.where(per_class_acc < overall_acc - 0.1)[0] if len(low_classes) > 0: print(f"Simpson's paradox warning: classes {low_classes} have accuracy " f"{per_class_acc[low_classes]} vs overall {overall_acc:.3f}") return overall_acc, per_class_acc # ------------------------------- # 5. Training loop — single pass per batch, equity-weighted loss # ------------------------------- def train(model, device, train_loader, optimizer, scaler, epoch, monitor, lambda_spectral=0.01, spectral_every=4): model.train() total_loss = 0.0 correct = torch.zeros(1, device=device) use_amp = device.type == 'cuda' for batch_idx, (data, target) in enumerate(train_loader): data = data.to(device, non_blocking=True) target = target.to(device, non_blocking=True) optimizer.zero_grad(set_to_none=True) # weights frozen for this step — computed from EMA, not this batch, # so the reweighting can't chase batch noise class_w = monitor.weights() with torch.autocast(device_type=device.type, enabled=use_amp): output = model(data) per_sample = F.cross_entropy(output, target, reduction='none') ce_loss = (class_w[target] * per_sample).mean() # feed the monitor the UNWEIGHTED losses (true difficulty signal) monitor.update(per_sample.detach().float(), target) if batch_idx % spectral_every == 0: spectral_loss = model.spectral_reg_loss() * spectral_every else: spectral_loss = 0.0 loss = ce_loss + lambda_spectral * spectral_loss scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() total_loss += loss.item() correct += (output.argmax(dim=1) == target).sum() avg_loss = total_loss / len(train_loader) acc = 100. * correct.item() / len(train_loader.dataset) print(f'Train Epoch {epoch}: Loss: {avg_loss:.4f}, Accuracy: {acc:.2f}%') return avg_loss, acc def test(model, device, test_loader): model.eval() test_loss = 0.0 correct = torch.zeros(1, device=device) use_amp = device.type == 'cuda' with torch.no_grad(): for data, target in test_loader: data = data.to(device, non_blocking=True) target = target.to(device, non_blocking=True) with torch.autocast(device_type=device.type, enabled=use_amp): output = model(data) test_loss += F.cross_entropy(output, target, reduction='sum').item() correct += (output.argmax(dim=1) == target).sum() test_loss /= len(test_loader.dataset) acc = 100. * correct.item() / len(test_loader.dataset) print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {acc:.2f}%') return acc # ------------------------------- # 6. Main # ------------------------------- def main(): 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) class_names = trainset.classes device = torch.device("cuda" if torch.cuda.is_available() else "cpu") pin = device.type == 'cuda' # Normal batch size again — the equity monitor replaces the giant batch # and the 20x same-batch inner loop train_loader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4, pin_memory=pin, persistent_workers=True) test_loader = DataLoader(testset, batch_size=512, shuffle=False, num_workers=4, pin_memory=pin, persistent_workers=True) model = InvariantCellNet(num_classes=10).to(device) try: model = torch.compile(model) except Exception: pass optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=5e-4) scaler = torch.amp.GradScaler(enabled=(device.type == 'cuda')) monitor = ClassEquityMonitor(num_classes=10, device=device, ema=0.99, eta=1.0, warn_z=1.5) print("=== METAFOUNDRY-ML Pre-Run Diagnosis ===") collision = detect_eigenvalue_collision(model) if collision: print("-> Eigenvalue collision detected before training. Spectral homotopy via regularization.") else: print("-> No severe collision, but regularization remains active.") best_acc = 0.0 for epoch in range(1, EPOCHS + 1): train(model, device, train_loader, optimizer, scaler, epoch, monitor, lambda_spectral=0.01) test_acc = test(model, device, test_loader) if test_acc > best_acc: best_acc = test_acc torch.save(model.state_dict(), 'best_invariant_cellnet.pth') # cheap online equity report every epoch monitor.report(class_names) if epoch % 10 == 0: print("\n--- Self-Diagnostic Report ---") detect_eigenvalue_collision(model) overall, per_class = simpson_paradox_check(model, test_loader, device) print(f"Overall test accuracy: {overall*100:.2f}%") print(f"Per-class accuracies: {per_class}") print("-------------------------------\n") model.load_state_dict(torch.load('best_invariant_cellnet.pth')) test(model, device, test_loader) print(f"\n=== Final Result ===") print(f"Best test accuracy: {best_acc:.2f}%") print("Non-sacrificial training: gradient budget continuously reallocated to lagging " "classes via EMA equity weights — same signal as 100 same-batch iterations, " "at the cost of one scatter-add per batch.") if __name__ == "__main__": main()