Reliable optimization is central to neural network (NN) training, yet Adam, the default optimizer for modern LLMs, rests on a fragile foundation. This thesis develops a principled grounding for Adam and motivates new designs. First, we revisit Adam's divergence--convergence debate and show the existence of a problem-dependent phase transition: with properly chosen, batch-size-dependent hyperparameters, Adam converges, whereas under small-$β_2$ regimes it can diverge. Second, we investigate why Adam substantially outperforms SGD on Transformers through Hessian structure. We find that the Hessian evolves toward a near-block-diagonal form along training, accompanied by strong block heterogeneity. We prove that this structure makes Adam's diagonal preconditioner effective. We further show that this special Hessian structure originates from consecutive multiplications of large matrix variables, and we provide a rigorous analysis based on random matrix theory. Finally, these insights motivate Adam-mini, a new optimizer that reduces Adam's memory footprint by 50\% while preserving its performance. Our results also have broader implications beyond Adam: they reveal new local structures in matrix-based nonconvex problems, and also help understand and improve recent NN optimizers, such as Muon.
We introduce the Deep Second-Order Stochastic Residual Method (D2SRM) for high-dimensional, Hessian-dependent fully nonlinear parabolic PDEs. A single scalar space--time network generates derivative-consistent approximations of the solution, gradient, and Hessian, which are trained jointly through second-order Brownian one-step residuals and terminal value and gradient penalties. For globally Lipschitz equations with identity diffusion and sufficiently weak Hessian coupling, we establish well-posedness in a Brownian occupation space and develop a population-level convergence theory. Under additional regularity, an a posteriori estimate bounds the squared full-jet occupation error of any admissible candidate by the time step and its population objective. For approximate population minimizers, the error bound separates time discretization, neural approximation, and population suboptimality; when the latter two terms are $O(h)$, the full-jet occupation norm is $O(h^{1/2})$. Experiments on a 100-dimensional manufactured benchmark compare terminal treatments, probe Hessian couplings inside and outside the proved small-gain range, and show decreasing errors as the time step decreases. The code is available at https://github.com/ZZHPKU/D2SRM.
Changxiao Cai, Yuchen Jiao, Gen Listat.ML cs.LG math.ST
Diffusion models are known to exploit unknown low-dimensional structure to accelerate sampling. However, existing convergence theory under low-dimensional data structure has largely focused on update rules with narrowly prescribed coefficient choices. This raises a fundamental question: is adaptation to low-dimensional structure sensitive to the precise choice of update coefficients? In this paper, we show that such adaptation is a robust property of diffusion models. For a broad class of update coefficients, we prove that $\widetilde{O}(k/\varepsilon)$ iterations suffice to generate an $\varepsilon$-accurate sample in total variation (TV) distance, independently of the ambient dimension. Our framework substantially broadens the class of diffusion samplers known to enjoy low dimensional adaptation and applies to several commonly used methods in practice. These results provide a theoretical justification for the empirical effectiveness of diffusion samplers across different coefficient choices when applied to structured, high-dimensional data.