import torch import torch.nn as nn import torch.nn.functional as F from torchvision import datasets, transforms from torch.utils.data import DataLoader # ============================================================================= # PNS CIRCUIT ELEMENTS # ============================================================================= class PNSCircuitLayer(nn.Module): """ A Layer that treats data as 'Particle Numbers'. Internal state is a multi-field tensor: [batch, channels, fields] Fields: 0: Linear (L), 1: Quadratic (Q), 2: Oscillatory (O), 3: Exponential (E), 4: Rest Mass (rho) """ def __init__(self, in_features, out_features): super().__init__() self.in_features = in_features self.out_features = out_features # Stable initialization: smaller weights to prevent instant saturation self.func_library = nn.Parameter(torch.randn(3, out_features, in_features, 3) * 0.01) # Initialize routing to favor the Linear function (Index 0) # This allows the model to start 'Classical' and evolve into 'Quantum' init_routing = torch.zeros(out_features, in_features, 3) init_routing[..., 0] = 1.0 self.routing_map = nn.Parameter(init_routing) self.bias_particles = nn.Parameter(torch.zeros(out_features, 5)) def forward(self, x): # x shape: [batch, in_features, 5] psi_x = x.sum(dim=-1) # [batch, in_features] # 1. THE ROUTING (The "Self-Questioning" Process) # Use a temperature-scaled softmax to prevent early saturation T = 1.0 routing_weights = torch.softmax(self.routing_map / T, dim=-1) # [out, in, 3] # TUNNELING: Add a small amount of noise to routing during training # This prevents the 'weights stop updating' bug by forcing exploration of subroutines if self.training: noise = torch.randn_like(routing_weights) * 0.01 routing_weights = torch.softmax(routing_weights + noise, dim=-1) # 2. THE SUBROUTINE LIBRARY all_results = [] # Type 0: Linear a0 = self.func_library[0, ..., 0] all_results.append(torch.einsum('bn,on->bo', psi_x, a0)) # Type 1: Oscillatory a1, w1, p1 = self.func_library[1, ..., 0], self.func_library[1, ..., 1], self.func_library[1, ..., 2] wave_input = (w1.unsqueeze(0) * psi_x.unsqueeze(1)) + p1.unsqueeze(0) all_results.append(torch.sum(a1.unsqueeze(0) * torch.sin(wave_input), dim=-1)) # Type 2: Exponential a2, w2 = self.func_library[2, ..., 0], self.func_library[2, ..., 1] exp_input = w2.unsqueeze(0) * psi_x.unsqueeze(1) all_results.append(torch.sum(a2.unsqueeze(0) * torch.exp(-torch.abs(exp_input)), dim=-1)) # 3. DYNAMIC PATCHING (Corrected Gradient Path) # Instead of summing routing_weights, we perform a weighted average per interaction # routing_weights: [out, in, 3], all_results[t]: [batch, out] final_interaction = torch.zeros_like(all_results[0]) for t in range(3): # Combine the routing for this function type and the function's result # route_strength: [out] = mean probability of using this function in this neuron route_strength = routing_weights[..., t].mean(dim=-1) final_interaction += route_strength.unsqueeze(0) * all_results[t] # 4. map back to multi-field PNS state using SIGNED interactions # We use sign(x)*abs(x)^k to preserve the particle's charge sign = torch.sign(final_interaction) abs_val = torch.abs(final_interaction) fields = [ final_interaction, # Linear (L) sign * (abs_val**2), # Quadratic (Q) - Preserves Charge torch.sin(final_interaction) * abs_val, # Oscillatory (O) - Amplitude modulation torch.exp(-abs_val) * sign, # Exponential (E) - Signed decay torch.mean(final_interaction, dim=1, keepdim=True).expand_as(final_interaction) # rho ] out_particles = torch.stack(fields, dim=-1) out_particles += self.bias_particles return out_particles class PNSDiode(nn.Module): """ ReLU equivalent: Only allows 'positive charge' particles to flow. """ def forward(self, x): # Collapses the particle to its scalar value to determine the charge sign # psi = L + Q + O + E + rho psi = x.sum(dim=-1, keepdim=True) # Diode gating: if psi < 0, the particle is annihilated to vacuum (0) gate = (psi > 0).float() return x * gate class PNSImpedanceMatcher(nn.Module): """ LayerNorm equivalent: Normalizes the energy (Quadratic field) to 1. """ def forward(self, x): # Soft-Normalization: We only normalize if the norm exceeds a threshold # This prevents the 'vanishing signal' problem while still stopping explosions norm = torch.norm(x, p=2, dim=-1, keepdim=True) scale = torch.clamp(1.0 / (norm + 1e-6), max=1.0) return x * scale class PNSCollapse(nn.Module): """ The Measurement Operator: Final collapse of all fields into a scalar. """ def forward(self, x): # Psi(P) = sum of all fields return x.sum(dim=-1) # ============================================================================= # FULL CIRCUIT MODEL # ============================================================================= class PNSClassifier(nn.Module): def __init__(self): super().__init__() # Input: 28*28 = 784. We map each pixel to a Particle (5 fields) self.layer1 = PNSCircuitLayer(784, 128) self.diode1 = PNSDiode() self.matcher1 = PNSImpedanceMatcher() self.layer2 = PNSCircuitLayer(128, 64) self.diode2 = PNSDiode() self.matcher2 = PNSImpedanceMatcher() self.layer3 = PNSCircuitLayer(64, 10) self.collapse = PNSCollapse() def forward(self, x): # 1. Injection: Convert scalar pixels into Particle Numbers # IMPORTANT: Scale input to prevent the Sine/Exp fields from saturating # Typical MNIST normalized values are -1 to 1. x_scaled = x * 0.1 batch_size = x.shape[0] p_in = torch.zeros((batch_size, 784, 5), device=x.device) p_in[..., 0] = x_scaled # Linear field excitation p_in[..., 4] = 1.0 # Every pixel starts with unit rest mass rho=1 # 2. Circuit Flow x = self.layer1(p_in) x = self.diode1(x) x = self.matcher1(x) x = self.layer2(x) x = self.diode2(x) x = self.matcher2(x) x = self.layer3(x) # 3. Measurement (Collapse) return self.collapse(x) # ============================================================================= # TRAINING LOOP # ============================================================================= def train_mnist(): # Data transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) train_set = datasets.MNIST('./data', train=True, download=True, transform=transform) test_set = datasets.MNIST('./data', train=False, transform=transform) train_loader = DataLoader(train_set, batch_size=64, shuffle=True) test_loader = DataLoader(test_set, batch_size=1000, shuffle=False) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = PNSClassifier().to(device) # Increased learning rate: We need a bigger 'Voltage' to move the routing weights optimizer = torch.optim.Adam(model.parameters(), lr=5e-3) criterion = nn.CrossEntropyLoss() model.train() for epoch in range(2): # Small epoch count for demonstration total_loss = 0 for batch_idx, (data, target) in enumerate(train_loader): data, target = data.view(-1, 784).to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() total_loss += loss.item() if batch_idx % 200 == 0: print(f"Epoch {epoch} [{batch_idx*64}/{len(train_loader)*64}] Loss: {loss.item():.4f}") # Evaluation model.eval() correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.view(-1, 784).to(device), target.to(device) output = model(data) pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() print(f"\nTest Accuracy: {100. * correct / len(test_loader.dataset):.2f}%") if __name__ == "__main__": train_mnist()