1 citations · 2 across the 13 of their papers we have counts for
13 papers
Towards joint scaling laws with optimal batch size schedules
Jiaxiang Li, Zhiqi Bu, Shiyun Xu
Modern deep learning typically keeps the batch size static throughout training, thus overlooking the joint effect of learning rate and batch size on the training dynamics. In this…
Why Muon Outperforms Adam: A Curvature Perspective
Shuche Wang, Fengzhuo Zhang, Jiaxiang Li +2
Muon improves training efficiency over Adam in large language-model training by about two times, but the local geometric source of this advantage remains unclear. Our work takes a…
Demystifying Manifold Constraints in LLM Pre-training
Kang An, Jiaxiang Li, Donald Goldfarb +1
The recent success of matrix optimizers (e.g., Muon) suggests that specific normalization of momentum, such as orthogonalization and row-wise normalization, benefits both the stabi…
Federated Learning on Riemannian Manifolds: A Gradient-Free Projection-Based Approach
Hongye Wang, Zhaoye Pan, Chang He +2
Federated learning (FL) has emerged as a powerful paradigm for collaborative model training across distributed clients while preserving data privacy. However, existing FL algorithm…
On Relatively Smooth Optimization over Riemannian Manifolds
Chang He, Jiaxiang Li, Bo Jiang +2
We study optimization over Riemannian embedded submanifolds, where the objective function is relatively smooth in the ambient Euclidean space. Such problems have broad applications…
SLTrain: a sparse plus low-rank approach for parameter and memory efficient pretraining
Andi Han, Jiaxiang Li, Wei Huang +4
Large language models (LLMs) have shown impressive capabilities across various tasks. However, training LLMs from scratch requires significant computational power and extensive mem…
Getting More Juice Out of the SFT Data: Reward Learning from Human Demonstration Improves SFT for LLM Alignment
Jiaxiang Li, Siliang Zeng, Hoi-To Wai +3
Aligning human preference and value is an important requirement for contemporary foundation models. State-of-the-art techniques such as Reinforcement Learning from Human Feedback (…
Problem-Parameter-Free Decentralized Nonconvex Stochastic Optimization
Jiaxiang Li, Xuxing Chen, Shiqian Ma +1
Existing decentralized algorithms usually require knowledge of problem parameters for updating local iterates. For example, the hyperparameters (such as learning rate) usually requ…
Riemannian Bilevel Optimization
Jiaxiang Li, Shiqian Ma
In this work, we consider the bilevel optimization problem on Riemannian manifolds. We inspect the calculation of the hypergradient of such problems on general manifolds and thus e…
Zeroth-order Riemannian Averaging Stochastic Approximation Algorithms
Jiaxiang Li, Krishnakumar Balasubramanian, Shiqian Ma
We present Zeroth-order Riemannian Averaging Stochastic Approximation (\texttt{Zo-RASA}) algorithms for stochastic optimization on Riemannian manifolds. We show that \texttt{Zo-RAS…
A Riemannian ADMM
Jiaxiang Li, Shiqian Ma, Tejes Srivastava
We consider a class of Riemannian optimization problems where the objective is the sum of a smooth function and a nonsmooth function, considered in the ambient space. This class of…
Federated Learning on Riemannian Manifolds
Jiaxiang Li, Shiqian Ma
Federated learning (FL) has found many important applications in smart-phone-APP based machine learning applications. Although many algorithms have been studied for FL, to the best…
Stochastic Zeroth-order Riemannian Derivative Estimation and Optimization
Jiaxiang Li, Krishnakumar Balasubramanian, Shiqian Ma
We consider stochastic zeroth-order optimization over Riemannian submanifolds embedded in Euclidean space, where the task is to solve Riemannian optimization problem with only nois…