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

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

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

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

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

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

@@ -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, _sinkhorn_sq_grad, get_sinkhorn_wd_scaler
+from ..util.sinkhorn import apply_sr_sinkhorn, get_sinkhorn_wd_scaler
 from ..util.signed_util import apply_stochastic_sign_
 
 class SinkSGD_adv(torch.optim.Optimizer):
@@ -90,6 +90,8 @@ class SinkSGD_adv(torch.optim.Optimizer):
             raise ValueError(f"Momentum should be >= 0.0. Got {momentum}")
         if not (weight_decay >= 0.0):
             raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
+        if centered_vt and not normed_momentum:
+            raise NotImplementedError(f"centered_vt is intended to be used with normed_momentum")
 
         state_precision = state_precision.lower()
         valid_precisions = {"auto", "fp32", "factored", "bf16_sr", "fp8_sr", "int8_sr"}
@@ -183,9 +185,11 @@ class SinkSGD_adv(torch.optim.Optimizer):
                         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)
+                    # Align shapes with Sinkhorn's 2D flattening
+                    dim0 = p.shape[0]
+                    dim1 = p.numel() // dim0
+                    state['vt_row'] = torch.zeros(dim0, device=device, dtype=torch.float32)
+                    state['vt_col'] = torch.zeros(dim1, device=device, dtype=torch.float32)
 
             if group.get('spectral_normalization', False) and is_spectral(p):
                 init_spectral_norm(state, p)
@@ -280,7 +284,7 @@ class SinkSGD_adv(torch.optim.Optimizer):
                 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
+                    grad_vt_sq = grad_vt.mul_(grad_vt).view(grad.shape[0], -1)
                     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)
@@ -289,7 +293,6 @@ class SinkSGD_adv(torch.optim.Optimizer):
                         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
@@ -309,10 +312,11 @@ 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
+            # Align with Sinkhorn: Alternate row/col preconditioning
+            update_2d = update.view(update.shape[0], -1)
+            update_2d.div_(vt_row.clamp_min(1e-30).sqrt().unsqueeze(1))
+            update_2d.div_(vt_col.clamp_min(1e-30).sqrt().unsqueeze(0))
+            update = update_2d.atan_().view_as(p)
 
         if not group.get('normed_momentum', False):
             if not is_vector:

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

@@ -80,23 +80,6 @@ def ortho_normed(p_2d, update_2d, p_norm_sq, dim, target_norm):
     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, 
@@ -126,8 +109,8 @@ def get_sinkhorn_wd_scaler(
 
     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])
+        row_denom = row_denom.sqrt().view(p_2d.shape[0], 1)
+        col_denom = col_denom.sqrt().view(1, p_2d.shape[1])
 
         # High denom (noise) -> smaller angle (protects weights)
         # Low denom (confident) -> larger angle (decays weights)

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

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

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

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