adv-optm 1.2.dev6__tar.gz → 1.2.dev8__tar.gz

{adv_optm-1.2.dev6 → adv_optm-1.2.dev8}/PKG-INFO

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

{adv_optm-1.2.dev6 → adv_optm-1.2.dev8}/adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "AdaMuon_adv",
 ]
 
-__version__ = "1.2.dev6"
+__version__ = "1.2.dev8"

{adv_optm-1.2.dev6 → adv_optm-1.2.dev8}/adv_optm/optim/AdaMuon_adv.py

@@ -3,7 +3,6 @@ from typing import Optional, Callable
 
 from .AdamW_adv import AdamW_adv
 from ..util.MuonAdam_helper import MuonAdamHelper
-from ..util.Kourkoutas import KourkoutasHelper
 
 from ..util.BF16_Stochastic_Rounding import add_stochastic_
 from ..util.Newton_Schulz import _newton_schulz_iteration
@@ -64,22 +63,13 @@ class AdaMuon_adv(torch.optim.Optimizer):
             matrices for muon NewtonSchulz (default: False).
         vector_reshape (bool): whether to reshape 1D vectors into 2D
             matrices to apply low-rank compression (default: True).
+        low_rank_ortho (bool): If True, enables low-rank orthogonalization, which
+            projects the update to a lower rank before orthogonalization.
+            (default: False)
+        ortho_rank (int): The rank for low-rank orthogonalization.
+            (default: 128)
         nnmf_factor (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
-        kourkoutas_beta (bool): whether to enable the layer-wise dynamic β₂ logic.
-            If `False`, the optimizer behaves as standard AdamW. (default: False)
-        beta2_min (float): The minimum value for dynamic β₂, used during periods of
-            high gradient variance ("sunspikes"). Must be less than `betas[1]`.
-            (default: 0.88)
-        ema_alpha (float): The decay rate for the Exponential Moving Average (EMA) of
-            the pooled gradient norms. Corresponds to `α` in the paper.
-            (default: 0.93)
-        tiny_spike (float): A small constant added to the denominator of the
-            "sunspike" ratio calculation to prevent division by zero. Corresponds
-            to `ε_spike` in the paper. (default: 1e-9)
-        k_warmup_steps (int): The number of initial steps during which β₂ is held
-            at a fixed beta2 value before the
-            dynamic logic activates. (default: 0)
         MuonWithAuxAdam (bool): If True, enables the hybrid optimizer mode.
             Parameters designated by `layer_key_fn` will be optimized with
             AdamW_adv instead of Muon. (default: False)
@@ -110,15 +100,10 @@ class AdaMuon_adv(torch.optim.Optimizer):
         alpha_grad: float = 100.0,
         vector_reshape_muon: bool = False,
         vector_reshape: bool = False,
+        # Low-rank Muon
+        low_rank_ortho: bool = False,
+        ortho_rank: int = 128,
         nnmf_factor: bool = False,
-        # K-b parameters
-        kourkoutas_beta: bool = False,
-        beta2_min: float = 0.9,
-        ema_alpha: float = 0.95,
-        tiny_spike: float = 1e-9,
-        k_warmup_steps: int = 0,
-        k_logging: int = 0,
-        layer_key_kb_fn: Optional[Callable] = None,
         # hybrid optimizer mode
         MuonWithAuxAdam: bool = False,
         layer_key_fn: Optional[Callable] = None,
@@ -142,14 +127,11 @@ class AdaMuon_adv(torch.optim.Optimizer):
             "vector_reshape": vector_reshape,
             "vector_reshape_muon": vector_reshape_muon,
             "nesterov":nesterov, "use_atan2":use_atan2,
-            "Simplified_AdEMAMix": Simplified_AdEMAMix, "alpha_grad": alpha_grad, 
-            "_kourkoutas_beta": kourkoutas_beta, "beta2_min": beta2_min, "ema_alpha": ema_alpha,
-            "tiny_spike": tiny_spike, "k_warmup_steps": k_warmup_steps, "k_logging": k_logging,
+            "Simplified_AdEMAMix": Simplified_AdEMAMix, "alpha_grad": alpha_grad,
+            # Low-rank Ortho
+            "low_rank_ortho": low_rank_ortho, "ortho_rank": ortho_rank,
         }
         self.stochastic_rounding = stochastic_rounding
