NuMuon: Nuclear-Norm-Constrained Muon for Compressible LLM Training

TL;DR

NuMuon introduces a nuclear-norm constraint to the Muon optimizer, enhancing the low-rank structure of LLM weights, which improves compressibility and maintains training quality at large scales.

Contribution

This work reveals that Muon-trained models naturally exhibit low-rank structures and proposes NuMuon to explicitly enforce this, leading to better compression without sacrificing convergence.

Findings

NuMuon increases model weight compressibility.

NuMuon improves post-compression model quality.

Models trained with NuMuon retain Muon's convergence benefits.

Abstract

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.