adv-optm 2.4.dev4__tar.gz → 2.4.dev5__tar.gz

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/PKG-INFO

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

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm/__init__.py

@@ -22,4 +22,4 @@ __all__ = [
     "SignSGD_adv",
 ]
 
-__version__ = "2.4.dev4"
+__version__ = "2.4.dev5"

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm/optim/AdaMuon_adv.py

@@ -206,6 +206,8 @@ class AdaMuon_adv(torch.optim.Optimizer):
         if spectral_normalization and rms_rescaling:
             print("Warning: spectral_normalization is incompatible with rms_rescaling, Disabling rms_rescaling.")
             rms_rescaling = False
+        if spectral_normalization and accelerated_ns:
+            ValueError("spectral_normalization violates accelerated Newton-Schulz assumptions. Pick one of them.")
 
         defaults = {
             "lr": lr, "betas": betas, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
@@ -260,6 +262,8 @@ class AdaMuon_adv(torch.optim.Optimizer):
             if group.get('use_muon') is None: # Fallback
                  group['use_muon'] = group.get('optim_type') == 'muon'
 
+        self.init_step()
+
         self.kourkoutas_helper = None
         if any(group.get('adam_kourkoutas_beta', False) for group in self.param_groups):
             self.kourkoutas_helper = KourkoutasHelper(self)
@@ -419,7 +423,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
             step_size = group['lr'] / bias_correction1
 
             if is_compiled:
-                step_size = torch.as_tensor(step_size, dtype=torch.float64)
+                step_size = torch.as_tensor(step_size)
                 adam_step_param = self._compiled_adam_step_parameter
             else:
                 adam_step_param = Muon_AuxAdam._adam_step_parameter
@@ -430,7 +434,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
         else: # Muon path
             if is_compiled:
-                lr = torch.as_tensor(group['lr'], dtype=torch.float64)
+                lr = torch.as_tensor(group['lr'])
                 muon_step_param = self._compiled_muon_step_parameter
             else:
                 lr = group['lr']
@@ -467,7 +471,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
             else:
                 shape_for_scaling = p.shape
 
-            scaled_eps, adaptive_eps, spectral_target, wd_scale = get_spectral_scaling(shape_for_scaling, group['n_layers'])
+            scaled_eps, adaptive_eps, spectral_target, wd_scale = get_spectral_scaling(p, shape_for_scaling, group['n_layers'])
 
             weight_decay = group['weight_decay'] * wd_scale
             decoupled_wd = True

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm/optim/AdamW_adv.py

@@ -6,7 +6,7 @@ from typing import Optional, Callable
 
 from ..util import param_update
 from ..util.factorization_util import _get_effective_shape, _reconstruct_state, _factorize_state
-from ..util.update_util import _grams_update, _cautious_update
+from ..util.update_util import _grams_update, _cautious_update, _init_fisher_wd_scaler, _get_fisher_wd_scaler
 from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
 from ..util.scaled_optm import scale_update, is_spectral, init_spectral_norm
@@ -29,6 +29,9 @@ class AdamW_adv(torch.optim.Optimizer):
         eps (float): term added to the denominator to improve
             numerical stability (default: 1e-8)
         weight_decay (float): weight decay (L2 penalty) (default: 0).
+        fisher_wd (bool): whether to use Fisher Adam (FAdam) weight decay, mapping
+            the decay direction through the empirical Fisher information matrix and
+            clipping its RMS. (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).
@@ -103,6 +106,7 @@ class AdamW_adv(torch.optim.Optimizer):
         eps: float = 1e-8,
         # Decoupled/cautious weight decay
         weight_decay: float = 0.0,
+        fisher_wd: bool = False,
         cautious_wd: bool = False,
         # Adam's Bias Correction
         use_bias_correction: bool = True,
@@ -155,7 +159,8 @@ class AdamW_adv(torch.optim.Optimizer):
             cautious_mask = False
 
         defaults = {
-            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
+            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
+            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd,
             "use_atan2": use_atan2,
             "orthogonal_gradient": orthogonal_gradient, "use_bias_correction": use_bias_correction,
             "beta3_ema": beta3_ema, "alpha": alpha, "compiled_optimizer": compiled_optimizer,
@@ -273,6 +278,8 @@ class AdamW_adv(torch.optim.Optimizer):
 
             _init_anchor(p, state, group)
 
+            _init_fisher_wd_scaler(group, state, p)
+
         beta1, beta2 = group['betas']
 
         current_step = state['step']
@@ -294,7 +301,7 @@ class AdamW_adv(torch.optim.Optimizer):
         random_int_tensor = None
 
         if group.get('compiled_optimizer', False):
-            step_size = torch.as_tensor(step_size, dtype=torch.float64)
+            step_size = torch.as_tensor(step_size)
             if p.dtype == torch.bfloat16 and self.stochastic_rounding:
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
@@ -389,6 +396,9 @@ class AdamW_adv(torch.optim.Optimizer):
                 denom = vt.sqrt_()
                 denom.div_(sqrt_bias_correction2).add_(group['eps'])
                 update.div_(denom)
+
+            wd_scaler = _get_fisher_wd_scaler(group, state.get("wd_scaler"), p, denom, group['use_atan2'])
+
             del vt
 
             update = update.view(p.shape)
@@ -430,6 +440,9 @@ class AdamW_adv(torch.optim.Optimizer):
                 denom = exp_avg_sq.sqrt()
                 denom.div_(sqrt_bias_correction2).add_(group['eps'])
                 update.div_(denom)
+
+            wd_scaler = _get_fisher_wd_scaler(group, state.get("wd_scaler"), p, denom, group['use_atan2'])
+
             del denom
 
         update_scaling = step_size * A if group['use_atan2'] else step_size
@@ -438,7 +451,7 @@ class AdamW_adv(torch.optim.Optimizer):
         else:
             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.dev4 → adv_optm-2.4.dev5}/adv_optm/optim/Adopt_adv.py

@@ -7,7 +7,7 @@ from ..util import param_update
 from ..util.factorization_util import _get_effective_shape, _reconstruct_state, _factorize_state, _nnmf
 from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
-from ..util.update_util import _grams_update, _cautious_update, _scale_sim_AdEMAMix_update
+from ..util.update_util import _grams_update, _cautious_update, _scale_sim_AdEMAMix_update, _init_fisher_wd_scaler, _get_fisher_wd_scaler
 from ..util.scaled_optm import scale_update, is_spectral, init_spectral_norm
 from ..util.centered_decay import _init_anchor
 
@@ -33,6 +33,9 @@ class Adopt_adv(torch.optim.Optimizer):
         eps (float): term added to the denominator to improve
             numerical stability (default: 1e-6)
         weight_decay (float): weight decay (L2 penalty) (default: 0)
+        fisher_wd (bool): whether to use Fisher Adam (FAdam) weight decay, mapping
+            the decay direction through the empirical Fisher information matrix and
+            clipping its RMS. (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).
@@ -119,6 +122,7 @@ class Adopt_adv(torch.optim.Optimizer):
         eps: float = 1e-6,
         # Decoupled/cautious weight decay
         weight_decay: float = 0.0,
+        fisher_wd: bool = False,
         cautious_wd: bool = False,
         # ADOPT clipping
         clip_lambda: Optional[Callable[[int], float]] = lambda step: step**0.25,
@@ -181,7 +185,8 @@ class Adopt_adv(torch.optim.Optimizer):
             print("Warning: cautious is incompatible with Simplified_AdEMAMix, Disabling cautious.")
 
         defaults = {
-            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
+            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
+            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd,
             "beta3_ema": beta3_ema, "alpha": alpha,
             "alpha_grad": alpha_grad,
             "kourkoutas_beta": kourkoutas_beta, "beta2_min": beta2_min, "ema_alpha": ema_alpha,
@@ -312,6 +317,8 @@ class Adopt_adv(torch.optim.Optimizer):
 
             _init_anchor(p, state, group)
 
+            _init_fisher_wd_scaler(group, state, p)
+
         current_step = state['step']
 
         # The first step is for initialization only (skip when use_atan2 as it's scale invariant).
@@ -322,7 +329,7 @@ class Adopt_adv(torch.optim.Optimizer):
         random_int_tensor = None
 
         if group.get('compiled_optimizer', False):
-            lr = torch.as_tensor(group['lr'], dtype=torch.float64)
+            lr = torch.as_tensor(group['lr'])
             if p.dtype == torch.bfloat16 and self.stochastic_rounding:
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
@@ -366,6 +373,7 @@ class Adopt_adv(torch.optim.Optimizer):
 
             # ADOPT Step A: Decorrelate g_t using v_{t-1}
             denom = vt.sqrt()
+            wd_scaler = _get_fisher_wd_scaler(group, state.get("wd_scaler"), p, denom, self.use_atan2)
 
             # Update second moment v_t for the *next* step using raw g_t
             if isinstance(beta2, torch.Tensor) and beta2.dim() > 0:
@@ -444,6 +452,7 @@ class Adopt_adv(torch.optim.Optimizer):
 
             # ADOPT Step A: Decorrelate g_t using v_{t-1}
             denom = vt.sqrt()
+            wd_scaler = _get_fisher_wd_scaler(group, state.get("wd_scaler"), p, denom, self.use_atan2)
 
             if self.use_atan2:
                 normalized_grad = torch.atan2(grad, denom, out=denom)
@@ -499,7 +508,7 @@ class Adopt_adv(torch.optim.Optimizer):
             update.mul_(update_scaling)
 
         # Parameter Update
-        param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)
+        param_update.apply_parameter_update(self, p, group, update, lr, 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.dev4 → adv_optm-2.4.dev5}/adv_optm/optim/Lion_Prodigy_adv.py

@@ -225,8 +225,8 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
             # TODO, workaround until pytorch#169634 is fixed
-            d = torch.as_tensor(group['d'], dtype=torch.float64)
-            dlr = torch.as_tensor(dlr, dtype=torch.float64)
+            d = torch.as_tensor(group['d'])
+            dlr = torch.as_tensor(dlr)
             step_param_fn = self._compiled_step_parameter
         else:
             d = group['d']

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm/optim/Lion_adv.py

@@ -9,6 +9,7 @@ from ..util.lion_k import _get_lion_k_update
 from ..util.scaled_optm import scale_update, is_spectral, init_spectral_norm
 from ..util.centered_decay import _init_anchor
 from ..util.update_util import _get_l1_adaptive_lr
+from ..util.signed_util import apply_stochastic_sign
 
 
 class Lion_adv(torch.optim.Optimizer):
@@ -45,6 +46,7 @@ class Lion_adv(torch.optim.Optimizer):
             parameter dimensionality. Sets p=2.0 for 4D tensors (Conv2D) (Biases/Norms) to
             use Spherical updates, and p=1.0 for others (Linear/Embeddings) to use Sign
             updates. Overrides explicit kappa_p value. (default: False).
+        stochastic_sign (bool): whether to use the Stochastic Sign operator. (default: False)
         freeze_on_flip (bool): Projected SignGD One-hit freeze. Masks updates for
             coordinates where the gradient sign flips compared to the previous step. (default: False)
         l1_adaptive (bool): Scales learning rate dynamically
@@ -80,6 +82,8 @@ class Lion_adv(torch.optim.Optimizer):
         # Lion-k
         kappa_p: float = 1.0,
         auto_kappa_p: bool = False,
+        # Stochastic Sign Operator
+        stochastic_sign: bool = False,
         # Projected and adaptive sign
         freeze_on_flip: bool = False,
         l1_adaptive: bool = False,
@@ -111,6 +115,7 @@ class Lion_adv(torch.optim.Optimizer):
             clip_threshold=clip_threshold,
             kappa_p=kappa_p,
             auto_kappa_p=auto_kappa_p,
+            stochastic_sign=stochastic_sign,
             freeze_on_flip=freeze_on_flip,
             l1_adaptive=l1_adaptive,
             scaled_optm= scaled_optm,
@@ -202,19 +207,22 @@ class Lion_adv(torch.optim.Optimizer):
         lr = group["lr"]
 
         random_int_tensor = None
+        random_noise_tensor = None
 
         if group.get('compiled_optimizer', False):
             if p.dtype == torch.bfloat16 and self.stochastic_rounding:
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
-            lr = torch.as_tensor(lr, dtype=torch.float64)
+            if group.get('stochastic_sign', False):
+                random_noise_tensor = param_update._get_random_noise_for_sso(p)
+            lr = torch.as_tensor(lr)
             step_param_fn = self._compiled_step_parameter
         else:
             step_param_fn = self._step_parameter
 
-        step_param_fn(p, grad, state, group, lr, random_int_tensor)
+        step_param_fn(p, grad, state, group, lr, random_int_tensor, random_noise_tensor)
 
-    def _step_parameter(self, p, grad, state, group, lr, random_int_tensor):
+    def _step_parameter(self, p, grad, state, group, lr, random_int_tensor, random_noise_tensor):
         if grad.dtype != torch.float32 and state['factored']:
             grad = grad.float()
         if group["clip_threshold"] > 0.0:
@@ -252,8 +260,6 @@ class Lion_adv(torch.optim.Optimizer):
             # Compute update term c_t
             update = torch.lerp(grad_reshaped, exp_avg, beta1)
 
-            l1_mean = _get_l1_adaptive_lr(p, update, state, group, kappa_p)
-
             # Standard Lion momentum update
             # m_t = beta2 * m_{t-1} + (1-beta2) * g_t
             exp_avg.lerp_(grad_reshaped, 1 - beta2)
@@ -262,7 +268,12 @@ class Lion_adv(torch.optim.Optimizer):
             state['mu_m_nmf'], state['mv_m_nmf'], state['sign'] = _factorize_state(exp_avg, signed=True)
             del exp_avg
 
-            update = _get_lion_k_update(update, kappa_p)
+            if freeze_on_flip:
+                # Fast binary diff (XOR) from momentum sign directly
+                flipped_packed = prev_sign_packed ^ state['sign']
+                flipped_mask = _unpack_bools(flipped_packed, original_m=d2).view_as(update)
+                update = torch.where(flipped_mask, 0.0, update)
+                del prev_sign_packed, flipped_packed, flipped_mask
 
             if self.cautious_mask:
                 mask = (update * grad_reshaped > 0).to(grad_reshaped.dtype)
@@ -272,12 +283,12 @@ class Lion_adv(torch.optim.Optimizer):
 
             update = update.view(p.shape)
 
-            if freeze_on_flip:
-                # Fast binary diff (XOR) from momentum sign directly
-                flipped_packed = prev_sign_packed ^ state['sign']
-                flipped_mask = _unpack_bools(flipped_packed, original_m=d2).view_as(update)
-                update = torch.where(flipped_mask, 0.0, update)
-                del prev_sign_packed, flipped_packed, flipped_mask
+            l1_mean = _get_l1_adaptive_lr(p, update, state, group, kappa_p, rescale=False)
+
+            if group.get('stochastic_sign', False):
+                update = apply_stochastic_sign(update, noise=random_noise_tensor)
+            else:
+                update = _get_lion_k_update(update, kappa_p)
 
         else:
             # Fallback to standard Lion logic
@@ -286,9 +297,13 @@ class Lion_adv(torch.optim.Optimizer):
             # Compute update term
             update = torch.lerp(grad, exp_avg, beta1)
 
-            l1_mean = _get_l1_adaptive_lr(p, update, state, group, kappa_p)
+            # Standard Lion momentum update
+            exp_avg.lerp_(grad, 1 - beta2)
 
-            update = _get_lion_k_update(update, kappa_p)
+            if freeze_on_flip:
+                current_sign = (update > 0).to(torch.uint8)
+                update = torch.where(current_sign == state['prev_sign'], update, 0.0)
+                state['prev_sign'] = current_sign
 
             if self.cautious_mask:
                 mask = (update * grad > 0).to(grad.dtype)
@@ -296,13 +311,12 @@ class Lion_adv(torch.optim.Optimizer):
                 update.mul_(mask)
                 del mask
 
-            # Standard Lion momentum update
-            exp_avg.lerp_(grad, 1 - beta2)
+            l1_mean = _get_l1_adaptive_lr(p, update, state, group, kappa_p, rescale=False)
 
-            if freeze_on_flip:
-                current_sign = (update > 0).to(torch.uint8)
-                update = torch.where(current_sign == state['prev_sign'], update, 0.0)
-                state['prev_sign'] = current_sign
+            if group.get('stochastic_sign', False):
+                update = apply_stochastic_sign(update, noise=random_noise_tensor)
+            else:
+                update = _get_lion_k_update(update, kappa_p)
 
         if l1_mean is not None:
             update.mul_(l1_mean)
@@ -312,7 +326,7 @@ class Lion_adv(torch.optim.Optimizer):
         else:
             update.mul_(lr)
 
-        param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)
+        param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor, wd_scaler=l1_mean)
 
     def compile(self, *args, **kwargs):
         self._compiled_step_parameter = torch.compile(self._step_parameter, *args, **kwargs)

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm/optim/Muon_adv.py