-        self._kourkoutas_beta = kourkoutas_beta
-        self._kourkoutas_helper = None
-        self.layer_key_kb_fn = layer_key_kb_fn
         self.MuonWithAuxAdam = MuonWithAuxAdam
         self.helper = None
         self.aux_adam = None
@@ -182,14 +164,9 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
             for key, value in defaults_to_use.items():
                 new_group.setdefault(key, value)
-            if '_kourkoutas_beta' not in new_group:
-                 if optim_type == 'adam':
-                     new_group['_kourkoutas_beta'] = False
-                 else:
-                     new_group['_kourkoutas_beta'] = muon_defaults['_kourkoutas_beta']
             final_param_groups.append(new_group)
 
-        super().__init__(final_param_groups, {})
+        super().__init__(final_param_groups, muon_defaults)
         
         # Now that self is initialized, create the helper
         self.helper = MuonAdamHelper(self, layer_key_fn)
@@ -219,9 +196,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
     @torch.no_grad()
     def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
-        if group['_kourkoutas_beta'] and self._kourkoutas_helper is None:
-            self._kourkoutas_helper = KourkoutasHelper(self)
-
         if self.MuonWithAuxAdam:
             optim_type = self.helper.get_optimizer_type(p)
             if optim_type == 'adam':
@@ -277,7 +251,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
         # Retrieve hyperparameters
         beta1, beta2 = group['betas']
-        current_step = state['step']
         nesterov = group['nesterov']
         Simplified_AdEMAMix = group['Simplified_AdEMAMix']
         alpha_grad = group['alpha_grad']
@@ -303,20 +276,37 @@ class AdaMuon_adv(torch.optim.Optimizer):
                 signed_m_buf = torch.sign(mt_buf)
             del grad_reshaped
 
-            update = _newton_schulz_iteration(
-                signed_m_buf,
-                steps=group['ns_steps'],
-                eps=group['ns_eps'],
-                coeffs=group['ns_coeffs'],
-            )
-
-            if group['_kourkoutas_beta']:
-                # Call prepare_step() once at the beginning of the step for all params
-                self._kourkoutas_helper.maybe_prepare_step(current_step)
-                # Accumulate current sign-stabilized orthogonal update's norm for the *next* step
-                self._kourkoutas_helper.accumulate_gradient_sq_norm(p, update.view(p.shape))
-                # Get the dynamic beta2 calculated in prepare_step()
-                beta2 = self._kourkoutas_helper.get_beta2(p, group, current_step)
+            # Orthogonalization step
+            if group['low_rank_ortho']:
+                # Low-Rank Orthogonalization on the reconstructed matrix
+                M = signed_m_buf
+                r = min(group['ortho_rank'], M.shape[0], M.shape[1])
+                if r > 0:
+                    G_sketch = torch.randn(M.shape[1], r, device=M.device, dtype=M.dtype)
+                    MG = M @ G_sketch
+                    if MG.dtype != torch.float32:
+                        MG_dtype = M.dtype
+                        Q, _ = torch.linalg.qr(MG.float())
+                        Q = Q.to(MG_dtype)
+                    else:
+                        Q, _ = torch.linalg.qr(MG)
+                    projected_M = Q.T @ M
+                    ortho_projected_M = _newton_schulz_iteration(
+                        projected_M, steps=group['ns_steps'], eps=group['ns_eps'], coeffs=group['ns_coeffs']
+                    )
+                    update = Q @ ortho_projected_M
+                else: # Fallback for invalid rank
+                    update = _newton_schulz_iteration(
+                        signed_m_buf, steps=group['ns_steps'], eps=group['ns_eps'], coeffs=group['ns_coeffs']
+                    )
+            else:
+                # Original full Newton-Schulz
+                update = _newton_schulz_iteration(
+                    signed_m_buf,
+                    steps=group['ns_steps'],
+                    eps=group['ns_eps'],
+                    coeffs=group['ns_coeffs'],
+                )
 
             # Reconstruct second momentum from previous step's factors
             vt_buf = _unnmf((state['mu_v_nmf'], state['mv_v_nmf']))
