#!/usr/bin/env python3 """ MNIST training script using the Paradox Hidden Layer. Because MNIST images are static 2D grids, the temporal convolution is re-interpreted as a **spatial** convolution along the flattened scan-line. The layer therefore detects high-contrast / oscillatory pixel transitions (edges, strokes) and switches to the tanh pathway exactly at those locations, while the smooth background is routed through the ReLU consensus pathway. """ import torch import torch.nn as nn import torch.nn.functional as F from torchvision import datasets, transforms from torch.utils.data import DataLoader import numpy as np # ============================================================================= # 1. LAYER DEFINITIONS # ============================================================================= class ParadoxLayer(nn.Module): """ Temporal Paradox Layer (sequence mode). Accepts (B, T, D) where T is time. See the earlier theory for the full derivation. """ def __init__(self, in_features: int, out_features: int, eps: float = 1e-5): super().__init__() self.in_features = in_features self.out_features = out_features self.eps = eps self.consensus = nn.Linear(in_features, out_features) self.paradox = nn.Linear(in_features, out_features) d1 = torch.tensor([1.0, -1.0]).view(1, 1, 2) self.register_buffer("d1_kernel", d1.repeat(in_features, 1, 1)) d2 = torch.tensor([1.0, -2.0, 1.0]).view(1, 1, 3) self.register_buffer("d2_kernel", d2.repeat(in_features, 1, 1)) self.temperature = nn.Parameter(torch.ones(1) * 2.0) self.bias = nn.Parameter(torch.zeros(1)) def _paradox_score(self, x_seq: torch.Tensor) -> torch.Tensor: """x_seq: (B, T, D) -> p: (B, T, 1)""" B, T, D = x_seq.shape x = x_seq.permute(0, 2, 1) # (B, D, T) v = F.conv1d(F.pad(x, (1, 0), mode="replicate"), self.d1_kernel, groups=D) a = F.conv1d(F.pad(x, (2, 0), mode="replicate"), self.d2_kernel, groups=D) v2 = v.pow(2).mean(dim=1, keepdim=True) # (B, 1, T) a2 = a.pow(2).mean(dim=1, keepdim=True) # (B, 1, T) v_sign = torch.sign(v) zcr = (v_sign[:, :, 1:] != v_sign[:, :, :-1]).float().mean(dim=2, keepdim=True) zcr = F.pad(zcr, (1, 0), mode="replicate") # (B, 1, T) ratio = a2 / (v2 + self.eps) score = self.temperature * (ratio * zcr + self.bias) p = torch.sigmoid(score) # (B, 1, T) return p.permute(0, 2, 1) # (B, T, 1) def forward(self, x: torch.Tensor) -> torch.Tensor: if x.dim() == 2: # Single step: temporal derivatives are degenerate (T=1), # so we fall back to a learned mixture of ReLU and tanh. p = torch.sigmoid(self.bias).expand(x.size(0), 1) # (B, 1) else: p = self._paradox_score(x) # (B, T, 1) y_c = F.relu(self.consensus(x)) y_p = torch.tanh(self.paradox(x)) return (1 - p) * y_c + p * y_p class SpatialParadoxLayer(nn.Module): """ Spatial variant for flat vectors. Treats the feature dimension as a 1D spatial axis (e.g. a scan-line). Computes spatial derivatives (finite differences) along the feature axis to detect local oscillations / edges. """ def __init__(self, in_features: int, out_features: int, eps: float = 1e-5): super().__init__() self.in_features = in_features self.out_features = out_features self.eps = eps self.consensus = nn.Linear(in_features, out_features) self.paradox = nn.Linear(in_features, out_features) # Fixed finite-difference kernels operating along the feature axis self.register_buffer("d1_kernel", torch.tensor([1.0, -1.0]).view(1, 1, 2)) self.register_buffer("d2_kernel", torch.tensor([1.0, -2.0, 1.0]).view(1, 1, 3)) self.temperature = nn.Parameter(torch.ones(1) * 2.0) self.bias = nn.Parameter(torch.zeros(1)) def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, D) B, D = x.shape x1d = x.unsqueeze(1) # (B, 1, D) # Spatial 1st and 2nd derivatives along the feature dimension v = F.conv1d(F.pad(x1d, (1, 0), mode="replicate"), self.d1_kernel) # (B, 1, D) a = F.conv1d(F.pad(x1d, (2, 0), mode="replicate"), self.d2_kernel) # (B, 1, D) # Energy v2 = v.pow(2).mean(dim=2, keepdim=True) # (B, 1, 1) a2 = a.pow(2).mean(dim=2, keepdim=True) # (B, 1, 1) # Zero-crossing rate of the spatial gradient v_sign = torch.sign(v) zcr = (v_sign[:, :, 1:] != v_sign[:, :, :-1]).float().mean(dim=2, keepdim=True) # (B, 1, 1) ratio = a2 / (v2 + self.eps) score = self.temperature * (ratio * zcr + self.bias) p = torch.sigmoid(score).view(B, 1) # (B, 1) y_c = F.relu(self.consensus(x)) y_p = torch.tanh(self.paradox(x)) return (1 - p) * y_c + p * y_p # ============================================================================= # 2. MODELS # ============================================================================= class ParadoxMLP(nn.Module): def __init__(self): super().__init__() self.net = nn.Sequential( SpatialParadoxLayer(784, 256), SpatialParadoxLayer(256, 128), nn.Linear(128, 10), ) def forward(self, x): x = x.view(x.size(0), -1) return self.net(x) class BaselineMLP(nn.Module): """Same capacity, standard ReLU MLP for comparison.""" def __init__(self): super().__init__() self.net = nn.Sequential( nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, 10), ) def forward(self, x): x = x.view(x.size(0), -1) return self.net(x) # ============================================================================= # 3. TRAINING & TESTING UTILITIES # ============================================================================= def train_epoch(model, loader, optimizer, criterion, device): model.train() total_loss = 0 correct = 0 for batch_idx, (data, target) in enumerate(loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() total_loss += loss.item() pred = output.argmax(dim=1) correct += pred.eq(target).sum().item() return total_loss / len(loader), 100.0 * correct / len(loader.dataset) def test_epoch(model, loader, criterion, device): model.eval() total_loss = 0 correct = 0 with torch.no_grad(): for data, target in loader: data, target = data.to(device), target.to(device) output = model(data) total_loss += criterion(output, target).item() pred = output.argmax(dim=1) correct += pred.eq(target).sum().item() return total_loss / len(loader), 100.0 * correct / len(loader.dataset) # ============================================================================= # 4. MAIN # ============================================================================= def main(epochs=5, batch_size=128, lr=1e-3): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") # MNIST loaders transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_loader = DataLoader( datasets.MNIST("../data", train=True, download=True, transform=transform), batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True ) test_loader = DataLoader( datasets.MNIST("../data", train=False, transform=transform), batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True ) criterion = nn.CrossEntropyLoss() # --- Train Paradox MLP --- print("\n=== Paradox MLP ===") model_p = ParadoxMLP().to(device) optimizer_p = torch.optim.Adam(model_p.parameters(), lr=lr) for epoch in range(1, epochs + 1): tr_loss, tr_acc = train_epoch(model_p, train_loader, optimizer_p, criterion, device) te_loss, te_acc = test_epoch(model_p, test_loader, criterion, device) print(f"Epoch {epoch}/{epochs} | " f"Train loss {tr_loss:.4f}, acc {tr_acc:.2f}% | " f"Test loss {te_loss:.4f}, acc {te_acc:.2f}%") # --- Train Baseline MLP --- print("\n=== Baseline ReLU MLP ===") model_b = BaselineMLP().to(device) optimizer_b = torch.optim.Adam(model_b.parameters(), lr=lr) for epoch in range(1, epochs + 1): tr_loss, tr_acc = train_epoch(model_b, train_loader, optimizer_b, criterion, device) te_loss, te_acc = test_epoch(model_b, test_loader, criterion, device) print(f"Epoch {epoch}/{epochs} | " f"Train loss {tr_loss:.4f}, acc {tr_acc:.2f}% | " f"Test loss {te_loss:.4f}, acc {te_acc:.2f}%") # --- Inspect paradox gates on a single test batch --- print("\n=== Paradox Gate Statistics (test batch) ===") data, target = next(iter(test_loader)) data = data.to(device) with torch.no_grad(): # Hook into the first spatial paradox layer gate_values = [] def hook(m, inp, out): # Recompute gate for inspection x = inp[0].view(inp[0].size(0), -1) B, D = x.shape x1d = x.unsqueeze(1) v = F.conv1d(F.pad(x1d, (1, 0), mode="replicate"), m.d1_kernel) a = F.conv1d(F.pad(x1d, (2, 0), mode="replicate"), m.d2_kernel) v2 = v.pow(2).mean(dim=2, keepdim=True) a2 = a.pow(2).mean(dim=2, keepdim=True) v_sign = torch.sign(v) zcr = (v_sign[:, :, 1:] != v_sign[:, :, :-1]).float().mean(dim=2, keepdim=True) ratio = a2 / (v2 + m.eps) score = m.temperature * (ratio * zcr + m.bias) p = torch.sigmoid(score).view(-1) gate_values.append(p.cpu().numpy()) handle = model_p.net[0].register_forward_hook(hook) _ = model_p(data) handle.remove() gates = np.concatenate(gate_values) print(f" Gate mean: {gates.mean():.4f}") print(f" Gate std : {gates.std():.4f}") print(f" Gate >0.5: {(gates > 0.5).mean()*100:.1f}% of pixels") print(" (High gate = image region is spatially oscillatory / edgy)") return model_p, model_b if __name__ == "__main__": main(epochs=5, batch_size=128, lr=1e-3)