@@ -183,6 +183,8 @@ class Muon_adv(torch.optim.Optimizer):
         if spectral_normalization and rms_rescaling:
             print("Warning: spectral_normalization is incompatible with rms_rescaling, Disabling rms_rescaling.")
             rms_rescaling = False
+        if spectral_normalization and accelerated_ns:
+            ValueError("spectral_normalization violates accelerated Newton-Schulz assumptions. Pick one of them.")
 
         defaults = {
             "lr": lr, "beta1": beta1, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
@@ -239,6 +241,8 @@ class Muon_adv(torch.optim.Optimizer):
             if group.get('use_muon') is None: # Fallback
                  group['use_muon'] = group.get('optim_type') == 'muon'
 
+        self.init_step()
+
         self.kourkoutas_helper = None
         if any(group.get('adam_kourkoutas_beta', False) for group in self.param_groups):
             self.kourkoutas_helper = KourkoutasHelper(self)
@@ -393,7 +397,7 @@ class Muon_adv(torch.optim.Optimizer):
             step_size = group['lr'] / bias_correction1
 
             if is_compiled:
-                step_size = torch.as_tensor(step_size, dtype=torch.float64)
+                step_size = torch.as_tensor(step_size)
                 adam_step_param = self._compiled_adam_step_parameter
             else:
                 adam_step_param = Muon_AuxAdam._adam_step_parameter
@@ -404,7 +408,7 @@ class Muon_adv(torch.optim.Optimizer):
 
         else: # Muon path
             if is_compiled:
-                lr = torch.as_tensor(group['lr'], dtype=torch.float64)
+                lr = torch.as_tensor(group['lr'])
                 muon_step_param = self._compiled_muon_step_parameter
             else:
                 lr = group['lr']

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm/optim/Prodigy_adv.py

@@ -9,7 +9,7 @@ from ..util import param_update
 from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
 from ..util.factorization_util import _get_effective_shape, _reconstruct_state, _factorize_state
-from ..util.update_util import _grams_update, _cautious_update, _scale_sim_AdEMAMix_update
+from ..util.update_util import _grams_update, _cautious_update, _scale_sim_AdEMAMix_update, _init_fisher_wd_scaler, _get_fisher_wd_scaler
 from ..util.centered_decay import _init_anchor
 
 A = 4 / math.pi
@@ -29,6 +29,9 @@ class Prodigy_adv(torch.optim.Optimizer):
         eps (float): term added to the denominator to improve
             numerical stability (default: 1e-8)
         weight_decay (float): weight decay (L2 penalty) (default: 0)
+        fisher_wd (bool): whether to use Fisher Adam (FAdam) weight decay, mapping
+            the decay direction through the empirical Fisher information matrix and
+            clipping its RMS. (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).
@@ -133,6 +136,7 @@ class Prodigy_adv(torch.optim.Optimizer):
         eps: float = 1e-8,
         # Decoupled/cautious weight decay
         weight_decay: float = 0.0,
+        fisher_wd: bool = False,
         cautious_wd: bool = False,
         # Stochastic Rounding for BF16
         stochastic_rounding: bool = True,
@@ -206,7 +210,8 @@ class Prodigy_adv(torch.optim.Optimizer):
             raise ValueError(f"For Kourkoutas-β, betas[1] (as beta2_max) must be > beta2_min. Got {betas[1]} and {beta2_min}")
 
         defaults = {
-            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
+            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
+            "fisher_wd": fisher_wd, "cautious_wd": cautious_wd,
             "use_atan2": use_atan2,
             "orthogonal_gradient": orthogonal_gradient,
             "beta3_ema": beta3_ema, "alpha": alpha, "compiled_optimizer": compiled_optimizer,
@@ -354,6 +359,8 @@ class Prodigy_adv(torch.optim.Optimizer):
 
             _init_anchor(p, state, group)
 
+            _init_fisher_wd_scaler(group, state, p)
+
         if not hasattr(self, 'd_denom'):
             self.d_denom = torch.tensor(0.0, device=p.device)
             self.d_numerator = torch.tensor(group.get('d_numerator', 0.0), device=p.device)
@@ -376,8 +383,8 @@ class Prodigy_adv(torch.optim.Optimizer):
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
             # TODO, workaround until pytorch#169634 is fixed
-            d = torch.as_tensor(group['d'], dtype=torch.float64)
-            dlr = torch.as_tensor(dlr, dtype=torch.float64)
+            d = torch.as_tensor(group['d'])
+            dlr = torch.as_tensor(dlr)
             step_param_fn = self._compiled_step_parameter
         else:
             d = group['d']
@@ -478,6 +485,7 @@ class Prodigy_adv(torch.optim.Optimizer):
             else:
                 denom = vt.sqrt_()
                 update.div_(denom.add_(d * group['eps']))
+            wd_scaler = _get_fisher_wd_scaler(group, state.get("wd_scaler"), p, denom, group['use_atan2'])
             del vt
 
             update_scaling = dlr * A if group['use_atan2'] else dlr
@@ -528,6 +536,7 @@ class Prodigy_adv(torch.optim.Optimizer):
             else:
                 denom = exp_avg_sq.sqrt()
                 update.div_(denom.add_(d * group['eps']))
+            wd_scaler = _get_fisher_wd_scaler(group, state.get("wd_scaler"), p, denom, group['use_atan2'])
             del denom
 
             update_scaling = dlr * A if group['use_atan2'] else dlr
@@ -557,7 +566,7 @@ class Prodigy_adv(torch.optim.Optimizer):
             if 'p0' in state:
                 del state['p0']
 
-        param_update.apply_parameter_update(self, p, group, update, dlr, random_int_tensor=random_int_tensor)
+        param_update.apply_parameter_update(self, p, group, update, dlr, 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.dev4 → adv_optm-2.4.dev5}/adv_optm/optim/SignSGD_adv.py

@@ -9,6 +9,7 @@ from ..util.lion_k import _get_lion_k_update
 from ..util.update_util import _get_l1_adaptive_lr
 from ..util.scaled_optm import scale_update, is_spectral, init_spectral_norm
 from ..util.centered_decay import _init_anchor
+from ..util.signed_util import apply_stochastic_sign
 
 
 class SignSGD_adv(torch.optim.Optimizer):
@@ -39,6 +40,7 @@ class SignSGD_adv(torch.optim.Optimizer):
             parameter dimensionality. Sets p=2.0 for 4D tensors (Conv2D) (Biases/Norms) to
             use Spherical updates, and p=1.0 for others (Linear/Embeddings) to use Sign
             updates. Overrides explicit kappa_p value. (default: False).
+        stochastic_sign (bool): whether to use the Stochastic Sign operator. (default: False)
         Simplified_AdEMAMix (bool): whether to use the Simplified AdEMAMix update rule.
             This changes the EMA to accumulator and the update numerator to `alpha_grad * grad + mt`, which can be
             more responsive, especially for small batch sizes. (default: False)
@@ -79,6 +81,8 @@ class SignSGD_adv(torch.optim.Optimizer):
         # Projection-k
         kappa_p: float = 1.0,
         auto_kappa_p: bool = True,
+        # Stochastic Sign Operator
+        stochastic_sign: bool = False,
         # Simplified_AdEMAMix
         alpha_grad: float = 1.0,
         Simplified_AdEMAMix: bool = False,
@@ -112,6 +116,7 @@ class SignSGD_adv(torch.optim.Optimizer):
             orthogonal_gradient=orthogonal_gradient,
             kappa_p=kappa_p,
             auto_kappa_p=auto_kappa_p,
+            stochastic_sign=stochastic_sign,
             alpha_grad=alpha_grad,
             Simplified_AdEMAMix=Simplified_AdEMAMix,
             scaled_optm= scaled_optm,
@@ -203,23 +208,26 @@ class SignSGD_adv(torch.optim.Optimizer):
         lr = group["lr"]
 
         random_int_tensor = None
+        random_noise_tensor = None
 
         if group.get('compiled_optimizer', False):
             if p.dtype == torch.bfloat16 and self.stochastic_rounding:
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
-            lr = torch.as_tensor(lr, dtype=torch.float64)
+            if group.get('stochastic_sign', False):
+                random_noise_tensor = param_update._get_random_noise_for_sso(p)
+            lr = torch.as_tensor(lr)
             step_param_fn = self._compiled_step_parameter
         else:
             step_param_fn = self._step_parameter
 
-        step_param_fn(p, grad, state, group, lr, random_int_tensor)
+        step_param_fn(p, grad, state, group, lr, random_int_tensor, random_noise_tensor)
 
         if group.get("l1_adaptive", False):
             state["step"] += 1
 
-    def _step_parameter(self, p, grad, state, group, lr, random_int_tensor):
-        if grad.dtype != torch.float32 and state['factored']:
+    def _step_parameter(self, p, grad, state, group, lr, random_int_tensor, random_noise_tensor):
+        if grad.dtype != torch.float32 and state.get('factored', False):
             grad = grad.float()
 
         if group["orthogonal_gradient"]:
@@ -269,18 +277,23 @@ class SignSGD_adv(torch.optim.Optimizer):
                 if freeze_on_flip:
                     state['sign'] = _pack_bools(raw_update > 0)
 
-            l1_mean = _get_l1_adaptive_lr(p, raw_update, state, group, kappa_p)
 
-            update = _get_lion_k_update(raw_update, kappa_p)
-            update = update.view(p.shape)
+            raw_update = raw_update.view(p.shape)
 
             if freeze_on_flip:
                 # Fast binary diff (XOR) from momentum sign directly
                 flipped_packed = prev_sign_packed ^ state['sign']
-                flipped_mask = _unpack_bools(flipped_packed, original_m=d2).view_as(update)
-                update = torch.where(flipped_mask, 0.0, update)
+                flipped_mask = _unpack_bools(flipped_packed, original_m=d2).view_as(raw_update)
+                raw_update = torch.where(flipped_mask, 0.0, raw_update)
                 del prev_sign_packed, flipped_packed, flipped_mask
 
+            l1_mean = _get_l1_adaptive_lr(p, raw_update, state, group, kappa_p)
+
+            if group.get('stochastic_sign', False):
+                update = apply_stochastic_sign(raw_update, noise=random_noise_tensor)
+            else:
+                update = _get_lion_k_update(raw_update, kappa_p)
+
         else:
             # Fallback to standard SignSGD logic
             if momentum > 0:
@@ -294,15 +307,18 @@ class SignSGD_adv(torch.optim.Optimizer):
             else:
                 raw_update = grad.clone()
 
-            l1_mean = _get_l1_adaptive_lr(p, raw_update, state, group, kappa_p)
-
-            update = _get_lion_k_update(raw_update, kappa_p)
-
             if freeze_on_flip:
                 current_sign = (raw_update > 0).to(torch.uint8)
-                update = torch.where(current_sign == state['prev_sign'], update, 0.0)
+                raw_update = torch.where(current_sign == state['prev_sign'], raw_update, 0.0)
                 state['prev_sign'] = current_sign
 
+            l1_mean = _get_l1_adaptive_lr(p, raw_update, state, group, kappa_p)
+
+            if group.get('stochastic_sign', False):
+                update = apply_stochastic_sign(raw_update, noise=random_noise_tensor)
+            else:
+                update = _get_lion_k_update(raw_update, kappa_p)
+
         if l1_mean is not None:
             update.mul_(l1_mean)
 
@@ -311,7 +327,7 @@ class SignSGD_adv(torch.optim.Optimizer):
         else:
             update.mul_(lr)
 
-        param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)
+        param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor, wd_scaler=l1_mean)
 
     def compile(self, *args, **kwargs):
         self._compiled_step_parameter = torch.compile(self._step_parameter, *args, **kwargs)

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm/optim/Simplified_AdEMAMix.py

