adv-optm 2.4.dev12__tar.gz → 2.4.dev13__tar.gz

{adv_optm-2.4.dev12 → adv_optm-2.4.dev13}/PKG-INFO

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

{adv_optm-2.4.dev12 → adv_optm-2.4.dev13}/adv_optm/__init__.py

@@ -24,4 +24,4 @@ __all__ = [
     "SinkSGD_adv",
 ]
 
-__version__ = "2.4.dev12"
+__version__ = "2.4.dev13"

{adv_optm-2.4.dev12 → adv_optm-2.4.dev13}/adv_optm/optim/AdamW_adv.py

@@ -232,12 +232,13 @@ class AdamW_adv(torch.optim.Optimizer):
     def supports_flat_params(self):
         return False
 
-    @torch.no_grad()
-    def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
-        if p.grad is None:
-            return
+    def init_step(self):
+        for group in self.param_groups:
+            for i, p in enumerate(group['params']):
+                self.__init_state(p, group)
 
-        grad = p.grad
+    @torch.no_grad()
+    def __init_state(self, p, group):
         state = self.state[p]
 
         # State Initialization
@@ -303,6 +304,15 @@ class AdamW_adv(torch.optim.Optimizer):
 
             _init_fisher_wd_scaler(group, state, p)
 
+    @torch.no_grad()
+    def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
+        if p.grad is None:
+            return
+
+        grad = p.grad
+        state = self.state[p]
+        self.__init_state(p, group)
+
         beta1, beta2 = group['betas']
 
         current_step = state['step']

{adv_optm-2.4.dev12 → adv_optm-2.4.dev13}/adv_optm/optim/SinkSGD_adv.py

@@ -116,6 +116,8 @@ class SinkSGD_adv(torch.optim.Optimizer):
             for device in devices:
                 param_update.set_seed(device)
 
+        self.init_step()
+
         self._compiled_step_parameter = None
         if compiled_optimizer:
             self.compile(fullgraph=True)
@@ -136,14 +138,14 @@ class SinkSGD_adv(torch.optim.Optimizer):
     def supports_flat_params(self):
         return False
 
-    @torch.no_grad()
-    def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
-        if p.grad is None:
-            return
+    def init_step(self):
+        for group in self.param_groups:
+            for i, p in enumerate(group['params']):
+                self.__init_state(p, group)
 
-        grad = p.grad
+    @torch.no_grad()
+    def __init_state(self, p, group):
         state = self.state[p]
-
         # State Initialization
         if 'step' not in state:
             state['step'] = 0
@@ -180,6 +182,15 @@ class SinkSGD_adv(torch.optim.Optimizer):
 
             _init_anchor(p, state, group)
 
+    @torch.no_grad()
+    def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
+        if p.grad is None:
+            return
+
+        grad = p.grad
+        state = self.state[p]
+        self.__init_state(p, group)
+
         step_size = group['lr']
 
         random_int_tensor = None
@@ -219,7 +230,7 @@ class SinkSGD_adv(torch.optim.Optimizer):
 
             if momentum != 0:
                 buf = _reconstruct_state((state['mu_b_nmf'], state['mv_b_nmf'], state['sign'], d2), signed=True)
-                buf.mul_(momentum).add_(grad_reshaped, alpha=1 - momentum)
+                buf.lerp_(grad_reshaped, 1 - momentum)
 
                 # Factorize updated buffer
                 state['mu_b_nmf'], state['mv_b_nmf'], state['sign'] = _factorize_state(buf.clone(), signed=True)
@@ -239,9 +250,7 @@ class SinkSGD_adv(torch.optim.Optimizer):
 
             if momentum != 0:
                 buf = get_state(state, 'momentum_buffer', actual_precision)
-
-                buf.mul_(momentum).add_(grad, alpha=1 - momentum)
-
+                buf.lerp_(grad, 1 - momentum)
 
                 set_state(state, 'momentum_buffer', buf, actual_precision, random_int_state_tensor)
 

{adv_optm-2.4.dev12 → adv_optm-2.4.dev13}/adv_optm/util/sinkhorn.py

@@ -36,13 +36,10 @@ def apply_sr_sinkhorn(update: torch.Tensor, p: torch.Tensor, ortho_project: bool
 
 
     # Precompute scaling factors. 
-    scale_first = m if scale_cond else n
-    scale_second = n if scale_cond else m
+    scale_first = math.sqrt(m if scale_cond else n)
+    scale_second = math.sqrt(n if scale_cond else m)
 
     if ortho_project:
-        # Pre-compute squares for the mathematical trick in ortho_normed
-        target_norm_sq_first = scale_first ** 2
-        target_norm_sq_second = scale_second ** 2
         param_2d = p.float().view(p.shape[0], -1)
         p_norm_sq_dim = torch.sum(param_2d * param_2d, dim=dim, keepdim=True).add_(1e-30)
         p_norm_sq_adim = torch.sum(param_2d * param_2d, dim=1-dim, keepdim=True).add_(1e-30)
@@ -53,23 +50,17 @@ def apply_sr_sinkhorn(update: torch.Tensor, p: torch.Tensor, ortho_project: bool
         norm1 = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(1e-12)
         update_2d.mul_(scale_first / norm1)
         if ortho_project:
-            update_2d = ortho_normed(param_2d, update_2d, p_norm_sq_dim, dim, scale_first, target_norm_sq_first)
+            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)
         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, target_norm_sq_second)
-
-    # Final step
-    norm1 = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(1e-12)
-    update_2d.mul_(scale_first / norm1)
-    if ortho_project:
-        update_2d = ortho_normed(param_2d, update_2d, p_norm_sq_dim, dim, scale_first, target_norm_sq_first)
+            update_2d = ortho_normed(param_2d, update_2d, p_norm_sq_adim, 1-dim, scale_second)
 
     return update_2d.view(original_shape).to(original_dtype)
 
-def ortho_normed(p_2d, update_2d, p_norm_sq, dim, target_norm, target_norm_sq):
+def ortho_normed(p_2d, update_2d, p_norm_sq, dim, target_norm):
     """
     Projects the update to be orthogonal to p along 'dim' and restores the original norm.
     """
@@ -80,10 +71,7 @@ def ortho_normed(p_2d, update_2d, p_norm_sq, dim, target_norm, target_norm_sq):
     # In-place subtraction: update_2d = update_2d - (proj * p_2d)
     update_2d.addcmul_(proj, p_2d, value=-1.0)
 
-    # Magnitude Preservation via Pythagorean theorem 
-    # ||g_orth||^2 = ||g||^2 - ||proj * p||^2
-    proj_norm_sq = (dot_prod ** 2) / p_norm_sq
-    g_orth_norm_sq = (target_norm_sq - proj_norm_sq).clamp_min_(1e-30)
-
-    scale_factor = target_norm / torch.sqrt(g_orth_norm_sq)
+    # Magnitude Preservation
+    g_orth_norm = update_2d.norm(p=2, dim=dim, keepdim=True).clamp_min_(1e-12)
+    scale_factor = target_norm / g_orth_norm
     return update_2d.mul_(scale_factor)

{adv_optm-2.4.dev12 → adv_optm-2.4.dev13}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-2.4.dev12 → adv_optm-2.4.dev13}/setup.py

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