""" MNIST classifier in ODE-LANG, implemented with PyTorch + torchdiffeq. Read of this file: Encoder -> 'state h(0) = encoder(image)' (initial condition) ODEFunc -> 'dh/dt = f(h, t)' (the declared vector field) odeint_adjoint-> 'evolve t in [0,1]' (the adaptive integrator) readout -> 'precipitate h(1)' (read off the final state) The "Rheo contract" — that the numerical trajectory h̃(t) satisfies |h̃(t) - h(t)| <= max(atol, rtol * |h(t)|) —is enforced by choosing method='dopri8' with rtol=1e-3 and atol=1e-4, the Dormand-Prince 8(7) embedded pair (free error estimate, free dense output). """ 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): """ The vector field the programmer declares. Inside Rheo this would be: dh/dt = mlp(h, t) Here in PyTorch it is a small MLP with a ReZero-style residual scale. ReZero (Bachlechner et al. 2020) initializes the gate to zero, so initially dh/dt = 0 and the ODE flow is the identity. This makes the model start as a stable linear classifier and only acquire flow as training proceeds — the same trick that lets Neural ODEs train stably. """ 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), ) # Rheo's "compile time constant" -> a torch.nn.Parameter here. self.gate = nn.Parameter(torch.zeros(1)) # ReZero init = 0 def forward(self, t: torch.Tensor, h: torch.Tensor) -> torch.Tensor: # nn.GroupNorm: Lipschitz stabilizer. Picard-Lindelöf needs Lipschitz f. return self.gate * self.net(h) # ---------- the encoder: initial condition h(0) = encoder(image) ---------- class Encoder(nn.Module): """ state h(0) = encoder(image) -- the initial condition of the IVP. """ 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) # ---------- the readout: precipitate h(1) ---------- class Readout(nn.Module): """ precipitate h(1) -> class logits """ 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): """ Whole program (Rheo form): state h ~ (initial condition) dh/dt = mlp(h,t) evolve t in [0,1] with integrator = dopri8, abs_tol = 1e-4, rel_tol = 1e-3 precipitate h(1) -> class """ 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 forward(self, x: torch.Tensor) -> torch.Tensor: # h(0) <- encoder(x) h0 = self.encoder(x) # 'evolve 0..1 with integrator = dopri8, abs_tol = 1e-4, rel_tol = 1e-3' t = torch.tensor([0.0, 1.0], device=x.device) # odeint_adjoint: same accuracy, O(1) memory — the "compiler's invariant" traj = odeint( self.func, h0, t, method="dopri8", atol=1e-4, rtol=1e-3, adjoint_params=tuple(self.func.parameters()), ) # traj shape: [T=2, batch, hidden_dim]; trajectory is the dense output. hT = traj[-1] # precipitate h(1) -> class logits return self.readout(hT) # ---------- training loop ---------- def train(model, loader, opt, device, epoch): model.train() for i, (x, y) in enumerate(loader): x, y = x.to(device), y.to(device) opt.zero_grad(set_to_none=True) logits = model(x) loss = F.cross_entropy(logits, y) loss.backward() # soft-clip the gradient into the vector field to keep f Lipschitz-ish torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() if i % 200 == 0: print(f"epoch {epoch} step {i:4d}/{len(loader)} loss {loss.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 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}%") if __name__ == "__main__": main()