@@ -288,7 +288,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
             # TODO, workaround until pytorch#169634 is fixed
-            lr = torch.as_tensor(lr, dtype=torch.float64)
+            lr = torch.as_tensor(lr)
             step_param_fn = self._compiled_step_parameter
         else:
             step_param_fn = self._step_parameter

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm/util/Muon_util.py

@@ -1,5 +1,7 @@
 import torch
 
+import math
+
 @torch.no_grad()
 def _newton_schulz_iteration(
     G: torch.Tensor,
@@ -359,11 +361,11 @@ def rms_adjustment(update: torch.Tensor, rms_rescaling: bool, lr):
         # This is slower due to norm calculations but it worked the best for t2i models.
         rms_target = 0.2 # default (Adam) value for RMS
         update_norm = torch.linalg.vector_norm(update)
-        return update.mul_(lr * rms_target * (update.numel()**0.5) / update_norm.clamp_min_(1e-8))
+        return update.mul_(lr * rms_target * (math.sqrt(update.numel())) / update_norm.clamp_min_(1e-8))
     else:
         # Original Muon scaling
         r, c = update.size(-2), update.size(-1)
-        scaling_factor = max(1, r / c) ** 0.5
+        scaling_factor = math.sqrt(max(1, r / c))
         return update.mul_(lr * scaling_factor)
 
 def _auto_projection_for_adamuon(raw_update: torch.Tensor, kappa_p: float) -> torch.Tensor:
@@ -474,15 +476,15 @@ def get_spectral_scaling(shape: torch.Size, n_layers: int):
     # A) Newton-Schulz Damping
     # This ensures the matrix orthogonalization is stable across scales.
     # Formula: (1/L) * sqrt(d_in / d_out)
-    ns_eps = (1.0 / L) * (d_in / d_out) ** 0.5
+    ns_eps = (1.0 / L) * math.sqrt(d_in / d_out)
 
     # B) Adaptive Denominator Epsilon
     # This ensures the Adam-style division doesn't explode or vanish.
     # Formula: (1/L) * (1 / sqrt(d_in * d_out))
-    adaptive_eps = (1.0 / L) * (1.0 / (d_in * d_out)**0.5)
+    adaptive_eps = (1.0 / L) * (1.0 / math.sqrt(d_in * d_out))
 
     # Spectral Target (Section F) -> sqrt(d_out/d_in)
-    spectral_target = (d_out / d_in) ** 0.5
+    spectral_target = math.sqrt(d_out / d_in)
 
     # Weight Decay (Section 3.4) -> 1/width
     wd_scale = 1.0 / d_in

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm/util/param_update.py

@@ -4,7 +4,7 @@ from torch.optim import Optimizer
 
 from typing import Dict, Any
 
-from .scaled_optm import scale_wds
+from .scaled_optm import adjust_wds, scale_wds
 from .centered_decay import dequantize_anchor
 
 _generators: Dict[torch.device, torch.Generator] = {}
@@ -29,11 +29,17 @@ def _apply_weight_decay(
         # 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)
-            p_calc.addcmul_(p_calc, mask, value=-scaled_wd)
+            if isinstance(scaled_wd, Tensor):
+                p_calc.addcmul_(p_calc, mask * scaled_wd, value=-1.0)
+            else:
+                p_calc.addcmul_(p_calc, mask, value=-scaled_wd)
             del mask
         else:
             # Standard decoupled weight decay
-            p_calc.add_(p_calc, alpha=-scaled_wd)
+            if isinstance(scaled_wd, Tensor):
+                p_calc.addcmul_(p_calc, scaled_wd, value=-1.0)
+            else:
+                p_calc.add_(p_calc, alpha=-scaled_wd)
 
     # Centered Weight Decay (pulls toward anchor)
     if scaled_cwd is not None and 'anchor_type' in state:
@@ -43,15 +49,20 @@ 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)
-            p_calc.addcmul_(decay_target, mask, value=-scaled_cwd)
+            if isinstance(scaled_cwd, Tensor):
+                p_calc.addcmul_(decay_target, mask * scaled_cwd, value=-1.0)
+            else:
+                p_calc.addcmul_(decay_target, mask, value=-scaled_cwd)
             del mask
         else:
             # Standard decoupled weight decay
-            p_calc.add_(decay_target, alpha=-scaled_cwd)
+            if isinstance(scaled_cwd, Tensor):
+                p_calc.addcmul_(decay_target, scaled_cwd, value=-1.0)
+            else:
+                p_calc.add_(decay_target, alpha=-scaled_cwd)
 
         del anchor, decay_target
 
-
 def apply_parameter_update(
     self,
     p: Tensor,
@@ -61,6 +72,7 @@ def apply_parameter_update(
     wd: float | None = None,
     random_int_tensor: Tensor | None = None,
     decoupled: bool = False,
+    wd_scaler: float | Tensor | None = None,
 ) -> None:
     """
     Applies decoupled weight decay (standard, cautious, centered) and the final
@@ -75,13 +87,16 @@ def apply_parameter_update(
         random_int_tensor: Optional pre-generated random tensor for stochastic
             rounding. Required for the `torch.compile` path.
         decoupled: Whenever to use the true decoupled weight decay.
+        wd_scaler: A multiplier/tensor to scale the calculated wd/cwd magnitude (e.g. for Fisher Adam WD).
     """
     wd = group["weight_decay"] if wd is None else wd
     cwd = group.get("centered_wd", 0.0)
 
     if group.get('scaled_optm', False):
         decoupled = True
-        wd, cwd = scale_wds(wd, cwd, p)
+        wd, cwd = adjust_wds(wd, cwd, p)
+        if wd_scaler is None:
+            wd, cwd = scale_wds(wd, cwd, p)
 
     # Calculate global decay factor for decoupled vs standard
     decay_factor = (lr / self._init_lr) if decoupled else lr
@@ -89,6 +104,12 @@ def apply_parameter_update(
     scaled_wd = (wd * decay_factor) if wd != 0 else None
     scaled_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
+
     state = self.state[p]
 
     # Compute full update in float32 if using bfloat16 with stochastic rounding
@@ -284,3 +305,24 @@ def post_process_loaded_state(optimizer: Optimizer) -> None:
                 # Ensure device match
                 if state[key].device != p.device:
                     state[key] = state[key].to(p.device)
+
+
+def _get_random_noise_for_sso(source: torch.Tensor) -> torch.Tensor:
+    """
+    Generates a random noise tensor for Stochastic Sign operator.
+    This function is not torch.compile-path friendly due to its use of torch.Generator.
+    """
+    global _generators
+    device = source.device
+    if device not in _generators:
+        set_seed(device)
+    # TODO, this is a workaround until torch compile error
+    # NotImplementedError: UserDefinedObjectVariable(generator) is fixed
+    generator = _generators[device]
+    # create a random noise tensor
+    return torch.randint(
+        size=source.shape,
+        device=source.device,
+        dtype=source.dtype,
+        generator=generator,
+    )

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm/util/scaled_optm.py

@@ -46,35 +46,39 @@ def scale_update(
     return update.mul_(lr)
 
 
-def scale_wds(wd: float, cwd: float, p: torch.Tensor) -> tuple[float, float]:
+def adjust_wds(wd: float, cwd: float, p: torch.Tensor) -> tuple[float, float]:
     """
-    Adjusts standard weight decay and centered weight decay based on the parameter's
-    shape and type to maintain effective regularization strength.
+    Adjusts standard weight decay and centered weight decay.
     """
     # DoRA Scale (Magnitude Vector)
     if getattr(p, '_is_dora_scale', False):
         return wd, cwd
 
-    conflict = cwd != 0
-
     if getattr(p, '_is_oft', False):
-        # Fallback to standard WD (using cwd value) if both are active.
-        return (cwd if conflict else wd), 0.0
+        return wd, 0.0
 
     if p.ndim >= 2:
-        fan_in = p.numel() // p.shape[0]
-
-        # When both WDs are active on LoRA, fallback to standard WD (using cwd value)
-        # Reverts the behavior for better DoRA tuning.
         is_lora = getattr(p, '_is_lora_A', False) or getattr(p, '_is_lora_B', False)
-        if conflict and is_lora:
-            return cwd / fan_in, 0.0
+        if is_lora:
+            return wd, 0.0
+
+    else:
+        # 1D Biases or generic 1D parameters
+        # Centered WD safely regularizes the delta without collapsing base feature variance.
+        return 0.0, cwd
+
 
+def scale_wds(wd: float, cwd: float, p: torch.Tensor) -> tuple[float, float]:
+    """
+    Scales standard weight decay and centered weight decay based on the parameter's
+    shape and type to maintain effective regularization strength.
+    """
+    if p.ndim >= 2:
+        fan_in = p.numel() // p.shape[0]
         return wd / fan_in, cwd / fan_in
 
-    # 1D Biases or generic 1D parameters
-    # Centered WD safely regularizes the delta without collapsing base feature variance.
-    return 0.0, cwd
+    # 1D tensors (like DoRA scale and Biases)
+    return wd, cwd
 
 
 @torch.no_grad()

adv_optm-2.4.dev5/adv_optm/util/signed_util.py

@@ -0,0 +1,13 @@
+import torch
+
+
+def apply_stochastic_sign(update: torch.Tensor, noise: torch.Tensor | None) -> torch.Tensor:
+    """
+    Applies the Stochastic Sign operator S_R(v).
+    Uses uniform noise injection to compute the stochastic sign
+    """
+    R = update.abs().max().clamp_min(1e-12)
+
+    if noise is None:
+        noise = torch.rand_like(update) * 2.0 - 1.0
+    return torch.sign(update / R + noise, out=update)

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm/util/update_util.py

@@ -1,5 +1,7 @@
 import torch
 
+import math
+
 def _grams_update(mt: torch.Tensor, grad: torch.Tensor, inplace: bool=False):
     """
     Applies the update rule of "Gradient Descent with Adaptive Momentum Scaling"
@@ -31,27 +33,63 @@ def _scale_sim_AdEMAMix_update(beta: float, current_step: int, alpha_grad: float
     lr = lr * total_scale
     return lr
 
+
+def _init_fisher_wd_scaler(group: dict, state: dict, p: torch.Tensor) -> torch.Tensor | None:
+    if not group.get('fisher_wd', False):
+        return
+
+    state["wd_scaler"] = torch.tensor(1.0, device=p.device)
+
+def _get_fisher_wd_scaler(group: dict, stored_scaler: torch.Tensor, p: torch.Tensor, denom: torch.Tensor, atan2: bool) -> torch.Tensor | None:
+    """
+    Calculates the Fisher weight decay scaler.
+    Maps the decay direction through the empirical Fisher information matrix
+    and clips its RMS to ensure stability.
+    From the paper:
+    "FAdam: Adam is a natural gradient optimizer using diagonal empirical Fisher information"
+    """
+    if not group.get('fisher_wd', False):
+        return None
+
+    if atan2:
+        wd_scaler = torch.atan2(stored_scaler, denom).mul_(4 / math.pi)
+    else:
+        eps = group.get('eps', 1e-8)
+        wd_scaler = 1.0 / (denom + eps)
+
+    # Reshape scaler if necessary to match parameter shape (for factored states)
+    wd_scaler = wd_scaler.view(p.shape)
+
+    gw_rms = torch.sqrt(torch.mean((p * wd_scaler) ** 2))
+    clip_coef = torch.clamp(gw_rms / 1.0, min=1.0)
+    return wd_scaler / clip_coef
+
 def _get_l1_adaptive_lr(
     p: torch.Tensor,
     update: torch.Tensor,
     state: dict,
     group: dict,
-    kappa_p: float
+    kappa_p: float,
+    rescale: bool = False,
 ) -> torch.Tensor:
     """
     Calculates the L1 adaptive learning rate based on gradient heterogeneity.
     """
-    if not group.get("l1_adaptive", False) and kappa_p != 1:
+    if not group.get("l1_adaptive", False) or kappa_p != 1:
         return None
 
-    momentum = group["momentum"]
-    alpha_grad = group["alpha_grad"]
     update_view = update.view(p.shape)
 
-    # Calculate scale factor based on momentum/update magnitude
-    scale_factor = _scale_sim_AdEMAMix_update(
-        momentum, state["step"] + 1, alpha_grad, 1, False
-    )
+    if rescale:
+        momentum = group["momentum"]
+        alpha_grad = group["alpha_grad"]
+
+        # Calculate scale factor based on momentum/update magnitude
+        scale_factor = _scale_sim_AdEMAMix_update(
+            momentum, state["step"] + 1, alpha_grad, 1, False
+        )
+    else:
+        scale_factor = 1
 
     # Determine dimension for mean calculation based on parameter type
     if getattr(p, '_is_oft', False) or getattr(p, '_is_lora_A', False):

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/adv_optm.egg-info/SOURCES.txt

@@ -27,4 +27,5 @@ adv_optm/util/factorization_util.py
 adv_optm/util/lion_k.py
 adv_optm/util/param_update.py
 adv_optm/util/scaled_optm.py
+adv_optm/util/signed_util.py
 adv_optm/util/update_util.py

{adv_optm-2.4.dev4 → adv_optm-2.4.dev5}/setup.py

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