adv-optm 2.5.9__tar.gz → 2.6.1.dev2__tar.gz

{adv_optm-2.5.9 → adv_optm-2.6.1.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: adv_optm
-Version: 2.5.9
+Version: 2.6.1.dev2
 Summary: A family of highly efficient, lightweight yet powerful optimizers.
 Home-page: https://github.com/Koratahiu/Advanced_Optimizers
 Author: Koratahiu

{adv_optm-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "SinkSGD_adv",
 ]
 
-__version__ = "2.5.9"
+__version__ = "2.6.1.dev2"

{adv_optm-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm/optim/AdaMuon_adv.py

@@ -137,6 +137,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
         # 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 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
         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-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm/optim/AdamW_adv.py

@@ -98,6 +98,7 @@ class AdamW_adv(torch.optim.Optimizer):
         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 @@ class AdamW_adv(torch.optim.Optimizer):
 
         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-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm/optim/Adopt_adv.py

@@ -101,6 +101,7 @@ class Adopt_adv(torch.optim.Optimizer):
         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 @@ class Adopt_adv(torch.optim.Optimizer):
             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-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm/optim/Lion_adv.py

@@ -64,6 +64,7 @@ class Lion_adv(torch.optim.Optimizer):
         # 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 @@ class Lion_adv(torch.optim.Optimizer):
             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-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm/optim/Muon_adv.py

@@ -111,6 +111,7 @@ class Muon_adv(torch.optim.Optimizer):
         # 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 @@ class Muon_adv(torch.optim.Optimizer):
         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-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm/optim/Prodigy_adv.py

@@ -115,6 +115,7 @@ class Prodigy_adv(torch.optim.Optimizer):
         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 @@ class Prodigy_adv(torch.optim.Optimizer):
 
         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-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm/optim/SignSGD_adv.py

@@ -59,6 +59,7 @@ class SignSGD_adv(torch.optim.Optimizer):
         # 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 @@ class SignSGD_adv(torch.optim.Optimizer):
             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-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm/optim/SinkSGD_adv.py

@@ -66,6 +66,7 @@ class SinkSGD_adv(torch.optim.Optimizer):
         # 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 @@ class SinkSGD_adv(torch.optim.Optimizer):
         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-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm/util/OrthoGrad.py

@@ -43,7 +43,7 @@ def iterative_ortho_project(p: torch.Tensor, grad: torch.Tensor, iters: int = 3)
     # 1D Vector Case fallback to the standard OrthoGrad
     is_vector = p.ndim < 2 or getattr(p, '_is_dora_scale', False) or getattr(p, 'is_vector', False)
     if is_vector:
-        return _orthogonalize_gradient(p, grad)
+        return flattened_ortho_project(p, grad)
 
     original_shape = grad.shape
 

{adv_optm-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm/util/param_update.py

@@ -6,7 +6,7 @@ import torch.nn.functional as F
 
 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
 
     # 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 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:
+            if scaled_wd:
+                eff_wd = scale_wd(eff_wd, p)
+            eff_wd = eff_wd * wd_scaler
+        if eff_cwd is not None:
+            if scaled_wd:
+                eff_cwd = scale_wd(eff_cwd, p)
+            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-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm/util/scaled_optm.py

@@ -7,14 +7,14 @@ import math
 _OFT_INDICES_CACHE = {}
 _OFT_IDENTITY_CACHE = {}
 
-def get_cached_structural_tensors(b: int, dtype: torch.dtype, device: torch.device):
+def get_cached_structural_tensors(b: int, device: torch.device):
     """
-    Retrieves or creates structural tensors (indices and Identity) for OFT exact geometry.
+    Retrieves or creates structural tensors (indices) for OFT exact geometry.
     Caches them globally to prevent redundant memory allocation across thousands of layers.
     """
-    global _OFT_INDICES_CACHE, _OFT_IDENTITY_CACHE
+    global _OFT_INDICES_CACHE
 
-    # Cache for Indices (Dtype independent, only depends on block size and device)
+    # Cache for Indices
     idx_key = (b, device)
     if idx_key not in _OFT_INDICES_CACHE:
         rows, cols = torch.triu_indices(b, b, 1, device=device)
@@ -22,15 +22,8 @@ def get_cached_structural_tensors(b: int, dtype: torch.dtype, device: torch.devi
     else:
         rows, cols = _OFT_INDICES_CACHE[idx_key]
 
-    # Cache for Identity Matrix (Depends on block size, dtype, and device)
-    id_key = (b, dtype, device)
-    if id_key not in _OFT_IDENTITY_CACHE:
-        I = torch.eye(b, dtype=dtype, device=device).unsqueeze(0)
-        _OFT_IDENTITY_CACHE[id_key] = I
-    else:
-        I = _OFT_IDENTITY_CACHE[id_key]
 
-    return rows, cols, I
+    return rows, cols
 
 def scale_update(
     p: torch.Tensor,
@@ -59,7 +52,7 @@ def scale_update(
         return max_abs_normalization(update, dim=None, lr=lr)
 
     # OFT Block Parameters: shape (k, C(b,2))
-    # Direct spectral normalization on the skew-symmetric blocks, followed by Riemannian preconditioning.
+    # Direct spectral normalization on the skew-symmetric blocks.
     if is_oft:
         return apply_spectral_riemannian_oft(p, update, lr, state)
 
@@ -104,6 +97,20 @@ 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, p: torch.Tensor) -> float:
+    """
+    Scale-invariant, dimension-scaled weight decay.
+    """
+    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
+
+    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."""
@@ -176,34 +183,26 @@ def apply_spectral_riemannian_oft(
     state: dict
 ) -> torch.Tensor:
     """
-    Applies Spectral Normalization directly on the skew-symmetric gradient,
-    then uses True Matrix Preconditioning: M @ G @ M where M = (I - Q^2).
-    Neutralizes the derivative shrinkage of the Cayley transform.
+    Applies Spectral Normalization directly on the skew-symmetric gradient.
     """
     n_el = p.shape[-1]
     block_size = int((1 + math.sqrt(1 + 8 * n_el)) / 2)
     device, dtype = p.device, p.dtype
-    rows, cols, I = get_cached_structural_tensors(block_size, dtype, device)
+    rows, cols = get_cached_structural_tensors(block_size, device)
 
     # Flatten any prepended batch dimensions for processing
     orig_shape = p.shape
 
     # Align the scale of p with the forward pass
     scale_factor = getattr(p, '_oft_scale_factor', 1.0)
-    p_flat = p.view(-1, n_el) / scale_factor
 
     update_flat = update.view(-1, n_el)
-    batch_size = p_flat.shape[0]
+    batch_size = update_flat.shape[0]
 
     # Initialize matrices
-    Q = torch.zeros(batch_size, block_size, block_size, device=device, dtype=dtype)
     G = torch.zeros(batch_size, block_size, block_size, device=device, dtype=dtype)
     batch_idx = torch.arange(batch_size, device=device)[:, None]
 
-    # Construct skew-symmetric parameter matrix Q
-    Q = Q.index_put((batch_idx, rows, cols), p_flat)
-    Q = Q - Q.transpose(-2, -1)
-
     # Construct skew-symmetric gradient matrix G
     G = G.index_put((batch_idx, rows, cols), update_flat)
     G = G - G.transpose(-2, -1)
@@ -235,21 +234,9 @@ def apply_spectral_riemannian_oft(
     target_scale = 0.5 * scale_factor
     spectral_eps = 1.0 / (2.0 * math.sqrt(block_size))
 
-    # Rescale G
     scale = lr * (target_scale / max_sigma.clamp_min(spectral_eps))
-    G = G * scale
-
-    # Apply Riemannian Preconditioning
-    # Compute True Matrix Preconditioner M = I - Q^2
-    M = I - torch.bmm(Q, Q)
-
-    # Apply exact preconditioning: G_prec = M @ G @ M
-    G_prec = torch.bmm(torch.bmm(M, G), M)
-
-    # Extract the preconditioned upper-triangular elements
-    update_prec_flat = G_prec[batch_idx, rows, cols]
 
-    return update_prec_flat.view(orig_shape)
+    return update_flat.mul_(scale).view(orig_shape)
 
 
 @torch.no_grad()

{adv_optm-2.5.9 → adv_optm-2.6.1.dev2}/adv_optm.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: adv_optm
-Version: 2.5.9
+Version: 2.6.1.dev2
 Summary: A family of highly efficient, lightweight yet powerful optimizers.
 Home-page: https://github.com/Koratahiu/Advanced_Optimizers
 Author: Koratahiu

{adv_optm-2.5.9 → adv_optm-2.6.1.dev2}/setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
 
 setup(
     name="adv_optm",
-    version="2.5.9",
+    version="2.6.1.dev2",
     author="Koratahiu",
     author_email="hiuhonor@gmail.com",
     license='Apache 2.0',