Better LMO-based Momentum Methods with Second-Order Information

The use of momentum in stochastic optimization algorithms has shown empirical success across a range of machine learning tasks. Recently, a new class of stochastic momentum algorithms has emerged within the Linear Minimization Oracle (LMO) framework–leading to state-of-the-art methods, such as Muon, Scion, and Gluon, that effectively solve deep neural network training problems. However, traditional stochastic momentum methods offer convergence guarantees no better than the ${O}(1/K^{1/4})$ rate. While several approaches–such as Hessian-Corrected Momentum (HCM)–have aimed to improve this rate, their theoretical results are generally restricted to the Euclidean norm setting. This limitation hinders their applicability in problems, where arbitrary norms are often required. In this paper, we extend the LMO-based framework by integrating HCM, and provide convergence guarantees under relaxed smoothness and arbitrary norm settings. We establish improved convergence rates of ${O}(1/K^{1/3})$ for HCM, which can adapt to the geometry of the problem and achieve a faster rate than traditional momentum. Experimental results on training Multi-Layer Perceptrons (MLPs) and Long Short-Term Memory (LSTM) networks verify our theoretical observations.

💡 Research Summary

This paper presents a significant advancement in stochastic optimization by integrating second-order information into momentum-based methods within the Linear Minimization Oracle (LMO) framework. The primary goal is to overcome the limitations of existing fast-converging momentum variants, such as Hessian-Corrected Momentum (HCM), whose theoretical guarantees have been largely confined to the Euclidean norm and standard smoothness assumptions. These limitations restrict their applicability to modern machine learning problems, which often involve arbitrary norms (e.g., in normalized or sign-based gradient methods) and violate classical Lipschitz smoothness conditions.

The authors’ key contribution is the generalization of two HCM variants (from Salehkalaeybar et al. and Tran & Cutkosky) to the LMO-based algorithmic template. The LMO framework, which updates parameters via (x_{k+1} = x_k + \eta_k \cdot \text{LMO}(m_k)), is known for encapsulating optimizers like Muon and Scion and for its flexibility with arbitrary norms. By embedding HCM-style updates—which use stochastic Hessian-vector products to correct the momentum term and reduce gradient estimator variance—into this framework, the authors create a new class of algorithms that marry geometric flexibility with accelerated convergence.

Theoretical analysis forms the core of the work. The authors prove that under relaxed smoothness assumptions—namely ((L_0, L_1))-smoothness (for the gradient) and ((M_0, M_1))-smoothness (for the Hessian)—the proposed LMO-based methods with second-order momentum achieve a convergence rate of (O(1/K^{1/3})) in terms of the expected gradient norm for non-convex problems. This rate improves upon the (O(1/K^{1/4})) rate of standard LMO-momentum methods (e.g., Gluon) and the (O(1/K^{2/7})) rate of LMO with extrapolated momentum (IGT). Notably, the (O(1/K^{1/3})) rate matches the known lower bound and is thus optimal under the considered settings. The analysis carefully handles the interplay between the Euclidean norm (used for variance bounds in Assumption 1) and the arbitrary primal/dual norms (used for smoothness in Assumptions 3 & 4) via norm equivalence constants.

A comprehensive comparison table (Table 1 in the paper) situates this work within the landscape of LMO-based methods, clearly highlighting the improvements in convergence rate and the generality of the assumptions. To validate the theory, the authors conduct experiments on training Multi-Layer Perceptrons (MLPs) and Long Short-Term Memory (LSTM) networks. The empirical results demonstrate that the proposed second-order momentum methods converge faster and often to better final test accuracies compared to their first-order momentum counterparts, confirming the practical utility of the theoretical advancements.

In summary, this paper successfully bridges a gap between high-performance second-order momentum techniques and the flexible LMO optimization framework. It provides both theoretical guarantees for improved convergence under more realistic and general conditions, and empirical evidence of effectiveness on practical deep learning tasks, offering a powerful and broadly applicable optimization tool.