adv-optm 2.4.dev17__tar.gz → 2.4.dev18__tar.gz

{adv_optm-2.4.dev17 → adv_optm-2.4.dev18}/PKG-INFO

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

{adv_optm-2.4.dev17 → adv_optm-2.4.dev18}/adv_optm/__init__.py

@@ -24,4 +24,4 @@ __all__ = [
     "SinkSGD_adv",
 ]
 
-__version__ = "2.4.dev17"
+__version__ = "2.4.dev18"

{adv_optm-2.4.dev17 → adv_optm-2.4.dev18}/adv_optm/optim/SinkSGD_adv.py

@@ -1,6 +1,6 @@
 import torch
 
-from typing import Optional, Callable
+import math
 
 from ..util import param_update
 from ..util.factorization_util import _get_effective_shape, _reconstruct_state, _factorize_state
@@ -9,7 +9,7 @@ from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.scaled_optm import scale_update, is_spectral, init_spectral_norm
 from ..util.centered_decay import _init_anchor
 from ..util.state_util import init_state_tensor, get_state, set_state, upcast_grad_for_precision
-from ..util.sinkhorn import apply_sr_sinkhorn
+from ..util.sinkhorn import apply_sr_sinkhorn, _sinkhorn_sq_grad, get_sinkhorn_wd_scaler
 from ..util.signed_util import apply_stochastic_sign_
 
 class SinkSGD_adv(torch.optim.Optimizer):
@@ -26,8 +26,6 @@ class SinkSGD_adv(torch.optim.Optimizer):
         weight_decay (float): weight decay (L2 penalty or decoupled) (default: 0).
         nesterov (bool): enables Nesterov momentum. Only applicable when momentum
             is non-zero. (default: False)
-        decoupled_wd (bool): whether to apply decoupled weight decay (like AdamW) 
-            instead of standard L2 penalty. (default: False)
         cautious_wd (bool): Enables Cautious Weight Decay. If True, weight decay is
             applied only to parameter coordinates where the sign of the parameter
             and the sign of the optimizer update align (default: False).
@@ -61,11 +59,13 @@ class SinkSGD_adv(torch.optim.Optimizer):
         orthogonal_sinkhorn: bool = False,
         # Normalization then Momentum
         normed_momentum: bool = False,
+        # Centered Variance Precondition
+        centered_vt: bool = False,
         # Nesterov Momentum
         nesterov: bool = False,
         nesterov_coef: float | None = None,
-        # Decoupled/cautious weight decay
-        decoupled_wd: bool = False,
+        # weight decay features
+        geometric_wd: bool = False,
         cautious_wd: bool = False,
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
@@ -101,8 +101,8 @@ 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,
-            "decoupled_wd": decoupled_wd, "cautious_wd": cautious_wd,
+            "weight_decay": weight_decay, "nesterov": nesterov, "nesterov_coef": nesterov_coef, "normed_momentum": normed_momentum, "centered_vt": centered_vt,
+            "geometric_wd": geometric_wd, "cautious_wd": cautious_wd,
             "orthogonal_gradient": orthogonal_gradient, 
             "compiled_optimizer": compiled_optimizer,
             "sinkhorn_iterations": sinkhorn_iterations,
@@ -182,6 +182,11 @@ class SinkSGD_adv(torch.optim.Optimizer):
                     if group['momentum'] != 0:
                         init_state_tensor(state, 'momentum_buffer', p.shape, actual_precision, p.device, dtype)
 
+                if group.get('centered_vt', False):
+                    p_shape = p.shape
+                    state['vt_row'] = torch.zeros(p_shape[:-1], device=device, dtype=torch.float32)
+                    state['vt_col'] = torch.zeros(p_shape[:-2] + p_shape[-1:], device=device, dtype=torch.float32)
+
             if group.get('spectral_normalization', False) and is_spectral(p):
                 init_spectral_norm(state, p)
 
@@ -237,7 +242,7 @@ class SinkSGD_adv(torch.optim.Optimizer):
         if group.get('normed_momentum', False):
             if not is_vector:
                 # Sinkhorn iterative normalization
