contrib: add SCION optimizer (norm-constrained LMO generalisation of Muon) - #1718
Open
irhyl wants to merge 4 commits into
Open
contrib: add SCION optimizer (norm-constrained LMO generalisation of Muon)#1718irhyl wants to merge 4 commits into
irhyl wants to merge 4 commits into
Conversation
Gram Newton-Schulz (GramNS) avoids recomputing XX^T at every iteration by maintaining the Gram matrix R = XX^T and an accumulated orthogonal transform Q, both in n x n space. This cuts per-step cost from O(nm^2) to O(n^3) when the matrix aspect ratio is large. Key changes: - Add _POLAR_EXPRESS_NS_COEFFS (Amsel et al., 2505.16932) and a 'polar_express' entry in _NS_COEFFS_PRESET_DICT. - Add gram=False parameter to orthogonalize_via_newton_schulz, scale_by_muon, and muon. When gram=True the Gram NS path runs instead of the standard Newton-Schulz loop; the restart at step 2 materialises Q back into X to keep the Gram matrix condition number bounded (Amsel et al., 2025). - Support both scalar (fori_loop) and per-step (scan) coefficient formats in the Gram path, mirroring the existing NS structure. - Add four parameterised tests covering wide/tall/square matrices, preset resolution, end-to-end muon integration, and consistency between gram=True and gram=False with the same coefficients.
…Muon) Implements scale_by_scion and scion from Pethick et al. (arXiv:2502.07529). SCION selects update directions via a linear minimization oracle over a norm ball, unifying Muon (spectral LMO / polar factor) with sign updates (Lion, l-inf norm) and normalised-gradient descent (Frobenius norm). With lmo='auto', parameters with ndim >= 2 are routed to the spectral LMO via Newton-Schulz and 1D parameters use the sign LMO, using combine.partition in the same pattern established by the Muon implementation. The spectral path is fully compatible with the gram=True Gram Newton-Schulz option and the polar_express / dion coefficient presets added in the GramMuon PR, making scion a strict superset of muon for matrix-shaped parameters.
Contributor
Author
|
Hi all, just following up on this PR. I believe it's ready for review, but please let me know if there's anything else I should address. Thanks for your time! |
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
1 participant
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
This PR adds
scionandscale_by_sciontooptax.contrib, implementing the SCION optimizer from Pethick et al. (arXiv:2502.07529, LIONS-EPFL).SCION is a direct generalisation of Muon: instead of always applying the spectral-norm LMO (the polar factor, via Newton-Schulz), it makes the choice of norm ball explicit and applies the corresponding linear minimization oracle (LMO) per parameter:
lmo'spectral'UV^T'sign'sign(G)element-wise'frobenius'G / ‖G‖_FWith
lmo='auto'(the default), routing is by parameter rank usingcombine.partition:ndim >= 2→ spectral,ndim == 1→ sign. This mirrors howmuon()routes 1D params to Adam, but stays within a single principled framework.Connection to previous work
This builds directly on #1717 (Gram Newton-Schulz + Polar Express coefficients for Muon). The spectral path in
scale_by_sciongoes throughorthogonalize_via_newton_schulzwith the samegram,ns_coeffs, andpreconditioningoptions, so SCION withlmo='spectral'is a strict superset of Muon for matrix-shaped parameters. Thens_coeffs='polar_express'preset from that PR works out of the box here.Changes
optax/contrib/_muon.py:_apply_lmo,ScaleByScionsState,scale_by_scion,scionoptax/contrib/__init__.py: export all three new public namesoptax/contrib/_muon_test.py: 6 new tests (spectral/sign/frobenius paths, auto routing, preset string, agreement withmuon)References