Independent wd for orthogonalized optimizer

emerging_optimizers/orthogonalized_optimizers/muon.py

@@ -67,6 +67,7 @@ def __init__(
         use_nesterov: bool = False,
         weight_decay: float = 0.01,
         use_decoupled_weight_decay: bool = True,
+        use_independent_weight_decay: bool = False,
         fp32_matmul_prec: str = "medium",
         coefficient_type: str = "quintic",
         num_ns_steps: int = 5,
@@ -107,6 +108,7 @@ def scaled_orthogonalize_fn(grad: torch.Tensor) -> torch.Tensor:
             use_nesterov,
             weight_decay,
             use_decoupled_weight_decay,
+            use_independent_weight_decay,
             fp32_matmul_prec,
             scaled_orthogonalize_fn,
         )

emerging_optimizers/orthogonalized_optimizers/orthogonalized_optimizer.py

@@ -36,6 +36,8 @@
         weight_decay: The weight decay used by the optimizer, default to be decoupled weight decay.
             See Decoupled Weight Decay Regularization: https://arxiv.org/abs/1711.05101
         use_decoupled_weight_decay: Whether to use decoupled weight decay, default to be True.
+        use_independent_weight_decay: Whether to use independent weight decay (https://arxiv.org/abs/2510.19093),
+            default to be False.
         fp32_matmul_prec: Precision of the matmul operations in optimizer states GEMM operations.
 """
 
@@ -100,6 +102,7 @@ def __init__(
         use_nesterov: bool,
         weight_decay: float,
         use_decoupled_weight_decay: bool,
+        use_independent_weight_decay: bool,
         fp32_matmul_prec: str,
         scaled_orthogonalize_fn: Callable | None = None,
         **kwargs: Any,
@@ -115,6 +118,7 @@ def __init__(
             use_nesterov=use_nesterov,
             weight_decay=weight_decay,
             use_decoupled_weight_decay=use_decoupled_weight_decay,
+            use_independent_weight_decay=use_independent_weight_decay,
             **kwargs,
         )
 
@@ -154,7 +158,12 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                 if group["weight_decay"] > 0.0:
                     if group["use_decoupled_weight_decay"]:
                         # Apply decoupled weight decay
-                        p.add_(p, alpha=(-group["lr"] * group["weight_decay"]))
+                        if group["use_independent_weight_decay"]:
+                            # use independent weight decay
+                            weight_decay_scale = group["weight_decay"]
+                        else:
+                            weight_decay_scale = group["weight_decay"] * group["lr"]
+                        p.add_(p, alpha=(-weight_decay_scale))
                     else:
                         # add l2 regularization before preconditioning (i.e. adding a squared loss term)
                         grad += group["weight_decay"] * p

tests/test_orthogonalized_optimizer.py

@@ -43,6 +43,7 @@ def test_orthogonalized_optimizer_core_matches_sgd(self, shape) -> None:
             use_nesterov=False,
             weight_decay=0.5,
             use_decoupled_weight_decay=True,
+            use_independent_weight_decay=False,
             fp32_matmul_prec="highest",
         )
 
@@ -84,6 +85,7 @@ def test_orthogonalized_optimizer_core_matches_sgd_with_momentum(self, shape) ->
             use_nesterov=False,
             weight_decay=0.0,
             use_decoupled_weight_decay=False,
+            use_independent_weight_decay=False,
             fp32_matmul_prec="highest",
         )
 
@@ -134,6 +136,7 @@ def dummy_interleaved_split_orth_fn(x: torch.Tensor) -> torch.Tensor:
             use_nesterov=False,
             weight_decay=0.0,
             use_decoupled_weight_decay=False,
+            use_independent_weight_decay=False,
             fp32_matmul_prec="highest",
             scaled_orthogonalize_fn=dummy_interleaved_split_orth_fn,
         )