提出MuonQ框架,实现4比特量化下稳定训练大模型。
MuonQ: Enhancing Low-Bit Muon Quantization via Directional Fidelity Optimization

- 通过方向保真优化,防止量化误差累积导致方向偏移。
- 4比特量化下,训练损失和下游准确率接近全精度水平。
- 适合需要降低内存占用的大模型训练场景。
Muon优化器通过梯度正交化实现了大语言模型训练的显著计算节省,但其优化器状态对量化误差更敏感:因正交化丢弃了奇异值大小仅保留方向信息,微小的方向量化误差会在更新中被放大。本文提出MuonQ,一种基于方向保真优化的低比特训练框架。首先,采用预量化归一化,使每步量化误差幅度一致,避免误差积累产生偏好方向;其次,引入结构分解,通过幂迭代分别量化主奇异分量,确保量化误差仅影响奇异值大小而非旋转奇异向量方向;第三,采用μ律压缩量化,为密集动量值分配更高分辨率,将量化目标从异常值保持转向密集区域可区分性。这些技术共同实现穆翁优化器状态的稳定4比特量化。在GPT风格与LLaMA风格模型上的预训练实验表明,MuonQ在4比特精度下恢复了全精度Muon的大部分训练损失与下游准确率,同时优化器状态内存最高降低7.3倍,在精度与内存之间取得良好权衡。代码已开源。
原文摘要 · Abstract (English)
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 $μ$-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 recovers most of full-precision Muon's training loss and downstream accuracy while reducing optimizer-state memory by up to 7.3$\times$, offering a favorable point on the accuracy-memory trade-off. Our code is available at https://github.com/YupengSu/MuonQ.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。