Open
Conversation
tinker495
force-pushed
the
feature/add-normuon
branch
from
85ef69e to
a4080cf
Compare
This commit introduces the NorMuon optimizer, which enhances the Muon algorithm by incorporating neuron-wise normalization of update magnitudes. The implementation includes the core functionality in `_normuon.py`, along with a dedicated test suite in `_normuon_test.py` to validate its behavior. Additionally, references to the new optimizer have been added in the documentation and the main `__init__.py` file. Key changes: - New optimizer: NorMuon with associated scaling and state management. - Tests for NorMuon functionality and behavior. - Documentation updates to include NorMuon in the API reference. PiperOrigin-RevId: [insert-rev-id-here]
tinker495
force-pushed
the
feature/add-normuon
branch
from
a4080cf to
a312387
Compare
Collaborator
|
Hey, we really appreciate the contribution! This is a lot of code which might be difficult to maintain with our small team of maintainers and so this might take us a while to merge. I'll try to get to this as soon as possible, thanks! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
normuonandscale_by_normuontooptax.contrib, implementing NorMuon (Neuron-wise Normalized Muon) (Li et al., 2025)combine.partitionChanges
optax/contrib/_normuon.pyscale_by_normuon(core transform) andnormuon(full optimizer with AdamW fallback)optax/contrib/_normuon_test.pyoptax/contrib/__init__.pynormuon,scale_by_normuon,NorMuonStateoptax/contrib/_common_test.pynormuonto shared contrib optimizer testsdocs/api/contrib.rstDesign
orthogonalize_via_newton_schulzandMuonDimensionNumbersfor consistencyvis stored as(batch, output)shape, matching the paper's Algorithm 1weight_dimension_numbersfor non-2D tensors (same interface as Muon)Test plan
pytest optax/contrib/_normuon_test.py— dedicated unit testspytest optax/contrib/_common_test.py— shared contrib optimizer testsrufflint check passes