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 # ---------- FluxTensor (same as before, simplified) ---------- class FluxTensor: def __init__(self, v, s=None, t=None): self.v = v self.s = torch.zeros_like(v) if s is None else s self.t = torch.zeros_like(v) if t is None else t def __add__(self, other): if isinstance(other, FluxTensor): return FluxTensor(self.v + other.v, torch.sqrt(self.s**2 + other.s**2 + 1e-8), self.t + other.t) else: return FluxTensor(self.v + other, self.s, self.t) def __sub__(self, other): if isinstance(other, FluxTensor): return FluxTensor(self.v - other.v, torch.sqrt(self.s**2 + other.s**2 + 1e-8), self.t - other.t) else: return FluxTensor(self.v - other, self.s, self.t) def __mul__(self, other): if isinstance(other, FluxTensor): return FluxTensor(self.v * other.v, torch.sqrt((self.v * other.s)**2 + (other.v * self.s)**2 + 1e-8), self.v * other.t + other.v * self.t) else: return FluxTensor(self.v * other, self.s * abs(other), self.t * other) def __rmul__(self, other): return self.__mul__(other) def collapse(self, work): return FluxTensor(self.v, torch.clamp(self.s - work, min=0.0), self.t) def to(self, device): return FluxTensor(self.v.to(device), self.s.to(device), self.t.to(device)) def clone(self): return FluxTensor(self.v.clone(), self.s.clone(), self.t.clone()) # ---------- Anti‑Resonance Operator ---------- class AntiResonance: def __init__(self, beta=0.9, gamma=0.5): self.beta = beta self.gamma = gamma self.buffer = None # will hold the low‑pass filtered error def __call__(self, grad_flux): """ grad_flux: FluxTensor whose .v is the raw gradient. Returns: FluxTensor after anti‑resonance. """ if self.buffer is None: self.buffer = grad_flux.clone() # initialize else: # Update momentum buffer: buffer = (1-beta)*grad + beta*buffer self.buffer = (1 - self.beta) * grad_flux + self.beta * self.buffer # Anti‑resonance: subtract gamma times the low‑pass component e_anti = grad_flux - self.gamma * self.buffer return e_anti # ---------- Iterative Flux MLP with Anti‑Resonance ---------- class IterativeFluxMLP_AR(nn.Module): def __init__(self, input_dim=784, hidden_dim=128, state_dim=64, num_classes=10, num_refinements=5, work_per_step=0.05, anti_beta=0.9, anti_gamma=0.5): super().__init__() self.state_dim = state_dim self.num_refinements = num_refinements self.work_per_step = work_per_step # Encoder: x -> embedding self.encoder = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, state_dim) ) # Refiner: takes [x_emb, y, e_anti] -> new y refiner_input_dim = state_dim + state_dim + state_dim # x_emb + y + e_anti self.refiner = nn.Sequential( nn.Linear(refiner_input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, state_dim) ) self.classifier = nn.Linear(state_dim, num_classes) # Anti‑resonance module for the error signal self.anti_res = AntiResonance(beta=anti_beta, gamma=anti_gamma) def forward(self, x, compute_error=False): batch = x.shape[0] device = x.device # Encode input x_emb = self.encoder(x) # [batch, state_dim] x_flux = FluxTensor(x_emb, torch.zeros_like(x_emb), torch.zeros_like(x_emb)) # Initial hidden state y (high entropy) y = FluxTensor( v=torch.zeros(batch, self.state_dim, device=device), s=torch.full((batch, self.state_dim), 0.5, device=device), t=torch.zeros(batch, self.state_dim, device=device) ) # We'll store each refinement step's y for gradient computation y_steps = [y] for step in range(self.num_refinements): # Concatenate x_emb and current y cat_v = torch.cat([x_flux.v, y.v], dim=-1) cat_s = torch.cat([x_flux.s, y.s], dim=-1) cat_t = torch.cat([x_flux.t, y.t], dim=-1) cat = FluxTensor(cat_v, cat_s, cat_t) # Apply refiner (we need to simulate a linear layer with flux arithmetic) # For simplicity, we use plain nn.Linear and treat the result as FluxTensor # with zero entropy. A full flux refiner would use FluxLinear. # Here we cheat: we take the plain tensor output and wrap it. # Pad with zeros for e_anti to match refiner input dim (x_emb + y + e_anti) e_anti_pad = torch.zeros(batch, self.state_dim, device=device) cat_full = torch.cat([cat_v, e_anti_pad], dim=-1) h = self.refiner[0](cat_full) # first linear h = F.relu(h) new_y_v = self.refiner[2](h) # second linear # We treat the new y as having the same entropy as old y (simplified) new_y = FluxTensor(new_y_v, y.s.clone(), y.t.clone()) # Apply collapse and evolution new_y = new_y.collapse(self.work_per_step) # (Evolution omitted for brevity) y = new_y y_steps.append(y) # Final classification using the last y's value logits = self.classifier(y.v) # If compute_error is True, also return the gradient of loss w.r.t. each y_step if compute_error: return logits, y_steps else: return logits def get_error_gradients(self, x, target): """Compute gradients of the loss w.r.t. each refinement step's y.""" # Forward pass with gradient tracking logits, y_steps = self.forward(x, compute_error=True) loss = F.cross_entropy(logits, target) # Skip y_steps[0] (initial y) since it doesn't require grad grads = torch.autograd.grad(loss, [step.v for step in y_steps[1:]], create_graph=False) # Return zero grad for initial y, then actual grads for refined steps zero_flux = FluxTensor(torch.zeros_like(y_steps[0].v), torch.zeros_like(y_steps[0].v), torch.zeros_like(y_steps[0].v)) grad_fluxes = [zero_flux] + [FluxTensor(g, torch.zeros_like(g), torch.zeros_like(g)) for g in grads] return grad_fluxes def refinement_with_anti_resonance(self, x, target): """Perform the iterative refinement while using anti‑resonant error as extra input.""" batch = x.shape[0] device = x.device x_emb = self.encoder(x) x_flux = FluxTensor(x_emb, torch.zeros_like(x_emb), torch.zeros_like(x_emb)) # Initial y y = FluxTensor( v=torch.zeros(batch, self.state_dim, device=device), s=torch.full((batch, self.state_dim), 0.5, device=device), t=torch.zeros(batch, self.state_dim, device=device) ) # We need the error gradient after each step. For simplicity, # we perform a forward pass to get the final loss and then backpropagate # through the whole unrolled graph. Then we can extract gradients at each step. # This is done in `get_error_gradients`. Here we use those gradients # inside the refinement loop (requires re‑computing graph). # For a true online version, we would use implicit differentiation. # For demonstration, we do one full forward/backward to get gradients, # then use them in a second forward pass. This is inefficient but clear. # Compute gradients w.r.t. each y step (using the current model state) grad_fluxes = self.get_error_gradients(x, target) # list of FluxTensor gradients # Now re‑run the refinement, but this time we feed the anti‑resonant error # as an additional input. We'll ignore the original classifier loss during this pass. y = FluxTensor( v=torch.zeros(batch, self.state_dim, device=device), s=torch.full((batch, self.state_dim), 0.5, device=device), t=torch.zeros(batch, self.state_dim, device=device) ) self.anti_res.buffer = None # reset anti‑resonance buffer for step in range(self.num_refinements): # Get the raw error gradient for this step e_raw = grad_fluxes[step] # FluxTensor # Apply anti‑resonance e_anti = self.anti_res(e_raw) # Concatenate x_emb, y, and e_anti cat_v = torch.cat([x_flux.v, y.v, e_anti.v], dim=-1) cat_s = torch.cat([x_flux.s, y.s, e_anti.s], dim=-1) cat_t = torch.cat([x_flux.t, y.t, e_anti.t], dim=-1) cat = FluxTensor(cat_v, cat_s, cat_t) # Refiner (again, using plain linear layers for simplicity) h = self.refiner[0](cat.v) h = F.relu(h) new_y_v = self.refiner[2](h) new_y = FluxTensor(new_y_v, y.s.clone(), y.t.clone()) new_y = new_y.collapse(self.work_per_step) y = new_y # Final logits from the refined y (no extra error input) logits = self.classifier(y.v) return logits # ---------- Training loop ---------- def train(model, device, train_loader, optimizer, epoch): model.train() correct = 0 total_loss = 0 for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) data = data.view(data.size(0), -1) optimizer.zero_grad() # Use the anti‑resonance‑aware refinement logits = model.refinement_with_anti_resonance(data, target) loss = F.cross_entropy(logits, target) loss.backward() optimizer.step() total_loss += loss.item() pred = logits.argmax(dim=1) correct += pred.eq(target).sum().item() if batch_idx % 100 == 0: print(f'Epoch {epoch} [{batch_idx*len(data)}/{len(train_loader.dataset)}] ' f'Loss: {loss.item():.6f}') acc = 100. * correct / len(train_loader.dataset) print(f'====> Epoch {epoch} Average loss: {total_loss/len(train_loader):.4f}, Acc: {acc:.2f}%') return total_loss/len(train_loader), acc 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) data = data.view(data.size(0), -1) logits = model(data) # standard forward without error input test_loss += F.cross_entropy(logits, target, reduction='sum').item() pred = logits.argmax(dim=1) correct += pred.eq(target).sum().item() test_loss /= len(test_loader.dataset) acc = 100. * correct / len(test_loader.dataset) print(f'Test set: Avg loss: {test_loss:.4f}, Acc: {acc:.2f}%\n') return test_loss, acc def main(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 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=64, shuffle=True) test_loader = DataLoader(datasets.MNIST('../data', train=False, transform=transform), batch_size=1000, shuffle=False) model = IterativeFluxMLP_AR(num_refinements=3, work_per_step=0.05, anti_beta=0.9, anti_gamma=0.5).to(device) optimizer = optim.Adam(model.parameters(), lr=0.001) for epoch in range(1, 6): train(model, device, train_loader, optimizer, epoch) test(model, device, test_loader) if __name__ == "__main__": main()