Rennala-NSGD: Asynchronous Stochastic Optimization Beyond Euclidean Geometry
Igor Sokolov ⋅ Alexander Tyurin ⋅ Peter Richtarik
Abstract
The recent empirical success of non-Euclidean optimizers, such as Muon (Jordan et al., 2024) and Shampoo (Gupta et al., 2018), has transformed the training of Large Language Models (LLMs) by exploiting the geometry of matrix spaces. Despite this progress, the theoretical understanding of stochastic optimization in general normed spaces remains limited, particularly in distributed settings. In this work, we present a novel analysis of Minibatch Stochastic Gradient Descent (SGD) in non-Euclidean geometries. Leveraging the framework of $\\kappa$-regular normed spaces (Juditsky & Nemirovski, 2008) and a light-tail noise model defined directly in the dual norm, we derive high-probability oracle complexity bounds for smooth non-convex optimization. Crucially, our analysis reveals that measuring stochastic variance within the intrinsic non-Euclidean geometry avoids the suboptimal dimension-dependent factors inherent in Euclidean norm equivalence. Building on this framework, we introduce Rennala-NSGD, a semi-asynchronous distributed algorithm, and establish, to the best of our knowledge, the first wall-clock time complexity guarantee for asynchronous non-Euclidean stochastic optimization.
Chat is not available.
Successful Page Load