Compute-Optimal Training as Stochastic Optimal Control

Choosing learning rate and batch size schedules for neural network training is typically performed via expensive grid search over heuristic families. We formulate compute-optimal training as a stochastic optimal control problem, where the state comprises per-eigenmode loss coefficients from the continuous-time scaling model of Bordelon et al. (2024) and the remaining compute budget, while the controls are the learning rate $\eta(t)$ and batch size $B(t)$. The resulting Hamilton-Jacobi-Bellman (HJB) equation characterizes the globally optimal schedule. We solve this HJB equation using the Deep Galerkin Method with policy iteration, a mesh-free approach that handles the 12-dimensional PDE. On the random feature model, the HJB framework yields two insights: (1) in the noise-dominated regime, a well-tuned constant schedule is near-optimal—the HJB solution explains this by showing that the value function's gradient structure leaves little room for improvement; (2) the optimal batch size grows monotonically, derived from first principles via the optimality conditions rather than empirical observation (Smith et al., 2018). Joint learning-rate-and-batch-size control provides 2–6% gains over learning-rate-only optimization, with the largest benefit at hard tasks where noise management is most impactful.