""" MNIST classifier in ODE-LANG (PyTorch + torchdiffeq), with addition: latent-flow consistency loss. Loss per batch: L_real = CE( neural_ode(x), y ) -- standard L_alt = CE( neural_ode( encoder(x) + ε ), y ) -- 'different seed' L = L_real + λ · L_alt The second term trains the *flow*: starting from many 'seeds' that are all plausibly h(0)'s of the same sample, the ODE must still classify them as y. So the latent space becomes basin-structured — exactly the picture Rheo draws: regions in latent space flow to the same attractor. """ import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms from torchdiffeq import odeint_adjoint as odeint # ---------- the declared vector field dh/dt = f(h, t) ---------- class ODEFunc(nn.Module): """ dh/dt = f(h, t). A small MLP with a ReZero-style residual gate. Initially zero so dh/dt = 0 at start; the ODE is the identity. """ def __init__(self, hidden_dim: int = 64): super().__init__() self.net = nn.Sequential( nn.Linear(hidden_dim, hidden_dim), nn.Tanh(), nn.Linear(hidden_dim, hidden_dim), nn.Tanh(), nn.Linear(hidden_dim, hidden_dim), ) self.gate = nn.Parameter(torch.zeros(1)) # ReZero def forward(self, t: torch.Tensor, h: torch.Tensor) -> torch.Tensor: return self.gate * self.net(h) # ---------- encoder: initial condition h(0) = encoder(image) ---------- class Encoder(nn.Module): def __init__(self, hidden_dim: int = 64): super().__init__() self.conv = nn.Sequential( nn.Conv2d(1, 16, 3, stride=2, padding=1), # 28 -> 14 nn.GroupNorm(8, 16), nn.SiLU(), nn.Conv2d(16, 32, 3, stride=2, padding=1), # 14 -> 7 nn.GroupNorm(8, 32), nn.SiLU(), nn.Conv2d(32, 64, 3, stride=2, padding=1), # 7 -> 4 nn.GroupNorm(8, 64), nn.SiLU(), ) self.fc = nn.Linear(64 * 4 * 4, hidden_dim) def forward(self, x: torch.Tensor) -> torch.Tensor: h = self.conv(x).flatten(1) return self.fc(h) # ---------- readout: precipitate h(1) ---------- class Readout(nn.Module): def __init__(self, hidden_dim: int = 64, num_classes: int = 10): super().__init__() self.norm = nn.LayerNorm(hidden_dim) self.fc = nn.Linear(hidden_dim, num_classes) def forward(self, hT: torch.Tensor) -> torch.Tensor: return self.fc(self.norm(hT)) # ---------- the full ODE-language program ---------- class NeuralODEClassifier(nn.Module): def __init__(self, hidden_dim: int = 64, num_classes: int = 10): super().__init__() self.encoder = Encoder(hidden_dim) self.func = ODEFunc(hidden_dim) self.readout = Readout(hidden_dim, num_classes) self.hidden_dim = hidden_dim def encode_state(self, x: torch.Tensor) -> torch.Tensor: return self.encoder(x) def flow(self, h0: torch.Tensor) -> torch.Tensor: """evolve 0..1 with integrator = dopri8, abs_tol = 1e-4, rel_tol = 1e-3""" t = torch.tensor([0.0, 1.0], device=h0.device) traj = odeint( self.func, h0, t, method="dopri8", atol=1e-4, rtol=1e-3, adjoint_params=tuple(self.func.parameters()), ) return traj[-1] def forward(self, x: torch.Tensor) -> torch.Tensor: h0 = self.encode_state(x) hT = self.flow(h0) return self.readout(hT) def forward_from_state(self, h0: torch.Tensor) -> torch.Tensor: """For alt-seed loss: classify from a given latent seed, not from an image.""" hT = self.flow(h0) return self.readout(hT) # ---------- training loop with the latent-flow consistency loss ---------- def train(model, loader, opt, device, epoch, latent_noise_std: float = 0.25, latent_aug_weight: float = 0.5): """ Train with two losses: CE on the actual image x CE on the trajectory evolved from a perturbed seed h(0)+ε Both must agree with y. """ model.train() for i, (x, y) in enumerate(loader): x, y = x.to(device), y.to(device) for j in range(3): opt.zero_grad(set_to_none=True) # (1) CE on the original sample x. logits_real = model(x) loss_real = F.cross_entropy(logits_real, y) # (2) Latent-flow consistency: # pick a different seed h̃(0) = h(0) + ε, then ODE-integrate it, # demand the readout agrees with y. h0 = model.encoder(x) h0_tilde = h0 + latent_noise_std * torch.randn_like(h0) logits_alt = model.forward_from_state(h0_tilde) loss_alt = F.cross_entropy(logits_alt, y) loss = loss_real + latent_aug_weight * loss_alt loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() if i % 10 == 0: print(f"epoch {epoch} step {i:4d}/{len(loader)} " f"loss_real {loss_real.item():.4f} " f"loss_alt_seed {loss_alt.item():.4f}") @torch.no_grad() def evaluate(model, loader, device): model.eval() correct, total = 0, 0 for x, y in loader: x, y = x.to(device), y.to(device) pred = model(x).argmax(dim=1) correct += (pred == y).sum().item() total += y.numel() return correct / total @torch.no_grad() def demo_latent_flow_consistency(model, loader, device, n_samples: int = 6, n_seeds: int = 16, perturbation_std: float = 0.40): """ For each sample in the test set, draw many perturbed seeds h~0 = h0 + ε in latent space, evolve them through the ODE, and check whether they collapse to the same class. This is the diagnostic the user asked for: 'what would the sample look like given it evolved from a different seed?' """ model.eval() x, y = next(iter(loader)) x, y = x[:n_samples].to(device), y[:n_samples].to(device) h0 = model.encoder(x) # canonical seeds eps = perturbation_std * torch.randn(n_seeds, *h0.shape, device=device) h0_alt = h0.unsqueeze(0) + eps # [S, B, H] t = torch.tensor([0.0, 1.0], device=device) h1_alt = odeint( model.func, h0_alt, t, method="dopri8", rtol=1e-3, atol=1e-4, )[-1] # [S, B, H] preds = model.readout(h1_alt).argmax(dim=-1) # [S, B] probs = model.readout(h1_alt).softmax(dim=-1) # [S, B, 10] print(f"\nTrue labels : {y.tolist()}") print(f"Predictions over {n_seeds} perturbed seeds, for each of {n_samples} samples:") for s in range(n_samples): agreement = (preds[:, s] == y[s]).float().mean().item() print(f" sample {s} (true={y[s].item()}) " f"agreement={agreement*100:.1f}% " f"top predictions per seed: {preds[:, s].tolist()}") # ---------- main ---------- def main(): device = "cuda" if torch.cuda.is_available() else "cpu" print(f"device: {device}") tfm = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)), ]) train_set = datasets.MNIST("../data", train=True, download=True, transform=tfm) test_set = datasets.MNIST("../data", train=False, download=True, transform=tfm) train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=2) test_loader = DataLoader(test_set, batch_size=512, shuffle=False, num_workers=2) model = NeuralODEClassifier(hidden_dim=64, num_classes=10).to(device) opt = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-4) for epoch in range(1, 4): train(model, train_loader, opt, device, epoch) acc = evaluate(model, test_loader, device) print(f"epoch {epoch} test accuracy: {acc*100:.2f}%") # The diagnostic the user asked for: print("\n=== latent flow consistency ===") demo_latent_flow_consistency(model, test_loader, device) if __name__ == "__main__": main()