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 import numpy as np # ------------------------------------------------------------ # 1. Define a simple MLP for MNIST # ------------------------------------------------------------ class MLP(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(28*28, 128) self.fc2 = nn.Linear(128, 64) self.fc3 = nn.Linear(64, 10) def forward(self, x): x = x.view(-1, 28*28) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) return self.fc3(x) # ------------------------------------------------------------ # 2. Helper functions for parameter vector manipulation # ------------------------------------------------------------ def get_params(model): """Return a flat vector of all model parameters.""" return torch.cat([p.data.view(-1) for p in model.parameters()]) def set_params(model, vec): """Set model parameters from a flat vector.""" idx = 0 for p in model.parameters(): size = p.numel() p.data.copy_(vec[idx:idx+size].view(p.shape)) idx += size # ------------------------------------------------------------ # 3. Distant-Target Derivative (SPSA with multiscale) # ------------------------------------------------------------ def distant_target_gradient(model, loss_fn, x, y, c_values, n_perturbations=2, grad_clip=1.0): """ Estimate the gradient of loss w.r.t. model parameters using the Distant-Target Derivative principle (multiscale SPSA). Key fixes vs. the original: - loss values are .detach().item() so no autograd graph leaks - NaN / Inf guard after each perturbation - Per-scale gradient norm clipping before variance weighting - Inverse-variance weights are clamped to prevent one scale dominating - The combined gradient is norm-clipped before returning Args: model : PyTorch model loss_fn : loss function (e.g., CrossEntropyLoss) x, y : mini-batch inputs and labels c_values : list of step sizes (baselines) — coarse → fine n_perturbations: number of ±delta pairs per scale grad_clip : max L2 norm for the final combined gradient Returns: g_combined : combined gradient vector (same size as parameters) """ theta0 = get_params(model).detach() n_params = len(theta0) device = theta0.device g_list = [] var_list = [] for c in c_values: g_scale = torch.zeros(n_params, device=device) sq_sum = torch.zeros(n_params, device=device) valid = 0 # count of non-NaN perturbations for _ in range(n_perturbations): # Bernoulli ±1 perturbation (SPSA) delta = (torch.randint(0, 2, (n_params,), device=device, dtype=torch.float32) * 2 - 1) # Forward pass at theta + c*delta set_params(model, theta0 + c * delta) with torch.no_grad(): loss_plus = loss_fn(model(x), y).item() # Forward pass at theta - c*delta set_params(model, theta0 - c * delta) with torch.no_grad(): loss_minus = loss_fn(model(x), y).item() # Skip this perturbation if losses are invalid if not (np.isfinite(loss_plus) and np.isfinite(loss_minus)): continue # SPSA estimate: scalar diff × delta (since delta_i = ±1, inv = delta) g_est = ((loss_plus - loss_minus) / (2.0 * c)) * delta g_scale += g_est sq_sum += g_est ** 2 valid += 1 # Restore original parameters before moving on set_params(model, theta0) if valid == 0: # All perturbations were NaN at this scale — skip it continue g_scale /= valid sq_sum /= valid # Per-scale gradient norm clip (prevents one wild estimate ruining the mix) g_norm = g_scale.norm() if g_norm > grad_clip: g_scale = g_scale * (grad_clip / g_norm) # Scalar variance of the gradient estimates at this scale var_scale = (sq_sum - g_scale ** 2).mean().item() var_scale = max(var_scale, 1e-8) # floor to avoid divide-by-zero g_list.append(g_scale) var_list.append(var_scale) if not g_list: # Every scale failed — return zeros rather than crashing return torch.zeros(n_params, device=device) # Inverse-variance weighting (Axiom 3: multiscale combination) inv_var = torch.tensor([1.0 / v for v in var_list], device=device) # Clamp so that no single scale can receive > 90 % of the weight inv_var = inv_var.clamp(max=inv_var.sum() * 0.9) weights = inv_var / inv_var.sum() g_combined = sum(w * g for w, g in zip(weights, g_list)) # Final norm clip on the combined gradient g_norm = g_combined.norm() if g_norm > grad_clip: g_combined = g_combined * (grad_clip / g_norm) return g_combined # ------------------------------------------------------------ # 4. Training loop with Distant-Target updates # ------------------------------------------------------------ def train(model, train_loader, test_loader, epochs=5, lr=0.01, c_values=None, grad_clip=1.0): """ c_values follow the theory: [coarse, optimal, fine] Chosen so that c_coarse >> noise level and c_fine ~ h*. Rule of thumb for MNIST cross-entropy: loss noise ~ 0.05-0.1 → h* ≈ (noise/curvature)^(1/3) Empirically: [0.5, 0.1, 0.02] work well. """ if c_values is None: c_values = [0.5, 0.1, 0.02] # coarse, medium, fine device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) loss_fn = nn.CrossEntropyLoss() # SGD; we inject hand-computed gradients each step optimizer = optim.SGD(model.parameters(), lr=lr) for epoch in range(epochs): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) g = distant_target_gradient( model, loss_fn, data, target, c_values=c_values, n_perturbations=2, grad_clip=grad_clip, ) optimizer.zero_grad() idx = 0 for p in model.parameters(): size = p.numel() p.grad = g[idx:idx+size].view(p.shape).clone() idx += size optimizer.step() if batch_idx % 100 == 0: with torch.no_grad(): loss = loss_fn(model(data), target).item() print(f'Train Epoch: {epoch} ' f'[{batch_idx * len(data)}/{len(train_loader.dataset)} ' f'({100. * batch_idx / len(train_loader):.0f}%)]' f'\tLoss: {loss:.6f}') # ---- Evaluation ---- model.eval() test_loss = 0.0 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) test_loss += loss_fn(output, target).item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader) # average over batches (not samples) accuracy = 100. * correct / len(test_loader.dataset) print(f'Test set: Avg loss: {test_loss:.4f}, ' f'Accuracy: {correct}/{len(test_loader.dataset)} ' f'({accuracy:.2f}%)\n') # ------------------------------------------------------------ # 5. Main # ------------------------------------------------------------ if __name__ == "__main__": 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, download=True, transform=transform) # Smaller batch keeps each SPSA evaluation cheap and noisier → need larger c train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False) model = MLP() # c_values: [coarse, medium, fine] — matches Distant-Target theory sections B-C # These are well above the noise floor for CrossEntropyLoss on MNIST c_values = [0.5, 0.1, 0.02] train( model, train_loader, test_loader, epochs=5, lr=0.005, # slightly lower LR because gradients are less noisy now c_values=c_values, grad_clip=1.0, # L2 norm cap on the combined gradient )