import time import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms from torch.utils.data import DataLoader # ========================================== # 1. THE CAUSAL-RESONANCE DOT PRODUCT LAYER # ========================================== class CRDPLinear(nn.Module): """ Causal-Resonance Dot Product Linear Layer. Causal-Resonance Formula: Out = (x @ W) + Sum_{t=1 to N}(MemoryKernel[t] * Trajectory[t]) + Resonance """ def __init__(self, in_features, out_features, memory_size=5): super(CRDPLinear, self).__init__() self.in_features = in_features self.out_features = out_features self.memory_size = memory_size # Stationary Weights (The 'Laws') -- fan-in scaled init instead of # an arbitrary *0.01, so layers start in a well-conditioned regime self.weight = nn.Parameter( torch.empty(in_features, out_features) ) nn.init.kaiming_uniform_(self.weight, a=5 ** 0.5) self.bias = nn.Parameter(torch.zeros(out_features)) # Memory Kernel M(t): initialize as a small, near-uniform weighted # average (not raw randn) so it starts as a gentle smoothing term # rather than injecting noise on the same scale as the signal. self.memory_kernel = nn.Parameter( torch.full((1, memory_size), 1.0 / memory_size) * 0.1 ) # Resonance term - global attractor shift self.resonance = nn.Parameter(torch.zeros(out_features)) # The buffer stores the trajectory of the flow self.trajectory_buffer = None def reset_memory(self): """Clear trajectory state. Call between epochs and before eval so that unrelated batches (train vs. test, or across shuffles) don't leak into each other through the memory kernel.""" self.trajectory_buffer = None def forward(self, x): batch_size = x.size(0) # 1. Snapshot Projection (a . b) current_proj = torch.matmul(x, self.weight) + self.bias # 2. Trajectory Buffer Management if self.trajectory_buffer is None or self.trajectory_buffer.size(0) != batch_size: self.trajectory_buffer = torch.zeros( batch_size, self.memory_size, self.out_features, device=x.device, dtype=current_proj.dtype ) # Shift window and append current projection updated_buffer = torch.cat([ self.trajectory_buffer[:, 1:, :].detach(), current_proj.unsqueeze(1) ], dim=1) self.trajectory_buffer = updated_buffer # 3. Causal Integration (The 'Star Product') kernel_expanded = self.memory_kernel.unsqueeze(0).expand(batch_size, -1, -1) memory_contribution = torch.bmm(kernel_expanded, self.trajectory_buffer) memory_contribution = memory_contribution.squeeze(1) # 4. Combine: Current Flow + Memory History + Resonance Attractor out = current_proj + memory_contribution + self.resonance return out # ========================================== # 2. THE CRDP-MLP ARCHITECTURE # ========================================== class CRDP_MLP(nn.Module): def __init__(self): super(CRDP_MLP, self).__init__() self.layer1 = CRDPLinear(784, 128) self.layer2 = CRDPLinear(128, 64) self.layer3 = CRDPLinear(64, 10) self.relu = nn.ReLU() def forward(self, x): x = x.view(-1, 784) x = self.relu(self.layer1(x)) x = self.relu(self.layer2(x)) x = self.layer3(x) return x def reset_memory(self): self.layer1.reset_memory() self.layer2.reset_memory() self.layer3.reset_memory() # ========================================== # 3. TRAINING AND TESTING PIPELINE # ========================================== def get_device(): if torch.cuda.is_available(): return torch.device("cuda") if torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") def run_experiment(): device = get_device() print(f"Using device: {device}") batch_size = 128 learning_rate = 0.001 epochs = 5 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = datasets.MNIST('../data', train=True, download=True, transform=transform) test_dataset = datasets.MNIST('../data', train=False, transform=transform) use_cuda = device.type == "cuda" train_loader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=use_cuda ) test_loader = DataLoader( test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=use_cuda ) model = CRDP_MLP().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) model.train() print("Starting Training with CRDP...") start = time.time() for epoch in range(epochs): # Reset memory state at the start of every epoch: the DataLoader # reshuffles each epoch, so last epoch's trailing batches have no # causal relationship to this epoch's first batches. #model.reset_memory() total_loss = 0 for data, target in train_loader: data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True) for i in range(10): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() if i==0: total_loss += loss.item() print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(train_loader):.4f}") elapsed = time.time() - start print(f"Training completed in {elapsed:.1f}s") # Reset memory before eval: otherwise the model's last training batch # (which had gradient updates applied *after* it was seen) bleeds into # the first test predictions through the trajectory buffer. model.reset_memory() model.eval() correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True) output = model(data) pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() print(f"\nFinal Test Accuracy: {100. * correct / len(test_loader.dataset):.2f}%") if __name__ == "__main__": run_experiment()