#!/usr/bin/env python3 """ XYFLOW: Boundary Flux Classifier for MNIST ========================================== Theory (from ODE-CCT framework): The "missing information" for 100% boundary accuracy is NOT more data points — it is the TRANSVERSE FLOW DYNAMICS: the directional info that tells the trajectory which way to fall off the boundary. Instead of learning x → label (discrete, discontinuous on boundary), we learn: S(x) : potential landscape with basins (attractors) per class F(x) : flow field (the MISSING INFORMATION — the true dynamics) Classification = sign of flux: flux = ∇(S_a - S_b) · F The boundary is a DYNAMICAL SEPARATRIX, not a static set of points. ODE uniqueness guarantees deterministic escaping → 100% accuracy on the boundary (modulo measure-zero non-hyperbolic equilibria). Architecture: PotentialField : 784 → 10 (one scalar potential S_k per class) FlowField : 784 → 784 (vector field F in input space) Loss = L_potential + λ_flux * L_flux + λ_align * L_align + λ_reg * ||F||² L_potential : CrossEntropy on -S (landscape shaping — basins per class) L_flux : Hinge on flux sign (the core "missing info" loss) For true class k, competitor j: flux = (∇S_k - ∇S_j) · F must be < 0 (toward basin k) L_align : ||F + ∇S_true||² (F ≈ gradient descent toward true basin) L_reg : ||F||² (prevent degenerate F = 0) Inference: Far from boundary : argmin_k S_k(x) (standard — deepest basin) Near boundary : sign( ∇(S_a-S_b) · F ) (flux resolves ambiguity) Usage: python xyflow_mnist.py --epochs 15 python xyflow_mnist.py --epochs 20 --lambda-flux 1.0 --boundary-threshold 2.0 """ import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchvision import datasets, transforms from torch.utils.data import DataLoader import numpy as np import time import argparse # ============================================================ # Model Components # ============================================================ class PotentialField(nn.Module): """ S_k(x): scalar potential for each class k. Lower S_k → more attracted to class k (deeper basin). Boundary between classes i, j: S_i(x) = S_j(x) (separatrix). """ def __init__(self, dim=784, num_classes=10, hidden=256): super().__init__() self.net = nn.Sequential( nn.Linear(dim, hidden), nn.LayerNorm(hidden), nn.SiLU(), nn.Linear(hidden, hidden), nn.LayerNorm(hidden), nn.SiLU(), nn.Linear(hidden, hidden), nn.LayerNorm(hidden), nn.SiLU(), nn.Linear(hidden, num_classes), ) def forward(self, x): return self.net(x) class FlowField(nn.Module): """ F(x): vector field in input space — the "missing information". This is the transverse flow that resolves boundary ambiguity. Direction tells the trajectory which basin to fall into. """ def __init__(self, dim=784, hidden=256): super().__init__() self.net = nn.Sequential( nn.Linear(dim, hidden), nn.LayerNorm(hidden), nn.SiLU(), nn.Linear(hidden, hidden), nn.LayerNorm(hidden), nn.SiLU(), nn.Linear(hidden, dim), ) def forward(self, x): return self.net(x) class XYFlowClassifier(nn.Module): """ Full XYFLOW boundary flux classifier. Training : learns S(x) and F(x) jointly. Inference: uses flux sign for boundary (hard) cases. """ def __init__(self, dim=784, num_classes=10, hidden=256): super().__init__() self.potential_net = PotentialField(dim, num_classes, hidden) self.flow_net = FlowField(dim, hidden) self.dim = dim self.num_classes = num_classes # --- forward helpers --- def potential(self, x): """S_k(x) for all k. Shape: (batch, num_classes).""" return self.potential_net(x) def flow(self, x): """F(x). Shape: (batch, dim).""" return self.flow_net(x) # --- flux loss (the core "missing information" loss) --- def compute_flux_loss(self, x, y, margin=1.0): """ For each sample with true class k and nearest competitor j: flux = (∇S_k - ∇S_j) · F want flux < 0 (flow toward basin k where S_k < S_j) loss = ReLU(flux + margin) Weighted by proximity to boundary (closer = more important). Uses only 2 backward passes (∇S_true, ∇S_competitor), not 10. """ batch_size = x.shape[0] device = x.device # Forward with grad enabled on input x_req = x.clone().requires_grad_(True) S = self.potential(x_req) # (batch, num_classes) # Nearest competitor: lowest S excluding the true class S_masked = S.clone() batch_idx = torch.arange(batch_size, device=device) S_masked[batch_idx, y] = float("inf") competitor = S_masked.argmin(dim=1) # (batch,) # ∇S_true (gradient of true-class potential w.r.t. input) grad_true = torch.autograd.grad( S[batch_idx, y].sum(), x_req, create_graph=True, retain_graph=True, )[0] # (batch, dim) # ∇S_competitor grad_comp = torch.autograd.grad( S[batch_idx, competitor].sum(), x_req, create_graph=True, retain_graph=True, )[0] # (batch, dim) # Flow field F_x = self.flow(x_req) # (batch, dim) # --- flux = (∇S_true - ∇S_comp) · F --- normal = grad_true - grad_comp # boundary normal flux = (normal * F_x).sum(dim=1) # (batch,) # Boundary distance: S_comp - S_true (positive = true class deeper) S_true = S[batch_idx, y] S_comp = S[batch_idx, competitor] boundary_dist = (S_comp - S_true).detach() # no grad through weight # Weight: higher for boundary samples (small |boundary_dist|) boundary_weight = torch.exp(-boundary_dist.abs() / 2.0) # Flux loss: want flux < 0 (toward true basin) flux_loss = (F.relu(flux + margin) * boundary_weight).mean() # Alignment loss: F ≈ -∇S_true (gradient descent toward true attractor) align_loss = (F_x + grad_true).pow(2).mean() # Regularization: prevent degenerate F = 0 flow_reg = F_x.pow(2).mean() return flux_loss, align_loss, flow_reg # --- inference --- def classify_standard(self, x): """Standard: argmin of potential (deepest basin).""" with torch.no_grad(): S = self.potential(x) preds = S.argmin(dim=1) return preds, S def classify_flux(self, x, boundary_threshold=1.0): """ Flux-enhanced classification. For points far from boundary: argmin S (standard). For points near boundary: use sign(flux) to resolve. flux > 0 → S_a - S_b increasing → moving toward b → switch to b flux < 0 → S_a - S_b decreasing → staying in a → keep a """ with torch.no_grad(): S = self.potential(x) # Top-2 lowest potentials (most attracted) top2 = S.topk(2, dim=1, largest=False) a = top2.indices[:, 0] # best (lowest S) b = top2.indices[:, 1] # competitor gap = top2.values[:, 1] - top2.values[:, 0] # margin preds = a.clone() # --- boundary samples: resolve with flux --- bnd_mask = gap < boundary_threshold n_bnd = bnd_mask.sum().item() if n_bnd > 0: bnd_idx = bnd_mask.nonzero(as_tuple=True)[0] x_b = x[bnd_idx] a_b = a[bnd_idx] b_b = b[bnd_idx] # Compute flux for boundary samples only with torch.enable_grad(): x_b_req = x_b.clone().requires_grad_(True) S_b = self.potential(x_b_req) ri = torch.arange(len(bnd_idx), device=x.device) grad_a = torch.autograd.grad( S_b[ri, a_b].sum(), x_b_req, create_graph=False, retain_graph=True, )[0] grad_b = torch.autograd.grad( S_b[ri, b_b].sum(), x_b_req, create_graph=False, )[0] F_b = self.flow(x_b_req) flux = ((grad_a - grad_b) * F_b).sum(dim=1) # flux > 0 → moving toward competitor b → switch switch = flux > 0 preds[bnd_idx[switch]] = b_b[switch] return preds, S, gap, n_bnd # ============================================================ # Training # ============================================================ def train(model, device, train_loader, optimizer, epoch, lambda_flux=0.5, lambda_align=0.1, lambda_reg=0.01, flux_margin=1.0, flux_interval=3, grad_clip=1.0): """ Dual-loss training: L = L_potential + λ_flux * L_flux + λ_align * L_align + λ_reg * ||F||² flux_interval: compute expensive flux loss every N batches. """ model.train() total_loss = 0.0 total_pot = 0.0 total_flux = 0.0 correct = 0 total = 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) # flatten 28×28 → 784 optimizer.zero_grad() # --- Potential loss (landscape shaping) --- S = model.potential(data) # Negate S: lower S = more attracted, CE expects higher = more likely pot_loss = F.cross_entropy(-S, target) # --- Flux loss (the missing information) --- if batch_idx % flux_interval == 0: flux_loss, align_loss, flow_reg = model.compute_flux_loss( data, target, margin=flux_margin ) loss = ( pot_loss + lambda_flux * flux_loss + lambda_align * align_loss + lambda_reg * flow_reg ) total_flux += flux_loss.item() else: loss = pot_loss loss.backward() if grad_clip: torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) optimizer.step() # --- metrics --- total_loss += loss.item() total_pot += pot_loss.item() preds = S.argmin(dim=1) correct += preds.eq(target).sum().item() total += len(target) if batch_idx % 200 == 0: print( f" Epoch {epoch} [{batch_idx * len(data):5d}/" f"{len(train_loader.dataset)}] " f"Loss {loss.item():.4f} Pot {pot_loss.item():.4f} " f"Acc {100. * correct / total:.2f}%" ) nb = len(train_loader) nf = max(1, nb // flux_interval) print( f" Epoch {epoch} AVG " f"Loss {total_loss / nb:.4f} " f"Pot {total_pot / nb:.4f} " f"Flux {total_flux / nf:.4f} " f"Acc {100. * correct / total:.2f}%" ) # ============================================================ # Testing # ============================================================ def test(model, device, test_loader, boundary_threshold=1.0, verbose=True): """ Compare standard vs. flux-enhanced classification. Report overall accuracy and boundary-specific accuracy. """ model.eval() std_correct = 0 flux_correct = 0 total = 0 bnd_std_correct = 0 bnd_flux_correct = 0 bnd_total = 0 all_gaps = [] for data, target in test_loader: data, target = data.to(device), target.to(device) data = data.view(data.size(0), -1) # Standard std_preds, S = model.classify_standard(data) std_correct += std_preds.eq(target).sum().item() # Flux-enhanced flux_preds, _, gap, n_bnd = model.classify_flux( data, boundary_threshold=boundary_threshold ) flux_correct += flux_preds.eq(target).sum().item() # Boundary-specific bnd_mask = gap < boundary_threshold bnd_total += n_bnd if n_bnd > 0: bnd_std_correct += std_preds[bnd_mask].eq(target[bnd_mask]).sum().item() bnd_flux_correct += flux_preds[bnd_mask].eq(target[bnd_mask]).sum().item() all_gaps.extend(gap.cpu().numpy()) total += len(target) std_acc = 100.0 * std_correct / total flux_acc = 100.0 * flux_correct / total bnd_pct = 100.0 * bnd_total / total if verbose: print(f"\n{'=' * 64}") print(f" Test Results (boundary_threshold = {boundary_threshold})") print(f"{'=' * 64}") print(f" Total samples: {total}") print(f" Standard accuracy: {std_acc:.2f}% ({std_correct}/{total})") print(f" Flux accuracy: {flux_acc:.2f}% ({flux_correct}/{total})") print(f" Improvement: {flux_acc - std_acc:+.2f}%") print(f" {'─' * 44}") print(f" Boundary samples: {bnd_total} ({bnd_pct:.1f}% of test set)") if bnd_total > 0: bs = 100.0 * bnd_std_correct / bnd_total bf = 100.0 * bnd_flux_correct / bnd_total print(f" Boundary std acc: {bs:.2f}% ({bnd_std_correct}/{bnd_total})") print(f" Boundary flux acc: {bf:.2f}% ({bnd_flux_correct}/{bnd_total})") print(f" Boundary improvement: {bf - bs:+.2f}%") print(f" {'─' * 44}") print( f" Gap stats: mean={np.mean(all_gaps):.3f} " f"median={np.median(all_gaps):.3f} " f"min={np.min(all_gaps):.3f}" ) print(f"{'=' * 64}\n") return std_acc, flux_acc # ============================================================ # Main # ============================================================ def main(): parser = argparse.ArgumentParser(description="XYFLOW Boundary Flux MNIST") parser.add_argument("--batch-size", type=int, default=128) parser.add_argument("--epochs", type=int, default=15) parser.add_argument("--lr", type=float, default=1e-3) parser.add_argument("--hidden", type=int, default=256) parser.add_argument("--lambda-flux", type=float, default=0.5, help="Weight for flux (missing-info) loss") parser.add_argument("--lambda-align", type=float, default=0.1, help="Weight for flow alignment loss") parser.add_argument("--lambda-reg", type=float, default=0.01, help="Weight for flow regularization") parser.add_argument("--flux-margin", type=float, default=1.0, help="Hinge margin for flux loss") parser.add_argument("--boundary-threshold", type=float, default=1.0, help="Potential gap below which a sample is 'boundary'") parser.add_argument("--flux-interval", type=int, default=3, help="Compute flux loss every N batches") parser.add_argument("--no-cuda", action="store_true") args = parser.parse_args() use_cuda = not args.no_cuda and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") print(f"Device: {device}") # ---- Data ---- transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)), ]) train_ds = datasets.MNIST("../data", train=True, download=True, transform=transform) test_ds = datasets.MNIST("../data", train=False, transform=transform) train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True) test_loader = DataLoader(test_ds, batch_size=256, shuffle=False) # ---- Model ---- model = XYFlowClassifier( dim=784, num_classes=10, hidden=args.hidden ).to(device) optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) n_pot = sum(p.numel() for p in model.potential_net.parameters()) n_flow = sum(p.numel() for p in model.flow_net.parameters()) print(f"\nParameters: {n_pot + n_flow:,} (potential {n_pot:,} + flow {n_flow:,})") print(f"λ_flux={args.lambda_flux} λ_align={args.lambda_align} " f"λ_reg={args.lambda_reg} margin={args.flux_margin}") print(f"boundary_threshold={args.boundary_threshold} " f"flux_interval={args.flux_interval} epochs={args.epochs}\n") # ---- Train ---- for epoch in range(1, args.epochs + 1): t0 = time.time() train( model, device, train_loader, optimizer, epoch, lambda_flux=args.lambda_flux, lambda_align=args.lambda_align, lambda_reg=args.lambda_reg, flux_margin=args.flux_margin, flux_interval=args.flux_interval, ) scheduler.step() if epoch % 3 == 0 or epoch == args.epochs: test(model, device, test_loader, boundary_threshold=args.boundary_threshold) print(f" Time: {time.time() - t0:.1f}s\n") # ---- Final evaluation: boundary threshold sweep ---- print("=" * 64) print(" FINAL EVALUATION — Boundary Threshold Sweep") print("=" * 64) for thresh in [0.5, 1.0, 2.0, 5.0, float("inf")]: test(model, device, test_loader, boundary_threshold=thresh) # ---- Save model ---- torch.save(model.state_dict(), "xyflow_mnist.pt") print("Model saved to xyflow_mnist.pt") print("Done.") if __name__ == "__main__": main()