MuonIO: Principled Norm-Aware Descent for Embedding Tables and Language Model Heads
The Muon optimizer derives its update rule for hidden linear layers by solving a local linearization of the loss penalized by the spectral norm, motivated by an RMS-stability argument for dense linear layers. Standard Muon implementations, however, exclude the input (embedding table) and output (language model head) la...
Lin-Kai Ma, Xin-Yu Luo, Meng-Bo Wang et al.
· 0 citations