Feat: Add decoupled, dimension-scaled weight decay option

adv_optm/__init__.py

@@ -20,4 +20,4 @@
     "SinkSGD_adv",
 ]
 
-__version__ = "2.5.11"
+__version__ = "2.6.1.dev3"

adv_optm/optim/AdaMuon_adv.py

@@ -137,6 +137,7 @@ def __init__(
         # Decoupled/cautious weight decay
         weight_decay: float = 0,
         cautious_wd: bool = False,
+        scaled_wd: bool = False,
         # Nesterov momentum
         nesterov: bool = True,
         nesterov_coef: float | None = None,
@@ -227,7 +228,7 @@ def __init__(
 
         defaults = {
             "lr": lr, "betas": betas, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
-            "eps": eps, "rms_rescaling": rms_rescaling, "ns_steps": ns_steps,
+            "eps": eps, "rms_rescaling": rms_rescaling, "ns_steps": ns_steps, "scaled_wd": scaled_wd,
             "ns_eps": ns_eps, "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor,
             "vector_reshape": vector_reshape,
             "nesterov":nesterov, "nesterov_coef": nesterov_coef, "use_atan2":use_atan2,

adv_optm/optim/AdamW_adv.py

@@ -98,6 +98,7 @@ def __init__(
         weight_decay: float = 0.0,
         fisher_wd: bool = False,
         cautious_wd: bool = False,
+        scaled_wd: bool = False,
         # Adam's Bias Correction
         use_bias_correction: bool = True,
         # Stochastic Rounding for BF16
@@ -156,7 +157,7 @@ def __init__(
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
-            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd,
+            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd, "scaled_wd": scaled_wd,
             "use_atan2": use_atan2, "nesterov": nesterov, "nesterov_coef": nesterov_coef,
             "normed_momentum": normed_momentum,
             "orthogonal_gradient": orthogonal_gradient, "use_bias_correction": use_bias_correction,

adv_optm/optim/Adopt_adv.py

@@ -101,6 +101,7 @@ def __init__(
         weight_decay: float = 0.0,
         fisher_wd: bool = False,
         cautious_wd: bool = False,
+        scaled_wd: bool = False,
         # ADOPT clipping
         clip_lambda: Optional[Callable[[int], float]] = lambda step: step**0.25,
         # Adam_atan2 (scale invariant)
@@ -157,7 +158,7 @@ def __init__(
             state_precision = "factored"
 
         defaults = {
-            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
+            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "scaled_wd": scaled_wd,
             "fisher_wd": fisher_wd, "cautious_wd": cautious_wd, "orthogonal_gradient": orthogonal_gradient,
             "nesterov": nesterov, "nesterov_coef": nesterov_coef,
             "kourkoutas_beta": kourkoutas_beta, "beta2_min": beta2_min, "ema_alpha": ema_alpha,

adv_optm/optim/Lion_adv.py

@@ -64,6 +64,7 @@ def __init__(
         # Decoupled/cautious weight decay
         weight_decay: float = 0.0,
         cautious_wd: bool = False,
+        scaled_wd: bool = False,
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # OrthoGrad
@@ -96,6 +97,7 @@ def __init__(
             betas=betas,
             weight_decay=weight_decay,
             cautious_wd=cautious_wd,
+            scaled_wd=scaled_wd,
             vector_reshape=vector_reshape,
             orthogonal_gradient=orthogonal_gradient,
             kappa_p=kappa_p,

adv_optm/optim/Muon_adv.py

@@ -111,6 +111,7 @@ def __init__(
         # Decoupled/cautious weight decay
         weight_decay: float = 0.0,
         cautious_wd: bool = False,
+        scaled_wd: bool = False,
         # Nesterov momentum
         nesterov: bool = True,
         nesterov_coef: float | None = None,
@@ -201,7 +202,7 @@ def __init__(
         defaults = {
             "lr": lr, "beta1": beta1, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
             "nesterov": nesterov, "nesterov_coef": nesterov_coef, "ns_steps": ns_steps, "ns_eps": ns_eps,
-            "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor,
+            "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor, "scaled_wd": scaled_wd,
             "vector_reshape": vector_reshape,  "rms_rescaling": rms_rescaling,
             "orthogonal_gradient": orthogonal_gradient,
             'compiled_optimizer': compiled_optimizer,

adv_optm/optim/Prodigy_adv.py

@@ -115,6 +115,7 @@ def __init__(
         weight_decay: float = 0.0,
         fisher_wd: bool = False,
         cautious_wd: bool = False,
+        scaled_wd: bool = False,
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # Adam_atan2 (scale invariant)
@@ -181,7 +182,7 @@ def __init__(
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
-            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd,
+            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd, "scaled_wd": scaled_wd,
             "use_atan2": use_atan2,
             "orthogonal_gradient": orthogonal_gradient,
             "compiled_optimizer": compiled_optimizer,

adv_optm/optim/SignSGD_adv.py

@@ -59,6 +59,7 @@ def __init__(
         # weight decay features
         geometric_wd: bool = False,
         cautious_wd: bool = False,
+        scaled_wd: bool = False,
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # OrthoGrad
@@ -108,6 +109,7 @@ def __init__(
             momentum=momentum,
             weight_decay=weight_decay,
             cautious_wd=cautious_wd,
+            scaled_wd=scaled_wd,
             geometric_wd=geometric_wd,
             vector_reshape=vector_reshape,
             orthogonal_gradient=orthogonal_gradient,

adv_optm/optim/SinkSGD_adv.py

@@ -66,6 +66,7 @@ def __init__(
         # weight decay features
         geometric_wd: bool = False,
         cautious_wd: bool = False,
+        scaled_wd: bool = False,
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
         # OrthoGrad
@@ -103,7 +104,7 @@ def __init__(
         defaults = {
             "lr": lr, "momentum": momentum,
             "weight_decay": weight_decay, "nesterov": nesterov, "nesterov_coef": nesterov_coef, "normed_momentum": normed_momentum, "snr_cond": snr_cond,
-            "geometric_wd": geometric_wd, "cautious_wd": cautious_wd,
+            "geometric_wd": geometric_wd, "cautious_wd": cautious_wd, "scaled_wd": scaled_wd,
             "orthogonal_gradient": orthogonal_gradient, 
             "compiled_optimizer": compiled_optimizer,
             "sinkhorn_iterations": sinkhorn_iterations,

adv_optm/util/param_update.py

@@ -6,7 +6,7 @@
 
 from typing import Dict, Any
 
-from .scaled_optm import adjust_wds
+from .scaled_optm import adjust_wds, scale_wd
 from .centered_decay import dequantize_anchor
 
 _generators: Dict[torch.device, torch.Generator] = {}
@@ -18,8 +18,8 @@ def _apply_weight_decay(
     p: Tensor,
     state: Dict[str, Any],
     group: Dict[str, Any],
-    scaled_wd: float | Tensor | None,
-    scaled_cwd: float | Tensor | None,
+    eff_wd: float | Tensor | None,
+    eff_cwd: float | Tensor | None,
     wd_target: Tensor | None = None,
     cwd_target: Tensor | None = None,
 ) -> None:
@@ -29,26 +29,26 @@ def _apply_weight_decay(
     cautious = group.get('cautious_wd', False)
 
     # Standard Weight Decay (pulls toward zero)
-    if scaled_wd is not None:
+    if eff_wd is not None:
         if wd_target is None:
             wd_target = p_calc
         # Cautious Weight Decay: only decay if the update pushes in the same direction as the decay
         if cautious:
             mask = (update_calc * p_calc >= 0).to(p_calc.dtype)
-            if isinstance(scaled_wd, Tensor):
-                p_calc.addcmul_(wd_target, mask * scaled_wd, value=-1.0)
+            if isinstance(eff_wd, Tensor):
+                p_calc.addcmul_(wd_target, mask * eff_wd, value=-1.0)
             else:
-                p_calc.addcmul_(wd_target, mask, value=-scaled_wd)
+                p_calc.addcmul_(wd_target, mask, value=-eff_wd)
             del mask
         else:
             # Standard decoupled weight decay
-            if isinstance(scaled_wd, Tensor):
-                p_calc.addcmul_(wd_target, scaled_wd, value=-1.0)
+            if isinstance(eff_wd, Tensor):
+                p_calc.addcmul_(wd_target, eff_wd, value=-1.0)
             else:
-                p_calc.add_(wd_target, alpha=-scaled_wd)
+                p_calc.add_(wd_target, alpha=-eff_wd)
 
     # Centered Weight Decay (pulls toward anchor)
-    if scaled_cwd is not None and 'anchor_data' in state:
+    if eff_cwd is not None and 'anchor_data' in state:
         if cwd_target is not None:
             decay_target = cwd_target
         else:
@@ -59,17 +59,17 @@ def _apply_weight_decay(
         if cautious:
             # Cautious Weight Decay: only decay if the update pushes in the same direction as the decay
             mask = (update_calc * decay_target >= 0).to(p_calc.dtype)
-            if isinstance(scaled_cwd, Tensor):
-                p_calc.addcmul_(decay_target, mask * scaled_cwd, value=-1.0)
+            if isinstance(eff_cwd, Tensor):
+                p_calc.addcmul_(decay_target, mask * eff_cwd, value=-1.0)
             else:
-                p_calc.addcmul_(decay_target, mask, value=-scaled_cwd)
+                p_calc.addcmul_(decay_target, mask, value=-eff_cwd)
             del mask
         else:
             # Standard decoupled weight decay
-            if isinstance(scaled_cwd, Tensor):
-                p_calc.addcmul_(decay_target, scaled_cwd, value=-1.0)
+            if isinstance(eff_cwd, Tensor):
+                p_calc.addcmul_(decay_target, eff_cwd, value=-1.0)
             else:
-                p_calc.add_(decay_target, alpha=-scaled_cwd)
+                p_calc.add_(decay_target, alpha=-eff_cwd)
 
         if cwd_target is None:
             del decay_target
@@ -105,18 +105,24 @@ def apply_parameter_update(
     wd = group["weight_decay"] if wd is None else wd
     cwd = group.get("centered_wd", 0.0)
     wd, cwd = adjust_wds(wd, cwd, p)
+    scaled_wd = group.get("scaled_wd", False)
+    decoupled = scaled_wd or decoupled
 
     # Calculate global decay factor for decoupled vs standard
     decay_factor = (lr / self._init_lr) if decoupled else lr
 
-    scaled_wd = (wd * decay_factor) if wd != 0 else None
-    scaled_cwd = (cwd * decay_factor) if cwd != 0 else None
+    eff_wd = (wd * decay_factor) if wd != 0 else None
+    eff_cwd = (cwd * decay_factor) if cwd != 0 else None
+
+    if scaled_wd:
+        eff_wd = scale_wd(eff_wd, p, skip_vectors=True)
+        eff_cwd = scale_wd(eff_cwd, p, skip_vectors=False)
 
     if wd_scaler is not None:
-        if scaled_wd is not None:
-            scaled_wd = scaled_wd * wd_scaler
-        if scaled_cwd is not None:
-            scaled_cwd = scaled_cwd * wd_scaler
+        if eff_wd is not None:
+            eff_wd = eff_wd * wd_scaler
+        if eff_cwd is not None:
+            eff_cwd = eff_cwd * wd_scaler
 
     state = self.state[p]
 
@@ -129,8 +135,8 @@ def apply_parameter_update(
         cwd_t = cwd_target.float() if cwd_target is not None else None
 
         # Apply weight decay if needed
-        if scaled_wd is not None or scaled_cwd is not None:
-            _apply_weight_decay(p_fp32, update_fp32, p, state, group, scaled_wd, scaled_cwd, wd_t, cwd_t)
+        if eff_wd is not None or eff_cwd is not None:
+            _apply_weight_decay(p_fp32, update_fp32, p, state, group, eff_wd, eff_cwd, wd_t, cwd_t)
 
         # Apply main update
         p_fp32.add_(-update_fp32)
@@ -147,8 +153,8 @@ def apply_parameter_update(
 
     else:
         # Standard path for non-bfloat16 or without stochastic rounding
-        if scaled_wd is not None or scaled_cwd is not None:
-            _apply_weight_decay(p, update, p, state, group, scaled_wd, scaled_cwd, wd_target, cwd_target)
+        if eff_wd is not None or eff_cwd is not None:
+            _apply_weight_decay(p, update, p, state, group, eff_wd, eff_cwd, wd_target, cwd_target)
 
         # Apply main update
         p.add_(-update)

adv_optm/util/scaled_optm.py

@@ -97,6 +97,25 @@ def adjust_wds(wd: float, cwd: float, p: torch.Tensor) -> tuple[float, float]:
         # Centered WD safely regularizes the delta without collapsing base feature variance.
         return wd, cwd
 
+def scale_wd(wd: float | None, p: torch.Tensor, skip_vectors: bool = False) -> float:
+    """
+    Scale-invariant, dimension-scaled weight decay.
+    """
+    if wd is None:
+        return wd
+    if getattr(p, '_is_oft', False):
+        n_el = p.shape[-1]
+        b = (1.0 + math.sqrt(1.0 + 8.0 * n_el)) / 2.0
+        wd = (2 * wd) / (b - 1)
+        return wd
+
+    is_vector = p.ndim < 2 or getattr(p, '_is_dora_scale', False) or getattr(p, 'is_vector', False) 
+    if is_vector:
+        return 0.0 if skip_vectors else wd
+
+    if p.ndim >= 2:
+        width = p.numel() // p.shape[0]
+        return wd / width
 
 def is_spectral(p: torch.Tensor) -> bool:
     """Determines if a parameter should undergo spectral normalization updates."""

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="adv_optm",
-    version="2.5.11",
+    version="2.6.1.dev3",
     author="Koratahiu",
     author_email="hiuhonor@gmail.com",
     license='Apache 2.0',