-                grad = apply_sr_sinkhorn(grad, p, ortho_project=orthogonal_sinkhorn, iters=sinkhorn_iterations)
+                grad = apply_sr_sinkhorn(grad, iters=sinkhorn_iterations, p=p, ortho_project=orthogonal_sinkhorn)
             else:
                 # For vectors, apply adaptive stochastic sign
                 grad = apply_stochastic_sign_(grad, sign_noise, is_vector=is_vector)
@@ -271,6 +276,24 @@ class SinkSGD_adv(torch.optim.Optimizer):
 
             if momentum != 0:
                 buf = get_state(state, 'momentum_buffer', actual_precision)
+
+                if group.get('centered_vt', False):
+                    vt_row, vt_col = state['vt_row'], state['vt_col']
+                    grad_vt = grad - buf
+                    grad_vt_sq = grad_vt * grad_vt
+                    mean_row_grad = grad_vt_sq.mean(dim=-1)
+                    mean_col_grad = grad_vt_sq.mean(dim=-2)
+                    vt_row.mul_(momentum).add_(mean_row_grad, alpha=1.0 - momentum)
+                    vt_col.mul_(momentum).add_(mean_col_grad, alpha=1.0 - momentum)
+                    if nesterov:
+                        nv_coef = momentum if nesterov_coef is None else nesterov_coef
+                        vt_row = vt_row.lerp(mean_row_grad, 1.0 - nv_coef)
+                        vt_col = vt_col.lerp(mean_col_grad, 1.0 - nv_coef)
+                    vt = _sinkhorn_sq_grad(vt_row, vt_col)
+                else:
+                    vt_row = None
+                    vt_col = None
+
                 buf.lerp_(grad, 1 - momentum)
 
                 set_state(state, 'momentum_buffer', buf, actual_precision, random_int_state_tensor)
@@ -285,21 +308,34 @@ class SinkSGD_adv(torch.optim.Optimizer):
 
             del random_int_state_tensor
 
+        if group.get('centered_vt', False):
+            denom = vt
+            update.atan2_(denom)
+        else:
+            denom = None
+
         if not group.get('normed_momentum', False):
             if not is_vector:
                 # Sinkhorn iterative normalization
-                update = apply_sr_sinkhorn(update, p, ortho_project=orthogonal_sinkhorn, iters=sinkhorn_iterations)
+                update = apply_sr_sinkhorn(update, iters=sinkhorn_iterations, p=p, ortho_project=orthogonal_sinkhorn)
             else:
                 # For vectors, apply adaptive stochastic sign
                 update = apply_stochastic_sign_(update, sign_noise, is_vector=is_vector)
 
+        if group.get('geometric_wd', False):
+            wd_scaler = get_sinkhorn_wd_scaler(p, row_denom=vt_row, col_denom=vt_col)
+        else:
+            wd_scaler = None
+
         update_scaling = step_size
         if group.get('spectral_normalization', False):
             update = scale_update(p, update, update_scaling, state=state)
         else:
+            if group.get('centered_vt', False):
+                update_scaling = update_scaling * (4/math.pi)
             update.mul_(update_scaling)
 
-        param_update.apply_parameter_update(self, p, group, update, step_size, random_int_tensor=random_int_tensor)
+        param_update.apply_parameter_update(self, p, group, update, step_size, random_int_tensor=random_int_tensor, wd_scaler=wd_scaler)
 
     def compile(self, *args, **kwargs):
         self._compiled_step_parameter = torch.compile(self._step_parameter, *args, **kwargs)

{adv_optm-2.4.dev17 → adv_optm-2.4.dev18}/adv_optm/util/sinkhorn.py

@@ -1,7 +1,7 @@
 import math
 import torch
 