@@ -381,25 +371,41 @@ class AdaMuon_adv(torch.optim.Optimizer):
                 if len(p.shape) > 2:
                     signed_m_buf = signed_m_buf.view(p.shape[0], -1)
 
-                # NewtonSchulz
-                update = _newton_schulz_iteration(
-                    signed_m_buf,
-                    steps=group['ns_steps'],
-                    eps=group['ns_eps'],
-                    coeffs=group['ns_coeffs'],
-                )
+                # Orthogonalization step
+                if group['low_rank_ortho']:
+                    # Low-Rank Orthogonalization on the reconstructed matrix
+                    M = signed_m_buf
+                    r = min(group['ortho_rank'], M.shape[0], M.shape[1])
+                    if r > 0:
+                        G_sketch = torch.randn(M.shape[1], r, device=M.device, dtype=M.dtype)
+                        MG = M @ G_sketch
+                        if MG.dtype != torch.float32:
+                            MG_dtype = M.dtype
+                            Q, _ = torch.linalg.qr(MG.float())
+                            Q = Q.to(MG_dtype)
+                        else:
+                            Q, _ = torch.linalg.qr(MG)
+                        projected_M = Q.T @ M
+                        ortho_projected_M = _newton_schulz_iteration(
+                            projected_M, steps=group['ns_steps'], eps=group['ns_eps'], coeffs=group['ns_coeffs']
+                        )
+                        update = Q @ ortho_projected_M
+                    else: # Fallback for invalid rank
+                        update = _newton_schulz_iteration(
+                            signed_m_buf, steps=group['ns_steps'], eps=group['ns_eps'], coeffs=group['ns_coeffs']
+                        )
+                else:
+                    # Original full Newton-Schulz
+                    update = _newton_schulz_iteration(
+                        signed_m_buf,
+                        steps=group['ns_steps'],
+                        eps=group['ns_eps'],
+                        coeffs=group['ns_coeffs'],
+                    )
 
                 if len(p.shape) > 2 or state['reshaped_1d_muon']:
                     update = update.view(p.shape)
 
-            if group['_kourkoutas_beta']:
-                # Call prepare_step() once at the beginning of the step for all params
-                self._kourkoutas_helper.maybe_prepare_step(current_step)
-                # Accumulate current sign-stabilized orthogonal update's norm for the *next* step
-                self._kourkoutas_helper.accumulate_gradient_sq_norm(p, update)
-                # Get the dynamic beta2 calculated in prepare_step()
-                beta2 = self._kourkoutas_helper.get_beta2(p, group, current_step)
-
                 vt_buf = state['second_momentum_buffer']
                 vt_buf.mul_(beta2).addcmul_(update, update, value=1 - beta2)
 

{adv_optm-1.2.dev6 → adv_optm-1.2.dev8}/adv_optm/optim/Muon_adv.py

@@ -18,6 +18,10 @@ class Muon_adv(torch.optim.Optimizer):
     the hidden layers of neural networks. It applies SGD with momentum and then
     orthogonalizes the resulting update matrix using a Newton-Schulz iteration.
 
+    NorMuon (Neuron-wise Normalized Muon) extends this by adding neuron-level
+    adaptive learning rates, combining the benefits of orthogonalization with
+    second-order momentum statistics.
+
     This implementation is designed for 2D parameters (e.g., linear layers) and
     can handle other-dimensional parameters (e.g., 1D bias, 4D convolutional layers) by
     flattening/reshaping them.
