Equilibrium-Matched Nonlinear Momentum Optimizer

Integrate this as a replacement for the momentum state in SGD, initially without Adam's adaptive preconditioner. Maintain parameters theta and a momentum vector u. Choose a positive componentwise friction such as gamma(u) = gamma_0(1 + a u^2), with gamma_0 > 0 and a >= 0. At each minibatch, compute g = grad_theta L(theta), evaluate gamma(u), evaluate gamma'(u) = 2 gamma_0 a u, draw independent standard-normal noise xi, and apply the displayed Euler-Maruyama update before theta <- theta + h u. Use the same temperature T and step size h for all comparisons. The paper supplies the Fokker-Planck equation and friction-diffusion matching rule; gamma_0, a, T, and h are engineering hyperparameters. The first cheap experiment should use a two-layer MLP on MNIST, comparing ordinary momentum SGD, constant-friction Langevin momentum, and the nonlinear matched optimizer at equal minibatch size and compute. First run with gradients set to zero and measure the momentum histogram, variance, autocorrelation, displacement mean-squared growth, and excess kurtosis. The quantitative prediction is that matched dynamics converge to variance T and approximately Gaussian stationary momentum, while a control with the same nonlinear friction but constant noise has a variance and kurtosis that depend on a. For parameter trajectories, the prediction is finite long-time diffusion with non-Gaussian intermediate displacement and excess kurtosis tending toward zero at long times. Sweep h and locate the divergence boundary; the measured stable region should shrink when the maximum effective friction or noise amplitude makes h gamma(u) large. Training should then test whether the transient heavy-tailed displacement improves escape from sharp minibatch-loss basins without violating the equilibrium diagnostics.