[Feature] Support cautious variant for Muon optimizer (#417)

kozistr · web-flow · commit 6071805d3494 · 2025-08-13T00:09:21.000+09:00
* docs: v3.8.0 changelog

* feature: support cautious variant

* update: test recipe

* docs: v3.8.0 changelog

* update: visualize_optimizers

* docs: v3.8.0 changelog

* docs: visualization

* build(deps): mkdocs plugins

* refactor: awesome-nav

* ci: update labeler
diff --git a/.github/labeler.yml b/.github/labeler.yml
@@ -12,6 +12,9 @@ dependencies:
   - changed-files:
     - any-glob-to-any-file:
       - pyproject.toml
+      - requirements.txt
+      - requirements-dev.txt
+      - requirements-docs.txt
 
 optimizer:
   - changed-files:
diff --git a/docs/.nav.yml b/docs/.nav.yml
@@ -0,0 +1,21 @@
+preserve_directory_names: true
+
+flatten_single_child_sections: false
+
+sort:
+  direction: asc
+  sections: last
+  type: alphabetical
+  ignore_case: false
+
+nav:
+  - Home: index.md
+  - Optimizer: optimizer.md
+  - LR Scheduler: lr_scheduler.md
+  - Loss Function: loss.md
+  - Utilization: util.md
+  - Base: base.md
+  - Visualization: visualization.md
+  - Change Logs:
+    - changelogs/*.md
+  - Q&A: qa.md
diff --git a/docs/changelogs/v3.8.0.md b/docs/changelogs/v3.8.0.md
@@ -6,17 +6,24 @@
     * [Through the River: Understanding the Benefit of Schedule-Free Methods for Language Model Training](https://arxiv.org/abs/2507.09846)
     * You can use this variant by setting `decoupling_c` parameter in the `ScheduleFreeAdamW` optimizer.
 * Add more built-in optimizers, `NAdam`, `RMSProp`, and `LBFGS` optimizers. (#415)
+* Support `cautious` variant for `Muon` optimizer. (#417)
 
 ### Update
 
 * Re-implement `Muon` and `AdaMuon` optimizers based on the recent official implementation. (#408, #410)
     * Their definitions have changed from the previous version, so please check out the documentation!
 * Update the missing optimizers from `__init__.py`. (#415)
+* Add the HuggingFace Trainer example. (#415)
+* Optimize the visualization outputs and change the visualization document to a table layout. (#416)
+
+### Dependency
+
+* Update `mkdocs` dependencies. (#417)
 
 ### CI
 
 * Add some GitHub actions to automate some processes. (#411, #412, #413)
 
-### Example
+## Contributions
 
-* Add the HuggingFace Trainer example. (#415)
+thanks to @AidinHamedi
diff --git a/docs/visualization.md b/docs/visualization.md
diff --git a/examples/visualize_optimizers.py b/examples/visualize_optimizers.py
@@ -16,7 +16,6 @@
 
 filterwarnings('ignore', category=UserWarning)
 
-IMG_FORMAT: str = 'jpg'
 OPTIMIZERS_IGNORE: Tuple[str, ...] = (
     'lomo',
     'adalomo',
@@ -42,8 +41,8 @@
 SEARCH_SEED: int = 42
 LOSS_MIN_THRESHOLD: float = 0.0
 
-DEFAULT_SEARCH_SPACES: Dict[str, object] = {'lr': hp.uniform('lr', 0, 2)}
-SPECIAL_SEARCH_SPACES: Dict[str, Dict[str, object]] = {
+DEFAULT_SEARCH_SPACES: Dict = {'lr': hp.uniform('lr', 0, 2)}
+SPECIAL_SEARCH_SPACES: Dict = {
     'adafactor': {'lr': hp.uniform('lr', 0, 10)},
     'adams': {'lr': hp.uniform('lr', 0, 10)},
     'dadaptadagrad': {'lr': hp.uniform('lr', 0, 10)},
@@ -407,7 +406,7 @@ def execute_experiments(
     """
     for i, (optimizer_class, search_space) in enumerate(optimizers, start=1):
         optimizer_name = optimizer_class.__name__
-        output_path = output_dir / f'{experiment_name}_{optimizer_name}.{IMG_FORMAT}'
+        output_path = output_dir / f'{experiment_name}_{optimizer_name}.jpg'
         if output_path.exists():
             continue
 
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -2,16 +2,6 @@ site_name: pytorch-optimizer
 site_description: 'optimizer & lr scheduler & loss function collections in PyTorch'
 repo_name: 'kozistr/pytorch-optimizer'
 repo_url: 'https://github.com/kozistr/pytorch_optimizer'
-nav:
-  - index.md
-  - base.md
-  - optimizer.md
-  - lr_scheduler.md
-  - loss.md
-  - util.md
-  - visualization.md
-  - ... | changelogs/*.md
-  - qa.md
 theme:
   name: material
   highlightjs: true
@@ -21,7 +11,7 @@ extra_javascript:
   - javascripts/tables.js
 plugins:
   - search
-  - awesome-pages
+  - awesome-nav
   - mkdocstrings:
       handlers:
         python:
diff --git a/pytorch_optimizer/optimizer/muon.py b/pytorch_optimizer/optimizer/muon.py
@@ -123,7 +123,7 @@ def __init__(
 
             group['weight_decouple'] = group.get('weight_decouple', weight_decouple)
 
-        super().__init__(params, {})
+        super().__init__(params, kwargs)
 
     def __str__(self) -> str:
         return 'Muon'
@@ -192,6 +192,9 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                     update = zero_power_via_newton_schulz_5(update, num_steps=group['ns_steps'])
 
+                    if group.get('cautious'):
+                        self.apply_cautious(update, grad)
+
                     lr: float = get_adjusted_lr(group['lr'], p.size(), use_adjusted_lr=group['use_adjusted_lr'])
 
                     p.add_(update.reshape(p.shape), alpha=-lr)
@@ -308,7 +311,7 @@ def __init__(
             group['weight_decouple'] = group.get('weight_decouple', weight_decouple)
             group['eps'] = group.get('eps', eps)
 
-        super().__init__(params, {})
+        super().__init__(params, kwargs)
 
     def __str__(self) -> str:
         return 'AdaMuon'
diff --git a/requirements-docs.txt b/requirements-docs.txt
@@ -1,12 +1,12 @@
 --index-url https://pypi.org/simple
 --extra-index-url https://download.pytorch.org/whl/cpu
 numpy<2.0
-torch==2.6.0
+torch==2.8.0
 mkdocs==1.6.1
-mkdocs-material==9.5.45
-pymdown-extensions==10.12
+mkdocs-material==9.6.16
+pymdown-extensions==10.16.1
 mkdocstrings-python==1.12.2
 markdown-include==0.8.1
 mdx_truly_sane_lists==1.3
-mkdocs-awesome-pages-plugin==2.9.3
+mkdocs-awesome-nav==3.1.2
 griffe==1.5.1
diff --git a/tests/constants.py b/tests/constants.py
@@ -721,6 +721,18 @@
     (AdaShift, {'lr': 1e1, 'keep_num': 1}, 3),
     (MARS, {'lr': 5e-1, 'lr_1d': 5e-1, 'weight_decay': 1e-3}, 3),
     (MARS, {'lr': 5e-1, 'lr_1d': 5e-1, 'weight_decay': 1e-3, 'optimize_1d': True}, 3),
+    (
+        Muon,
+        {
+            'lr': 5e-1,
+            'weight_decay': 1e-3,
+            'use_adjusted_lr': True,
+            'adamw_lr': 5e-1,
+            'adamw_betas': (0.9, 0.98),
+            'adamw_wd': 1e-2,
+        },
+        7,
+    ),
 ]
 STABLE_ADAMW_SUPPORTED_OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
     (ADOPT, {'lr': 1e0, 'weight_decay': 1e-3, 'stable_adamw': True}, 5),
diff --git a/tests/test_optimizer_variants.py b/tests/test_optimizer_variants.py
@@ -8,7 +8,7 @@
     COPT_SUPPORTED_OPTIMIZERS,
     STABLE_ADAMW_SUPPORTED_OPTIMIZERS,
 )
-from tests.utils import build_model, ids, simple_parameter, tensor_to_numpy
+from tests.utils import build_model, build_optimizer_parameter, ids, simple_parameter, tensor_to_numpy
 
 
 @pytest.mark.parametrize('optimizer_config', ADANORM_SUPPORTED_OPTIMIZERS, ids=ids)
@@ -80,11 +80,14 @@ def test_adamd_variant(optimizer_config, environment):
 @pytest.mark.parametrize('optimizer_config', COPT_SUPPORTED_OPTIMIZERS, ids=ids)
 def test_cautious_variant(optimizer_config, environment):
     x_data, y_data = environment
+
     model, loss_fn = build_model()
 
     optimizer_class, config, num_iterations = optimizer_config
 
-    optimizer = optimizer_class(model.parameters(), **config, cautious=True)
+    parameters, config = build_optimizer_parameter(model.parameters(), optimizer_class.__name__, config)
+
+    optimizer = optimizer_class(parameters, **config, cautious=True)
 
     init_loss, loss = np.inf, np.inf
     for _ in range(num_iterations):
diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py
@@ -5,14 +5,14 @@
 
 from pytorch_optimizer.base.exception import NoClosureError, ZeroParameterSizeError
 from pytorch_optimizer.optimizer import DynamicLossScaler, load_optimizer
-from pytorch_optimizer.optimizer.alig import l2_projection
 from pytorch_optimizer.optimizer.grokfast import gradfilter_ema, gradfilter_ma
 from pytorch_optimizer.optimizer.scion import build_lmo_norm
 from tests.constants import COMPLEX_OPTIMIZERS, OPTIMIZERS
 from tests.utils import (
     Example,
     LogisticRegression,
     build_model,
+    build_optimizer_parameter,
     dummy_closure,
     ids,
     names,
@@ -24,29 +24,6 @@
 )
 
 
-def build_optimizer_parameter(parameters, optimizer_name, config):
-    if optimizer_name == 'AliG':
-        config.update({'projection_fn': lambda: l2_projection(parameters, max_norm=1)})
-    elif optimizer_name in ('Muon', 'AdaMuon'):
-        hidden_weights = [p for p in parameters if p.ndim >= 2]
-        hidden_gains_biases = [p for p in parameters if p.ndim < 2]
-
-        parameters = [
-            {'params': hidden_weights, 'use_muon': True},
-            {'params': hidden_gains_biases, 'use_muon': False},
-        ]
-    elif optimizer_name == 'AdamWSN':
-        sn_params = [p for p in parameters if p.ndim == 2]
-        regular_params = [p for p in parameters if p.ndim != 2]
-        parameters = [{'params': sn_params, 'sn': True}, {'params': regular_params, 'sn': False}]
-    elif optimizer_name == 'AdamC':
-        norm_params = [p for i, p in enumerate(parameters) if i == 1]
-        regular_params = [p for i, p in enumerate(parameters) if i != 1]
-        parameters = [{'params': norm_params, 'normalized': True}, {'params': regular_params}]
-
-    return parameters, config
-
-
 @pytest.mark.parametrize('optimizer_fp32_config', OPTIMIZERS, ids=ids)
 def test_f32_optimizers(optimizer_fp32_config, environment):
     def closure(x):
diff --git a/tests/utils.py b/tests/utils.py
@@ -7,6 +7,7 @@
 
 from pytorch_optimizer.base.type import LOSS
 from pytorch_optimizer.optimizer import AdamW, Lookahead, OrthoGrad, ScheduleFreeWrapper
+from pytorch_optimizer.optimizer.alig import l2_projection
 
 
 class LogisticRegression(nn.Module):
@@ -130,3 +131,26 @@ def sphere_loss(x: torch.Tensor) -> torch.Tensor:
 def build_model(use_complex: bool = False):
     torch.manual_seed(42)
     return ComplexLogisticRegression() if use_complex else LogisticRegression(), nn.BCEWithLogitsLoss()
+
+
+def build_optimizer_parameter(parameters, optimizer_name, config):
+    if optimizer_name == 'AliG':
+        config.update({'projection_fn': lambda: l2_projection(parameters, max_norm=1)})
+    elif optimizer_name in ('Muon', 'AdaMuon'):
+        hidden_weights = [p for p in parameters if p.ndim >= 2]
+        hidden_gains_biases = [p for p in parameters if p.ndim < 2]
+
+        parameters = [
+            {'params': hidden_weights, 'use_muon': True},
+            {'params': hidden_gains_biases, 'use_muon': False},
+        ]
+    elif optimizer_name == 'AdamWSN':
+        sn_params = [p for p in parameters if p.ndim == 2]
+        regular_params = [p for p in parameters if p.ndim != 2]
+        parameters = [{'params': sn_params, 'sn': True}, {'params': regular_params, 'sn': False}]
+    elif optimizer_name == 'AdamC':
+        norm_params = [p for i, p in enumerate(parameters) if i == 1]
+        regular_params = [p for i, p in enumerate(parameters) if i != 1]
+        parameters = [{'params': norm_params, 'normalized': True}, {'params': regular_params}]
+
+    return parameters, config