@@ -54,6 +58,19 @@ class Muon_adv(torch.optim.Optimizer):
             matrices to apply low-rank compression (default: True).
         nnmf_factor (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
+        low_rank_ortho (bool): If True, enables low-rank orthogonalization, which
+            projects the update to a lower rank before orthogonalization.
+            (default: False)
+        ortho_rank (int): The rank for low-rank orthogonalization.
+            (default: 128)
+        normuon_variant (bool): If True, enables the NorMuon update rule, which adds
+            neuron-wise normalization. (default: False)
+        beta2_normuon (float): The exponential decay rate for the second moment estimates
+            used in NorMuon. (default: 0.95)
+        normuon_eps (float): Epsilon for NorMuon normalization stability. (default: 1e-8)
+        normuon_lr_scale (float): Scaling factor for the NorMuon learning rate.
+            (default: 0.2)
+        normuon_atan2 (bool): whether to use the atan2 for NorMuon. (default: False)
         MuonWithAuxAdam (bool): If True, enables the hybrid optimizer mode.
             Parameters designated by `layer_key_fn` will be optimized with
             AdamW_adv instead of Muon. (default: False)
@@ -82,6 +99,15 @@ class Muon_adv(torch.optim.Optimizer):
         vector_reshape_muon: bool = False,
         vector_reshape: bool = False,
         nnmf_factor: bool = False,
+        # Low-rank Muon
+        low_rank_ortho: bool = False,
+        ortho_rank: int = 128,
+        # NorMuon additions
+        normuon_variant: bool = False,
+        beta2_normuon: float = 0.95,
+        normuon_eps: float = 1e-8,
+        normuon_lr_scale: float = 0.2,
+        normuon_atan2: bool = False,
         # hybrid optimizer mode
         MuonWithAuxAdam: bool = False,
         layer_key_fn: Optional[Callable] = None,
@@ -92,6 +118,8 @@ class Muon_adv(torch.optim.Optimizer):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
         if not (0.0 <= beta1 < 1.0):
             raise ValueError(f"beta1 should be in [0.0, 1.0). Got {beta1}")
+        if normuon_variant and not (0.0 <= beta2_normuon < 1.0):
+            raise ValueError(f"beta2_normuon should be in [0.0, 1.0) for NorMuon. Got {beta2_normuon}")
         if not (weight_decay >= 0.0):
             raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
         if not (ns_steps > 0):
@@ -106,10 +134,16 @@ class Muon_adv(torch.optim.Optimizer):
             "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor,
             "vector_reshape": vector_reshape,
             "vector_reshape_muon": vector_reshape_muon,
-            "Simplified_AdEMAMix": Simplified_AdEMAMix, "alpha_grad": alpha_grad, 
+            "Simplified_AdEMAMix": Simplified_AdEMAMix, "alpha_grad": alpha_grad,
+            # Low-rank Ortho
+            "low_rank_ortho": low_rank_ortho, "ortho_rank": ortho_rank,
+            # NorMuon
+            "normuon_variant": normuon_variant, "beta2_normuon": beta2_normuon,
+            "normuon_eps": normuon_eps, "normuon_lr_scale": normuon_lr_scale,
+            "normuon_atan2": normuon_atan2,
         }
         self.stochastic_rounding = stochastic_rounding
-        
+
         self.MuonWithAuxAdam = MuonWithAuxAdam
         self.helper = None
         self.aux_adam = None
@@ -144,7 +178,7 @@ class Muon_adv(torch.optim.Optimizer):
 
             final_param_groups.append(new_group)
 
-        super().__init__(final_param_groups, {})
+        super().__init__(final_param_groups, muon_defaults)
         
         # Now that self is initialized, create the helper
         self.helper = MuonAdamHelper(self, layer_key_fn)
@@ -223,6 +257,12 @@ class Muon_adv(torch.optim.Optimizer):
                 elif len(p.shape) == 1:
                     state['momentum_buffer'] = torch.zeros_like(p)
 