-def apply_sr_sinkhorn(update: torch.Tensor, p: torch.Tensor, ortho_project: bool, iters: int = 5) -> torch.Tensor:
+def apply_sr_sinkhorn(update: torch.Tensor, iters: int = 5, p: torch.Tensor | None = None, ortho_project: bool = False) -> torch.Tensor:
     """
     Applies Square-Root Sinkhorn (SR-Sinkhorn) multi-normalization.
     As described in 'Gradient Multi-Normalization for Efficient LLM Training'.
@@ -47,13 +47,16 @@ def apply_sr_sinkhorn(update: torch.Tensor, p: torch.Tensor, ortho_project: bool
     # In-place alternating Sinkhorn normalization steps
     for _ in range(iters):
         # First normalization step
-        norm1 = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(1e-12)
+        # Stability floor: equivalent to a single-element vector norm lower bound (lb)
+        norm1_lb = 1 / math.sqrt(update_2d.shape[dim])
+        norm1 = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(norm1_lb)
         update_2d.mul_(scale_first / norm1)
         if ortho_project:
             update_2d = ortho_normed(param_2d, update_2d, p_norm_sq_dim, dim, scale_first)
 
         # Second normalization step
-        norm2 = update_2d.norm(p=2, dim=1-dim, keepdim=True).clamp_min_(1e-12)
+        norm2_lb = 1 / math.sqrt(update_2d.shape[1-dim])
+        norm2 = update_2d.norm(p=2, dim=1-dim, keepdim=True).clamp_min_(norm2_lb)
         update_2d.mul_(scale_second / norm2)
         if ortho_project:
             update_2d = ortho_normed(param_2d, update_2d, p_norm_sq_adim, 1-dim, scale_second)
@@ -72,6 +75,69 @@ def ortho_normed(p_2d, update_2d, p_norm_sq, dim, target_norm):
     update_2d.addcmul_(proj, p_2d, value=-1.0)
 
     # Magnitude Preservation
-    g_orth_norm = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(1e-12)
+    norm_lb = 1 / math.sqrt(update_2d.shape[dim])
+    g_orth_norm = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(norm_lb)
     scale_factor = target_norm / g_orth_norm
     return update_2d.mul_(scale_factor)
+
+def _sinkhorn_sq_grad(
+    vt_row: torch.Tensor,
+    vt_col: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Reconstructs the variance precondition from its rank-1 factors.
+    Modified from:
+    https://github.com/jettify/pytorch-optimizer/blob/master/torch_optimizer/adafactor.py
+    """
+    r_factor = (
+        (vt_row / vt_row.mean(dim=-1).clamp_min_(1e-30))
+        .sqrt_()
+        .unsqueeze(-1)
+    )
+    c_factor = vt_col.unsqueeze(-2).sqrt()
+    return torch.mul(r_factor, c_factor) 
+
+def get_sinkhorn_wd_scaler(
+    p: torch.Tensor, 
+    row_denom: torch.Tensor | None = None, 
+    col_denom: torch.Tensor | None = None
+):
+    """
+    Computes a structural weight decay multiplier.
+    Penalizes parameters belonging to dominant rows/columns more heavily, 
+    while protecting parameters in under-utilized/noisy rows/columns from decay.
+    """
+    if p.ndim < 2:
+        return 1.0 
+
+    p_2d = p.view(p.shape[0], -1)
+
+    # Lower bounds based on the effective 2D shapes
+    row_lb = 1 / math.sqrt(p_2d.shape[1])
+    col_lb = 1 / math.sqrt(p_2d.shape[0])
+
+    # Get the norms
+    row_norms = torch.linalg.vector_norm(p_2d, ord=2, dim=1, keepdim=True).clamp_min_(row_lb)
+    col_norms = torch.linalg.vector_norm(p_2d, ord=2, dim=0, keepdim=True).clamp_min_(col_lb)
+
+    # Compute the structural scaler
+    row_factor = row_norms.sqrt_()
+    col_factor = col_norms.sqrt_()
+
+    if row_denom is not None and col_denom is not None:
+        # Reshape denominators to ensure safe in-place broadcasting
+        row_denom = row_denom.view(p_2d.shape[0], 1)
+        col_denom = col_denom.view(1, p_2d.shape[1])
+
+        # High denom (noise) -> smaller angle (protects weights)
+        # Low denom (confident) -> larger angle (decays weights)
+        row_factor.atan2_(row_denom)
+        col_factor.atan2_(col_denom)
+
+    # Outer product: merges the row and column confidences into a 2D matrix
+    wd_scaler = row_factor * col_factor
+
+    # Normalize the scaler so its mean is exactly 1.0
+    wd_scaler.div_(wd_scaler.mean().clamp_min_(1e-12))
+
+    return wd_scaler.view_as(p)

{adv_optm-2.4.dev17 → adv_optm-2.4.dev18}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-2.4.dev17 → adv_optm-2.4.dev18}/setup.py

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