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. MLP for MNIST (unchanged) # ------------------------------------------------------------ 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. Parameter vector helpers (unchanged) # ------------------------------------------------------------ def get_params(model): return torch.cat([p.data.view(-1) for p in model.parameters()]) def set_params(model, vec): 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 (multiscale SPSA) — CORRECTED # ------------------------------------------------------------ # KEY FIXES vs the original: # # (a) Removed ALL gradient norm clipping. # Original clip=1.0 on a 109 514-dim vector gave each # parameter ≤ 1/√109514 ≈ 0.003 gradient. With lr=0.005 # the update was ~1.5e-5 — too small to learn anything. # # (b) c_values are scaled by 1/√(n_params). # Original c=[0.5, 0.1, 0.02] produced perturbation # NORMS of 165, 33, 6.6 — 12× larger than the model's # own parameter norm (~13). Both L(θ+cΔ) and L(θ−cΔ) # collapsed to ~2.3 (random), so ΔL ≈ 0 and the gradient # estimate was pure noise. # Scaled c gives perturbation norms of ~1.0, 0.3, 0.1 # — small enough for the Taylor expansion to hold. # # (c) n_perturbations raised 2 → 5 (SNR improves √(5/2)≈1.6×). # # (d) Returns diagnostics (mean |ΔL| per scale) so you can # verify signal is above the noise floor. # def distant_target_gradient(model, loss_fn, x, y, c_values, n_perturbations=5): theta0 = get_params(model).detach() n_params = len(theta0) device = theta0.device g_list = [] var_list = [] diag_dL = [] # mean |ΔL| per scale (diagnostics) for c in c_values: g_scale = torch.zeros(n_params, device=device) sq_sum = torch.zeros(n_params, device=device) valid = 0 dL_abs = [] for _ in range(n_perturbations): delta = (torch.randint(0, 2, (n_params,), device=device, dtype=torch.float32) * 2 - 1).float() # Forward at θ + c·Δ set_params(model, theta0 + c * delta) with torch.no_grad(): loss_plus = loss_fn(model(x), y).item() # Forward at θ − c·Δ set_params(model, theta0 - c * delta) with torch.no_grad(): loss_minus = loss_fn(model(x), y).item() if not (np.isfinite(loss_plus) and np.isfinite(loss_minus)): continue dL = loss_plus - loss_minus dL_abs.append(abs(dL)) # SPSA estimate: (ΔL / 2c) · Δ (Δ_i = ±1 ⇒ Δ⁻¹ = Δ) g_est = (dL / (2.0 * c)) * delta g_scale += g_est sq_sum += g_est ** 2 valid += 1 # Restore original parameters set_params(model, theta0) if valid == 0: continue g_scale /= valid sq_sum /= valid # Scalar variance of gradient estimates at this scale var_scale = (sq_sum - g_scale ** 2).mean().item() var_scale = max(var_scale, 1e-12) g_list.append(g_scale) var_list.append(var_scale) diag_dL.append(np.mean(dL_abs) if dL_abs else 0.0) if not g_list: return torch.zeros(n_params, device=device), [0.0] * len(c_values) # Inverse-variance weighting (Axiom 3: multiscale combination) inv_var = torch.tensor([1.0 / v for v in var_list], device=device) inv_var = inv_var.clamp(max=inv_var.sum() * 0.9) # no scale > 90 % weights = inv_var / inv_var.sum() g_combined = sum(w * g for w, g in zip(weights, g_list)) # NOTE: no global norm clip — the per-parameter magnitude is # naturally ~|∇L|, which is what we want. return g_combined, diag_dL # ------------------------------------------------------------ # 4. Training loop — CORRECTED # ------------------------------------------------------------ def train(model, train_loader, test_loader, epochs=10, lr=0.01, momentum=0.9, c_values=None, n_perturbations=5, lr_decay=0.95): """ Corrections: • lr raised 0.005 → 0.01, momentum 0 → 0.9 (momentum averages noisy SPSA gradients over steps) • lr_decay: exponential schedule per epoch (SPSA theory recommends decreasing gain sequences) • Diagnostic prints: |g|, |ΔL| per scale """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) loss_fn = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum) for epoch in range(epochs): current_lr = lr * (lr_decay ** epoch) for pg in optimizer.param_groups: pg['lr'] = current_lr model.train() grad_norms = [] for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) g, dL_diag = distant_target_gradient( model, loss_fn, data, target, c_values=c_values, n_perturbations=n_perturbations, ) grad_norms.append(g.norm().item()) 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'Epoch {epoch} ' f'[{batch_idx*len(data)}/{len(train_loader.dataset)} ' f'({100.*batch_idx/len(train_loader):.0f}%)] ' f'lr={current_lr:.5f} loss={loss:.4f} ' f'|g|={g.norm().item():.1f} ' f'ΔL={[f"{d:.5f}" for d in dL_diag]}') print(f' → epoch {epoch} mean |g| = {np.mean(grad_norms):.1f}') # ---- Evaluation ---- model.eval() test_loss, correct = 0.0, 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) accuracy = 100. * correct / len(test_loader.dataset) print(f' Test: loss={test_loss:.4f} ' f'Acc={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) train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False) model = MLP() # ========================================================== # KEY FIX: scale c by 1/√(n_params) # ========================================================== # The theory's h* is the perturbation NORM, not the per- # parameter step. In n-dimensional space: # # ‖c·Δ‖ = c · √n ⇒ c = h / √n # # We choose h_norms = [1.0, 0.3, 0.1] (coarse → fine) # giving perturbation norms well below the model's # parameter norm (~13), so the Taylor expansion holds # and ΔL carries real gradient signal. # ========================================================== n_params = sum(p.numel() for p in model.parameters()) inv_sqrt_n = 1.0 / np.sqrt(n_params) h_norms = [1.0, 0.3, 0.1] # perturbation norms (theory's h) c_values = [h * inv_sqrt_n for h in h_norms] print(f"Parameters: {n_params}") print(f"1/sqrt(n): {inv_sqrt_n:.6f}") print(f"h_norms: {h_norms}") print(f"c_values: {[f'{c:.6f}' for c in c_values]}") print(f"Perturb norms: {[f'{c*np.sqrt(n_params):.4f}' for c in c_values]}") print(f"(param norm ~13, so ratios = " f"{[f'{c*np.sqrt(n_params)/13:.3f}' for c in c_values]})") print() train(model, train_loader, test_loader, epochs=10, lr=0.01, # was 0.005 — raised for usable step size momentum=0.9, # NEW — smooths noisy SPSA estimates c_values=c_values, n_perturbations=5, # was 2 — more averaging per scale lr_decay=0.95, # NEW — exponential LR schedule )