+            # NorMuon state initialization
+            if group['normuon_variant']:
+                if len(p.shape) >= 2 or state['reshaped_1d_muon']:
+                    num_rows = p.shape[0] if len(p.shape) >= 2 else state['effective_shape'][0]
+                    state['normuon_v'] = torch.zeros(num_rows, device=p.device, dtype=torch.float32)
+
         beta1 = group['beta1']
         nesterov = group['nesterov']
         Simplified_AdEMAMix = group['Simplified_AdEMAMix']
@@ -251,14 +291,60 @@ class Muon_adv(torch.optim.Optimizer):
                 update = mt_buf.clone()
             del grad_reshaped
 
-            update = _newton_schulz_iteration(
-                update,
-                steps=group['ns_steps'],
-                eps=group['ns_eps'],
-                coeffs=group['ns_coeffs'],
-            )
+            # Orthogonalization step
+            if group['low_rank_ortho']:
+                # Low-Rank Orthogonalization on the reconstructed matrix
+                M = update
+                r = min(group['ortho_rank'], M.shape[0], M.shape[1])
+                if r > 0:
+                    G_sketch = torch.randn(M.shape[1], r, device=M.device, dtype=M.dtype)
+                    MG = M @ G_sketch
+                    if MG.dtype != torch.float32:
+                        MG_dtype = M.dtype
+                        Q, _ = torch.linalg.qr(MG.float())
+                        Q = Q.to(MG_dtype)
+                    else:
+                        Q, _ = torch.linalg.qr(MG)
+                    projected_M = Q.T @ M
+                    ortho_projected_M = _newton_schulz_iteration(
+                        projected_M, steps=group['ns_steps'], eps=group['ns_eps'], coeffs=group['ns_coeffs']
+                    )
+                    update = Q @ ortho_projected_M
+                else: # Fallback for invalid rank
+                    update = _newton_schulz_iteration(
+                        update, steps=group['ns_steps'], eps=group['ns_eps'], coeffs=group['ns_coeffs']
+                    )
+            else:
+                # Original full Newton-Schulz
+                update = _newton_schulz_iteration(
+                    update,
+                    steps=group['ns_steps'],
+                    eps=group['ns_eps'],
+                    coeffs=group['ns_coeffs'],
+                )
 
