import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms import numpy as np EPS = 1e-8 # ------------------------------------------------------------ # 1. Differentiable operators (replace addition) # ------------------------------------------------------------ def op_max(signals, dim=1): """Max across input signals (subdifferentiable)""" return signals.max(dim=dim)[0] def op_min(signals, dim=1): """Min across input signals""" return signals.min(dim=dim)[0] def op_product(signals, dim=1, eps=1e-8): """Product across signals – numerically stable via log-sum-exp.""" log_signals = torch.log(torch.abs(signals) + eps) sign = torch.sign(signals).prod(dim=dim) return sign * torch.exp(log_signals.sum(dim=dim)) def op_geometric_mean(signals, dim=1, eps=1e-8): """Geometric mean = (∏ x_i)^(1/N)""" n = signals.size(dim) log_sum = torch.log(torch.abs(signals) + eps).sum(dim=dim) sign = torch.sign(signals).prod(dim=dim) return sign * torch.exp(log_sum / n) def op_harmonic_mean(signals, dim=1, eps=1e-8): """Harmonic mean = N / (∑ 1/x_i)""" n = signals.size(dim) inv_sum = (1.0 / (signals + eps)).sum(dim=dim) return n / (inv_sum + eps) def op_softmax_fusion(signals, dim=1, temperature=1.0): """Weighted sum with softmax attention across signals.""" weights = torch.softmax(signals / temperature, dim=dim) return (weights * signals).sum(dim=dim) def op_modular_addition(signals, dim=1, period=2*np.pi): """(∑ x_i) mod period – phase wrapping effect.""" total = signals.sum(dim=dim) return total % period def op_multiplication(signals, dim=1): """Plain multiplication (alias for op_product)""" return op_product(signals, dim=dim) def op_synchronization(signals, dim=1): """Simplified phase coupling: sum of sin(θ_i - mean_θ)""" # treat signals as phases, compute coherent sum mean = signals.mean(dim=dim, keepdim=True) diff = signals - mean return torch.sin(diff).sum(dim=dim) def op_hadamard_mix(signals, dim=1, threshold=0.0): """Only keep features where all signals > threshold (sparse masking)""" mask = (signals > threshold).all(dim=dim) return signals.mean(dim=dim).masked_fill(~mask, 0.0) # Collection of usable operators (differentiable or subdifferentiable) OPERATORS = ( ('max', op_max), ('min', op_min), ('product', op_product), ('geometric_mean', op_geometric_mean), ('harmonic_mean', op_harmonic_mean), ('softmax_fusion', op_softmax_fusion), ('modular_add', op_modular_addition), ('synchronization', op_synchronization), ('hadamard_mix', op_hadamard_mix), ('mean', lambda s, dim=1: s.mean(dim=dim)), ) # ------------------------------------------------------------ # 2. Operator-Gated Linear Layer (OGCL) # ------------------------------------------------------------ class OperatorGatedLinear(nn.Module): """ Fully connected layer where the summation over input features is replaced by a learned mixture of operators. """ def __init__(self, in_features, out_features, bias=True, temp=0.5): super().__init__() self.in_features = in_features self.out_features = out_features self.temp = temp # Weight matrix: for each output neuron, we compute weighted inputs self.weight = nn.Parameter(torch.Tensor(out_features, in_features)) if bias: self.bias = nn.Parameter(torch.Tensor(out_features)) else: self.register_parameter('bias', None) # Tiny controller that chooses a weighted blend of operators. self.controller = nn.Sequential( nn.Linear(1, 8), nn.ReLU(), nn.Linear(8, len(OPERATORS)) ) self.reset_parameters() def reset_parameters(self): nn.init.kaiming_uniform_(self.weight, a=np.sqrt(5)) if self.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / np.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.bias, -bound, bound) def forward(self, x): # x: (batch, in_features) batch_size = x.size(0) # Per-sample entropy is cheaper and more interpretable than mixing over the batch axis. abs_x = x.abs() p = abs_x / (abs_x.sum(dim=1, keepdim=True) + EPS) entropy = -(p * torch.log(p + EPS)).sum(dim=1, keepdim=True) gate_logits = self.controller(entropy) op_weights = F.softmax(gate_logits / self.temp, dim=1) chunk_size = 32 output = x.new_empty(batch_size, self.out_features) for i in range(0, self.out_features, chunk_size): end = min(i + chunk_size, self.out_features) chunk_w = self.weight[i:end] signals = x.unsqueeze(1) * chunk_w.unsqueeze(0) if self.bias is not None: signals = signals + self.bias[i:end].view(1, -1, 1) chunk_out = signals.new_zeros(batch_size, end - i) for j, (_, op_func) in enumerate(OPERATORS): chunk_out.add_(op_func(signals, dim=2) * op_weights[:, j:j + 1]) output[:, i:end] = chunk_out return output # ------------------------------------------------------------ # 3. Full Model for CIFAR-10 (no Conv2d) # ------------------------------------------------------------ class OperatorCollapseNet(nn.Module): """ Simple MLP with OperatorGatedLinear layers. Flatten 3x32x32 -> 3072 -> 256 -> 128 -> 64 -> 10 """ def __init__(self, input_dim=3072, hidden_dims=[256, 128, 64], num_classes=10): super().__init__() layers = [] prev_dim = input_dim for hdim in hidden_dims: layers.append(OperatorGatedLinear(prev_dim, hdim)) layers.append(nn.BatchNorm1d(hdim)) layers.append(nn.ReLU()) prev_dim = hdim layers.append(OperatorGatedLinear(prev_dim, num_classes)) self.net = nn.Sequential(*layers) def forward(self, x): x = x.view(x.size(0), -1) # flatten return self.net(x) # ------------------------------------------------------------ # 4. Training Setup (CIFAR-10) # ------------------------------------------------------------ def train(model, device, train_loader, optimizer, epoch): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = F.cross_entropy(output, target) loss.backward() optimizer.step() if batch_idx % 100 == 0: print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ' f'({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}') def test(model, device, test_loader): model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) test_loss += F.cross_entropy(output, target, reduction='sum').item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader.dataset) print(f'\nTest set: Average loss: {test_loss:.4f}, ' f'Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.0f}%)\n') # ------------------------------------------------------------ # 5. Main # ------------------------------------------------------------ def main(): use_cuda = torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") print(f"Using device: {device}") # CIFAR-10 data (no convolutions, so just flatten) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) train_dataset = datasets.CIFAR10(root='../data', train=True, download=True, transform=transform) test_dataset = datasets.CIFAR10(root='../data', train=False, download=True, transform=transform) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=0) test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=0) model = OperatorCollapseNet().to(device) optimizer = optim.Adam(model.parameters(), lr=0.001) for epoch in range(1, 11): # 10 epochs train(model, device, train_loader, optimizer, epoch) test(model, device, test_loader) if __name__ == "__main__": main()