NuMuon: Nuclear-Norm-Constrained Muon for Compressible LLM Training
This addresses memory and cost constraints for deploying LLMs, but is incremental as it builds on existing Muon and compression methods.
The paper tackled the problem of memory and deployment costs in large language models (LLMs) by showing that Muon-trained models exhibit low-rank weight structures, and proposed NuMuon, which enhances compressibility and improves post-compression quality while maintaining convergence.
The rapid progress of large language models (LLMs) is increasingly constrained by memory and deployment costs, motivating compression methods for practical deployment. Many state-of-the-art compression pipelines leverage the low-rank structure of trained weight matrices, a phenomenon often associated with the properties of popular optimizers such as Adam. In this context, Muon is a recently proposed optimizer that improves LLM pretraining via full-rank update steps, but its induced weight-space structure has not been characterized yet. In this work, we report a surprising empirical finding: despite imposing full-rank updates, Muon-trained models exhibit pronounced low-rank structure in their weight matrices and are readily compressible under standard pipelines. Motivated by this insight, we propose NuMuon, which augments Muon with a nuclear-norm constraint on the update direction, further constraining the learned weights toward low-rank structure. Across billion-parameter-scale models, we show that NuMuon increases weight compressibility and improves post-compression model quality under state-of-the-art LLM compression pipelines while retaining Muon's favorable convergence behavior.