-            update = update.view(p.shape).mul_(group['lr'])
+
+            if group['normuon_variant'] and 'normuon_v' in state:
+                v_t = state['normuon_v']
+                beta2_normuon = group['beta2_normuon']
+                # Update 2nd moment estimate
+                mean_squared_update = torch.mean(update.square(), dim=1)
+                v_t.mul_(beta2_normuon).add_(mean_squared_update, alpha=1 - beta2_normuon)
+                # Normalize update
+                if group['normuon_atan2']:
+                    a = 1.2732395
+                    update.atan2_(v_t.sqrt().unsqueeze(1)).mul_(a)
+                else:
+                    update.div_(v_t.sqrt().unsqueeze(1).add_(group['normuon_eps']))
+                # Scale learning rate
+                update_norm = torch.linalg.vector_norm(update)
+                if update_norm > 1e-12:
+                    scaled_lr = group['normuon_lr_scale'] * group['lr'] * (p.numel()**0.5) / update_norm
+                else:
+                    scaled_lr = 0.0
+                update = update.view(p.shape).mul_(scaled_lr)
+            else: # Original Muon learning rate application
+                update = update.view(p.shape).mul_(group['lr'])
 
             state['sign'] = _pack_bools(mt_buf > 0)
             _nnmf(mt_buf.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
@@ -298,19 +384,80 @@ class Muon_adv(torch.optim.Optimizer):
                 if len(p.shape) > 2:
                     update = update.view(p.shape[0], -1)
 
-                # NewtonSchulz
-                update = _newton_schulz_iteration(
-                    update,
-                    steps=group['ns_steps'],
-                    eps=group['ns_eps'],
-                    coeffs=group['ns_coeffs'],
-                )
+                # Orthogonalization step
+                if group['low_rank_ortho']:
+                    # Low-Rank Orthogonalization based on Gaussian Sketching
+                    M = update
+                    r = min(group['ortho_rank'], M.shape[0], M.shape[1])
+
+                    if r > 0:
+                        # 1. Sketch the matrix
+                        G_sketch = torch.randn(M.shape[1], r, device=M.device, dtype=M.dtype)
+                        MG = M @ G_sketch
+                        
+                        # 2. QR decomposition to get orthogonal basis Q
+                        if MG.dtype != torch.float32:
+                            MG_dtype = M.dtype
+                            Q, _ = torch.linalg.qr(MG.float())
+                            Q = Q.to(MG_dtype)
+                        else:
+                            Q, _ = torch.linalg.qr(MG)
+                        
+                        # 3. Project M onto the basis
+                        projected_M = Q.T @ M
+                        
+                        # 4. Orthogonalize the smaller projected matrix
+                        ortho_projected_M = _newton_schulz_iteration(
+                            projected_M,
+                            steps=group['ns_steps'],
+                            eps=group['ns_eps'],
+                            coeffs=group['ns_coeffs'],
+                        )
+                        
+                        # 5. Project back to the original space
+                        update = Q @ ortho_projected_M
+                    else: # Fallback for invalid rank
+                        update = _newton_schulz_iteration(
+                            update,
+                            steps=group['ns_steps'],
+                            eps=group['ns_eps'],
+                            coeffs=group['ns_coeffs'],
+                        )
+                else:
+                    # Original NewtonSchulz
+                    update = _newton_schulz_iteration(
+                        update,
+                        steps=group['ns_steps'],
+                        eps=group['ns_eps'],
+                        coeffs=group['ns_coeffs'],
+                    )
+
+                # NorMuon Logic
+                if group['normuon_variant'] and 'normuon_v' in state:
+                    v_t = state['normuon_v']
+                    beta2_normuon = group['beta2_normuon']
+                    # Update 2nd moment estimate
+                    mean_squared_update = torch.mean(update.square(), dim=1)
+                    v_t.mul_(beta2_normuon).add_(mean_squared_update, alpha=1 - beta2_normuon)
+                    # Normalize update
+                    if group['normuon_atan2']:
+                        a = 1.2732395
+                        update.atan2_(v_t.sqrt().unsqueeze(1)).mul_(a)
+                    else:
+                        update.div_(v_t.sqrt().unsqueeze(1).add_(group['normuon_eps']))
+                    # Scale learning rate
+                    update_norm = torch.linalg.vector_norm(update)
+                    if update_norm > 1e-12:
+                        scaled_lr = group['normuon_lr_scale'] * group['lr'] * (p.numel()**0.5) / update_norm
+                    else:
+                        scaled_lr = 0.0
+                    update.mul_(scaled_lr)
+                else: # Original Muon learning rate application
+                    update.mul_(group['lr'])
 
                 # Reshape back to original if we flattened or reshaped
                 if len(p.shape) > 2 or state['reshaped_1d_muon']:
                     update = update.view(p.shape)
-
-                update.mul_(group['lr'])
             
             else: # Fallback to standard SGD with momentum for 1D params (biases, etc.) when not reshaped
                 # Momentum update

{adv_optm-1.2.dev6 → adv_optm-1.2.dev8}/adv_optm/util/Kourkoutas.py

@@ -94,7 +94,7 @@ class KourkoutasHelper:
         for layer_key, info in self.layer_info.items():
             params, group = info['params'], info['group_ref']
 
-            if not group.get('kourkoutas_beta', False):
+            if not group.get('kourkoutas_beta', False) and not group.get('_kourkoutas_beta', False):
                 continue
 
             first_param_in_layer = info['params'][0]

{adv_optm-1.2.dev6 → adv_optm-1.2.dev8}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-1.2.dev6 → adv_optm-1.2.dev8}/setup.py

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