Demystifying Manifold Constraints in LLM Pre-training
Abstract
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 stability and acceleration of LLM training. Consequently, several recent studies have suggested that weights should also be normalized, leading to a Riemannian optimization problem. While such constrained training frameworks demonstrate superior performance, the effects of explicitly constraining weights, and their interaction with existing stabilization mechanisms, remain less understood. To bridge this gap, we study manifold constrained training dynamics through activation scales, rotational dynamics, and the update-to-weight ratio. We propose a Riemannian spectral steepest descent optimizer called MACRO, alongside a radius selection principle to serve as our testbed. Our analysis and numerical experiments reveal that RMSNorm and manifold constraints serve overlapping roles, and that weight decay can be completely eliminated when manifold constraints are applied. By controlling the update-to-weight ratio, constrained training significantly alleviates update cancellation, empirically demonstrating that MACRO is robust to low-precision computation and competitive with existing algorithms for standard LLM pre-training.
Related
- Proximal-Linearized NEP Method for Composite Optimization with Set-Inclusion and Manifold Constraints
- Projected Primal–Dual Gradient Preference Optimization Under Category-Wise Split-Conformal Safety Constraints for Large Language Model Post-Training
- Demystifying Prediction-Powered Inference.
- Linearity of Graph Manifold Groups
- $3$-manifold groups are linear
BibTeX
@misc{an2026demystifying,
title = {{Demystifying Manifold Constraints in LLM Pre-training}},
author = {Kang An and Jiaxiang Li and Donald Goldfarb and Shiqian Ma},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2605.04418},
eprint = {2605.04418},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2605.04418},
}
RIS
TY - UNPB
TI - Demystifying Manifold Constraints in LLM Pre-training
AU - Kang An
AU - Jiaxiang Li
AU - Donald Goldfarb
AU - Shiqian Ma
PY - 2026
DA - 2026/09/28
PB - arXiv
DO - 10.48550/arxiv.2605.04418
UR - https://arxiv.org/abs/2605.04418
AB - 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 stability and acceleration of LLM training. Consequently, several recent studies have suggested that weights should also be normalized, leading to a Riemannian optimization problem. While such constrained training frameworks demonstrate superior performance, the effects of explicitly constraining weights, and their interaction with existing stabilization mechanisms, remain less understood. To bridge this gap, we study manifold constrained training dynamics through activation scales, rotational dynamics, and the update-to-weight ratio. We propose a Riemannian spectral steepest descent optimizer called MACRO, alongside a radius selection principle to serve as our testbed. Our analysis and numerical experiments reveal that RMSNorm and manifold constraints serve overlapping roles, and that weight decay can be completely eliminated when manifold constraints are applied. By controlling the update-to-weight ratio, constrained training significantly alleviates update cancellation, empirically demonstrating that MACRO is robust to low-precision computation and competitive with existing algorithms for standard LLM pre-training.
ER -