MuonQ: Enhancing Low-Bit Muon Quantization via Directional Fidelity Optimization

TL;DR

MuonQ introduces a low-bit quantization framework for Muon's optimizer, employing directional fidelity techniques to maintain training stability and accuracy while significantly reducing memory usage.

Contribution

It presents a novel 4-bit quantization method for Muon optimizer states using normalization, structural decomposition, and $rac12;$-law companding, enabling efficient large language model training.

Findings

MuonQ achieves stable 4-bit quantization of Muon optimizer states.

Pre-training results show MuonQ matches full-precision Muon in loss and accuracy.

Optimizer state memory is reduced by up to 7.3 times.

Abstract

The Muon optimizer has emerged as a compelling alternative to Adam for training large language models, achieving remarkable computational savings through gradient orthogonalization. However, Muon's optimizer state is more sensitive to quantization errors: because the orthogonalization discards the magnitudes of singular values and retains only directional information, even small quantization errors in singular vector directions are amplified in the update. In this work, we propose MuonQ, a low-bit Muon training framework built on the principle of directional fidelity optimization. First, we apply a pre-quantization normalization so that each step introduces quantization errors of the same magnitude, preventing the accumulated error from developing a preferred direction. Second, we introduce a structural decomposition that separately quantizes the dominant singular components via power iteration, ensuring that quantization errors perturb only singular value magnitudes rather than rotating singular vector directions. Third, we adopt $\mu$-law companding quantization to allocate higher resolution to densely packed momentum values, shifting the quantization objective from outlier preservation to dense-region distinguishability. Together, these techniques enable stable 4-bit quantization of Muon's optimizer states. Pre-training experiments on GPT-style and LLaMA-style models demonstrate that MuonQ at 4-bit precision closely matches full-precision Muon in both training loss and downstream task accuracy, while reducing optimizer state memory by up to 7.3 $\times$. Our code is available at https://github.com/YupengSu/MuonQ.