torchzero 0.0.1__py3-none-any.whl

torchzero/modules/optimizers/adam.py

@@ -0,0 +1,118 @@
+from collections import abc
+
+import torch
+
+from ...tensorlist import TensorList
+from ...core import OptimizerModule
+
+def _adam_step(ascent: TensorList, exp_avg: TensorList, exp_avg_sq: TensorList, alpha, beta1, beta2, eps, step:int, max_exp_avg_sqs: TensorList | None, params: TensorList | None = None):
+    # Decay the first and second moment running average coefficient
+    exp_avg.lerp_compat_(ascent, 1 - beta1)
+    exp_avg_sq.mul_(beta2).addcmul_(ascent, ascent.conj(), value=1 - beta2)
+
+    bias_correction1 = 1 - beta1**step
+    bias_correction2 = 1 - beta2**step
+
+    if max_exp_avg_sqs is not None:
+        max_exp_avg_sqs.maximum_(exp_avg_sq)
+        denom = max_exp_avg_sqs.sqrt().div_(bias_correction2**0.5).add_(eps)
+    else:
+        denom = exp_avg_sq.sqrt().div_(bias_correction2**0.5).add_(eps)
+
+    if params is None:
+        return (exp_avg / denom).mul_(alpha / bias_correction1)
+
+    # else directly apply the update to params
+    params.addcdiv_(exp_avg, denom, value = -(alpha / bias_correction1))
+    return None
+
+
+
+class Adam(OptimizerModule):
+    """Adam. Combines momentum and RMSProp. Exactly matches `torch.optim.Adam`.
+
+    Args:
+        beta1 (float, optional): exponential decay rate of gradient moving average. Defaults to 0.9.
+        beta2 (float, optional): exponential decay rate of squared gradient moving average. Defaults to 0.999.
+        eps (float, optional): epsilon for numerical stability. Defaults to 1e-8.
+        amsgrad (bool, optional):
+            whether to use the AMSGrad variant of this algorithm from
+            On the Convergence of Adam and Beyond (default: False).
+        alpha (float, optional): learning rate. Defaults to 1.
+    """
+    def __init__(self, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8, alpha: float = 1, amsgrad=False):
+        defaults = dict(alpha = alpha, beta1=beta1, beta2=beta2, eps=eps)
+        super().__init__(defaults)
+
+        self.cur_step = 1
+        self.amsgrad = amsgrad
+
+    @torch.no_grad
+    def step(self, state):
+        # Adam step is a bit differet from other optimizer steps
+        # due to how common it is, I implemented two additional optimizations,
+
+        # 1st - if next module is None or if next module is LR and module after is None
+        # this will directly update parameters using `addcdiv_`
+
+        # 2nd - if next module is LR`, adam will "fuse" with it to avoid an additional add operation.
+
+        # the optimizations are quite verbose and seem to barely have any effect, so I probably won't implement
+        # this for other modules
+
+        settings = self.get_all_group_keys()
+
+        if self.amsgrad:
+            exp_avg, exp_avg_sq, max_exp_avg_sqs = self.get_state_keys('exp_avg', 'exp_avg_sq', 'max_exp_avg_sqs')
+        else:
+            exp_avg, exp_avg_sq = self.get_state_keys('exp_avg', 'exp_avg_sq')
+            max_exp_avg_sqs = None
+
+        params = None
+
+        # apply addcdiv if next module is None
+        if self.next_module is None: params = self.get_params()
+
+        # fuse with LR module if it is next
+        if self.next_module is not None and self.next_module.IS_LR_MODULE:
+            alpha = self.next_module.get_group_key('lr') * settings['alpha']
+            self.next_module._skip = True # type:ignore
+
+            # apply addcdiv if module after LR is None.
+            if self.next_module.next_module is None: params = self.get_params()
+
+        else:
+            alpha = settings['alpha']
+
+        # get params if ascent is None so we need params to access their gradient as initial ascent
+        if state.ascent is None:
+            if params is None: pg = self.get_params()
+            else: pg = params
+        else:
+            pg = None
+
+        ret = _adam_step(
+            ascent=state.maybe_use_grad_(pg),
+            exp_avg = exp_avg,
+            exp_avg_sq = exp_avg_sq,
+            alpha = alpha,
+            beta1 = settings['beta1'],
+            beta2 = settings['beta2'],
+            eps = settings['eps'],
+            step = self.cur_step,
+            max_exp_avg_sqs = max_exp_avg_sqs,
+            params = params
+        )
+
+        self.cur_step += 1
+        if params is None:
+            assert ret is not None
+            state.ascent = ret
+            return self._update_params_or_step_with_next(state)
+
+        # next module is either None or LR
+        if self.next_module is None: return state.get_loss()
+
+        # step with LR, which has _skip = True so it won't apply lr, but may step with the scheduler
+        self.next_module._update(state, None) # type:ignore
+        return state.get_loss()

torchzero/modules/optimizers/lion.py

@@ -0,0 +1,28 @@
+import torch
+
+from ...core import OptimizerModule
+from ...tensorlist import TensorList
+
+
+def _lion_step_(ascent: TensorList, ema: TensorList, beta1, beta2,):
+    update = ema.lerp_compat(ascent, 1-beta1).sign_()
+    ema.lerp_compat_(ascent, 1-beta2)
+    return update
+
+class Lion(OptimizerModule):
+    """Lion (EvoLved Sign Momentum) optimizer from https://arxiv.org/abs/2302.06675.
+
+    Args:
+        beta1 (float, optional): dampening for momentum. Defaults to 0.9.
+        beta2 (float, optional): momentum factor. Defaults to 0.99.
+    """
+
+    def __init__(self, beta1: float = 0.9, beta2: float = 0.99):
+        defaults = dict(beta1=beta1, beta2=beta2)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        beta1, beta2 = self.get_group_keys('beta1', 'beta2')
+        ema = self.get_state_key('ema')
+        return _lion_step_(ascent,ema,beta1,beta2)

torchzero/modules/optimizers/rmsprop.py

@@ -0,0 +1,51 @@
+from collections import abc
+
+import torch
+
+from ...tensorlist import TensorList
+from ...core import OptimizerModule
+
+
+def _rmsprop_step_(ascent: TensorList, mean_sqr: TensorList, smoothing, eps: TensorList):
+    mean_sqr.mul_(smoothing).addcmul_(ascent, ascent, value = 1 - smoothing)
+    return ascent.div_(mean_sqr.sqrt().add_(eps))
+
+def _centered_rmsprop_step_(ascent: TensorList, mean_sqr: TensorList, mean: TensorList, smoothing, eps: TensorList):
+    mean_sqr.mul_(smoothing).addcmul_(ascent, ascent, value = 1 - smoothing)
+    mean.lerp_compat_(ascent, 1-smoothing)
+    return ascent.div_(mean_sqr.addcmul(mean, mean, value=-1).sqrt_().add_(eps))
+
+class RMSProp(OptimizerModule):
+    """
+    Divides ascent direction by running average of its mean square root.
+
+    Exactly matches `torch.optim.RMSProp`.
+
+    Args:
+        smoothing (float, optional):
+            smoothing constant (decay of ascent mean square root running average).
+            Defaults to 0.99.
+        eps (float, optional): term added to the denominator to improve numerical stability. Defaults to 1e-8.
+        centered (float, optional):
+            if True, compute the centered RMSProp, the gradient is normalized by an estimation of its variance.
+            Defaults to False.
+
+    reference
+        https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf
+    """
+    def __init__(self, smoothing: float = 0.99, eps: float = 1e-8, centered=False):
+
+        defaults = dict(smoothing = smoothing, eps = eps)
+        super().__init__(defaults)
+        self.centered = centered
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        settings = self.get_all_group_keys()
+        if self.centered:
+            mean, mean_sqr = self.get_state_keys('mean', 'mean_sqr')
+            updated_direction = _centered_rmsprop_step_(ascent, mean_sqr, mean, settings['smoothing'], settings['eps'])
+        else:
+            mean_sqr = self.get_state_key('mean_sqr')
+            updated_direction = _rmsprop_step_(ascent, mean_sqr, settings['smoothing'], settings['eps'])
+        return updated_direction

torchzero/modules/optimizers/rprop.py

@@ -0,0 +1,99 @@
+from collections import abc
+
+import torch
+
+from ...tensorlist import TensorList, where
+from ...core import OptimizerModule
+
+
+def _bool_ones_like(x):
+    return torch.ones_like(x, dtype=torch.bool)
+
+class Rprop(OptimizerModule):
+    """
+    Resilient propagation. The update magnitude gets multiplied by `nplus` if gradient didn't change the sign,
+    or `nminus` if it did. Then the update is applied with the sign of the current gradient.
+
+    Additionally, if gradient changes sign, the update for that weight is reverted.
+    Next step, magnitude for that weight won't change.
+
+    Compared to pytorch this also implements backtracking update when sign changes.
+    To make this behave exactly the same as `torch.optim.Rprop`, set `backtrack` to False.
+
+    Args:
+        nplus (float): multiplicative increase factor for when ascent didn't change sign (default: 1.2).
+        nminus (float): multiplicative decrease factor for when ascent changed sign (default: 0.5).
+        lb (float): minimum step size, can be None (default: 1e-6)
+        ub (float): maximum step size, can be None (default: 50)
+        backtrack (float):
+            if True, when ascent sign changes, undoes last weight update, otherwise sets update to 0.
+            When this is False, this exactly matches pytorch Rprop. (default: True)
+        alpha (float): learning rate (default: 1).
+
+    reference
+        *Riedmiller, M., & Braun, H. (1993, March). A direct adaptive method for faster backpropagation learning:
+        The RPROP algorithm. In IEEE international conference on neural networks (pp. 586-591). IEEE.*
+    """
+    def __init__(
+        self,
+        nplus: float = 1.2,
+        nminus: float = 0.5,
+        lb: float | None = 1e-6,
+        ub: float | None = 50,
+        backtrack=True,
+        alpha: float = 1,
+    ):
+        defaults = dict(nplus = nplus, nminus = nminus, alpha = alpha, lb = lb, ub = ub)
+        super().__init__(defaults)
+        self.current_step = 0
+        self.backtrack = backtrack
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        params = self.get_params()
+
+        sign = ascent.sign_()
+        nplus, nminus, lb, ub = self.get_group_keys('nplus', 'nminus', 'lb', 'ub')
+        prev, allowed, magnitudes = self.get_state_keys(
+            'prev', 'allowed', 'magnitudes',
+            inits = [torch.zeros_like, _bool_ones_like, torch.zeros_like],
+            params=params
+        )
+
+        # initialize on 1st step
+        if self.current_step == 0:
+            magnitudes.fill_(self.get_group_key('alpha')).clamp_(lb, ub)
+            ascent = magnitudes * sign
+            prev.copy_(ascent)
+            self.current_step += 1
+            return ascent
+
+        mul = (sign * prev).mul_(allowed)
+
+        sign_changed = mul < 0
+        sign_same = mul > 0
+        zeroes = mul == 0
+
+        mul.fill_(1)
+        mul.masked_fill_(sign_changed, nminus)
+        mul.masked_fill_(sign_same, nplus)
+
+        # multiply magnitudes based on sign change and clamp to bounds
+        magnitudes.mul_(mul).clamp_(lb, ub)
+
+        # revert update if sign changed
+        if self.backtrack:
+            ascent = sign.mul_(magnitudes)
+            ascent.masked_set_(sign_changed, prev.neg_())
+        else:
+            ascent = sign.mul_(magnitudes * ~sign_changed)
+
+        # update allowed to only have weights where last update wasn't reverted
+        allowed.set_(sign_same | zeroes)
+
+        prev.copy_(ascent)
+        self.current_step += 1
+        return ascent
+
+
+

torchzero/modules/optimizers/sgd.py

@@ -0,0 +1,54 @@
+import typing as T
+
+import torch
+
+from ...core import OptimizerModule
+from ..momentum.momentum import _heavyball_step, _nesterov_step_
+
+class SGD(OptimizerModule):
+    """Same as `torch.optim.SGD` but as an optimizer module. Exactly matches `torch.optim.SGD`, except
+    nesterov momentum additionally supports dampening, and negative momentum is allowed.
+
+    Args:
+        momentum (float, optional): momentum. Defaults to 0.
+        dampening (float, optional): momentum dampening. Defaults to 0.
+        weight_decay (float, optional): weight decay (L2 regularization). Defaults to 0.
+        nesterov (bool, optional):
+            enables nesterov momentum, otherwise uses heavyball momentum. Defaults to False.
+        alpha (float, optional): learning rate. Defaults to 1e-3.
+    """
+    def __init__(
+        self,
+        momentum: float = 0,
+        dampening: float = 0,
+        weight_decay: float = 0,
+        nesterov: bool = False,
+        alpha: float = 1,
+    ):
+
+        defaults = dict(alpha=alpha, momentum=momentum, dampening=dampening, weight_decay=weight_decay,)
+        super().__init__(defaults)
+        self.nesterov = nesterov
+        self.current_step = 0
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        params = self.get_params()
+        settings = self.get_all_group_keys()
+
+        if any(i != 0 for i in settings['weight_decay']):
+            ascent += params * settings['weight_decay']
+
+        if any(i != 1 for i in settings['alpha']):
+            ascent *= settings['alpha']
+
+        if any(i != 0 for i in settings['momentum']):
+            velocity = self.get_state_key('velocity', init = torch.zeros_like if self.nesterov else ascent)
+            # consistency with pytorch which on first step only initializes momentum
+            if self.current_step > 0 or self.nesterov:
+                # nesterov step can be done in-place, polyak returns new direction
+                if self.nesterov: _nesterov_step_(ascent, velocity, settings['momentum'], settings['dampening'])
+                else: ascent = _heavyball_step(ascent, velocity, settings['momentum'], settings['dampening'])
+
+        self.current_step += 1
+        return ascent

torchzero/modules/orthogonalization/__init__.py

@@ -0,0 +1,2 @@
+from .svd import Orthogonalize, orthogonalize_grad_
+from .newtonschulz import ZeropowerViaNewtonSchulz, zeropower_via_newtonschulz_, DualNormCorrection

torchzero/modules/orthogonalization/newtonschulz.py

@@ -0,0 +1,159 @@
+"""
+Newton-Schulz iteration code is taken from https://github.com/KellerJordan/Muon
+
+Keller Jordan and Yuchen Jin and Vlado Boza and You Jiacheng and Franz Cecista and Laker Newhouse and Jeremy Bernstein.
+Muon: An optimizer for hidden layers in neural networks (2024). URL: https://kellerjordan.github.io/posts/muon
+"""
+from collections.abc import Iterable
+
+import torch
+
+from ...core import OptimizerModule, _Targets
+# from ...utils.compile import maybe_compile
+
+def _zeropower_via_newtonschulz5(G, steps):
+    """
+    code from https://github.com/KellerJordan/Muon
+
+    Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
+    quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
+    of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
+    zero even beyond the point where the iteration no longer converges all the way to one everywhere
+    on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
+    where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
+    performance at all relative to UV^T, where USV^T = G is the SVD.
+    """
+    assert len(G.shape) == 2
+    a, b, c = (3.4445, -4.7750,  2.0315)
+    X = G.bfloat16()
+    if G.size(0) > G.size(1):
+        X = X.T
+
+    # Ensure spectral norm is at most 1
+    X = X / (X.norm() + 1e-7)
+    # Perform the NS iterations
+    for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * A @ A # adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
+        X = a * X + B @ X
+
+    if G.size(0) > G.size(1):
+        X = X.T
+
+    return X
+
+_compiled_zeropower_via_newtonschulz5 = torch.compile(_zeropower_via_newtonschulz5)
+
+
+def zeropower_via_newtonschulz_(params: Iterable[torch.Tensor], steps: int = 6, adaptive = False, compiled = True):
+    """Uses newton-Schulz iteration to compute the zeroth power / orthogonalization of gradients of an iterable of parameters.
+
+    This sets gradients in-place.
+
+    Note that the Muon page says that embeddings and classifier heads should not be orthogonalized.
+
+    The orthogonalization code is taken from https://github.com/KellerJordan/Muon
+    Args:
+        params (abc.Iterable[torch.Tensor]): parameters that hold gradients to orthogonalize.
+        steps (int): The number of Newton-Schulz iterations to run. (6 is probably always enough).
+            The number of Newton-Schulz iterations to run. (6 is probably always enough). Defaults to 6.
+        adaptive (bool, optional):
+            Enables adaptation to scale of gradients (from https://github.com/leloykun/adaptive-muon). Defaults to False.
+        compiled (bool, optional):
+            Uses compiled newton-Schulz iteration function. Faster but won't work on windows. Defaults to True.
+
+
+    """
+    if compiled: fn = _compiled_zeropower_via_newtonschulz5
+    else: fn = _zeropower_via_newtonschulz5
+    for p in params:
+        if p.grad is not None and p.grad.ndim >= 2 and min(p.grad.shape) >= 2:
+            G = p.grad.view(p.grad.shape[0], -1)
+            X = fn(G, steps)
+
+            if adaptive:
+                # this is from https://github.com/leloykun/adaptive-muon
+                X = torch.einsum('ij,ij,ab->ab', G.type_as(X), X, X)  # Adaptive scaling,`(G * X).sum() * X` == (G.T @ X).trace() * X
+
+            p.grad = X.reshape_as(p.grad).to(p.grad, copy=False)
+
+
+class ZeropowerViaNewtonSchulz(OptimizerModule):
+    """Uses Newton-Schulz iteration to compute the zeroth power / orthogonalization of gradients of an iterable of parameters.
+
+    To disable orthogonalization for a parameter, put it into a parameter group with "newtonshultz" = False.
+    The Muon page says that embeddings and classifier heads should not be orthogonalized.
+
+    The orthogonalization code is taken from https://github.com/KellerJordan/Muon.
+
+    Note that unlike this module, Muon also uses Adam for gradients that are not orthogonalized,
+    so I'd still recommend using it. Maybe use `Wrap` to wrap it into a module (I will make muon
+    with selectable modules to optimize non-muon params soon)
+
+    However not using Adam, or putting Adam module after this to apply it to ALL updates, both seem
+    to work quite well too.
+
+    Args:
+        ns_steps (int, optional):
+            The number of Newton-Schulz iterations to run. (6 is probably always enough). Defaults to 6.
+        adaptive (bool, optional):
+            Enables adaptation to scale of gradients (from https://github.com/leloykun/adaptive-muon). Defaults to True.
+        compiled (bool, optional):
+            Uses compiled newton-Schulz iteration function. Faster but won't work on windows. Defaults to True.
+        target (str, optional):
+            determines what this module updates.
+
+            "ascent" - it updates the ascent
+
+            "grad" - it updates the gradient (and sets `.grad` attributes to updated gradient).
+
+            "closure" - it makes a new closure that sets the updated ascent to the .`grad` attributes.
+    """
+    def __init__(self, ns_steps = 6, adaptive = False, compiled=True, target:_Targets='ascent'):
+        defaults = dict(newtonshultz = True, ns_steps=ns_steps, adaptive=adaptive)
+        super().__init__(defaults, target=target)
+
+        if compiled: self._zeropower_via_newtonschulz5 = _compiled_zeropower_via_newtonschulz5
+        else: self._zeropower_via_newtonschulz5 = _zeropower_via_newtonschulz5
+
+    def _update(self, state, ascent):
+        toggle, ns_steps, adaptive = self.get_group_keys('newtonshultz', 'ns_steps', 'adaptive', cls=list)
+
+        for asc, enable, steps, ada in zip(ascent, toggle, ns_steps, adaptive):
+            if enable and len([i for i in asc.shape if i > 1]) != 0:
+                G = asc.view(asc.shape[0], -1)
+                X = self._zeropower_via_newtonschulz5(G, steps)
+
+                if ada:
+                    # this is from https://github.com/leloykun/adaptive-muon
+                    X = torch.einsum('ij,ij,ab->ab', G.type_as(X), X, X)  # Adaptive scaling,`(G * X).sum() * X` == (G.T @ X).trace() * X
+
+                asc.set_(X.reshape_as(asc).to(asc, copy=False)) # type:ignore
+
+        return ascent
+
+
+
+class DualNormCorrection(OptimizerModule):
+    """Dual norm correction from https://github.com/leloykun/adaptive-muon.
+
+    Description from the page:
+
+    Single-line modification to any (dualizer-based) optimizer that allows the optimizer to adapt to the scale of the gradients as they change during training.
+    This is done by scaling the dualized gradient by the clipped dual norm of the original gradient.
+    """
+    def __init__(self, adaptive_scale_min: int | None = -1, adaptive_scale_max: int | None = 1):
+        defaults = dict(adaptive_scale_min = adaptive_scale_min, adaptive_scale_max = adaptive_scale_max)
+        super().__init__(defaults)
+
+    def _update(self, state, ascent):
+        params = self.get_params()
+        adaptive_scale_min, adaptive_scale_max = self.get_group_keys('adaptive_scale_min', 'adaptive_scale_max')
+
+        for asc, grad, min, max in zip(ascent, state.maybe_compute_grad_(params), adaptive_scale_min, adaptive_scale_max):
+            if len([i for i in asc.shape if i > 1]) != 0:
+                scale = torch.einsum('ij,ij->', grad.view(grad.shape[0], -1), asc.view(asc.shape[0], -1))
+                if min is not None or max is not None: scale = scale.clip(min, max)
+                asc *= scale
+
+        return ascent

torchzero/modules/orthogonalization/svd.py

@@ -0,0 +1,86 @@
+"""Orthogonalization code adapted from https://github.com/MarkTuddenham/Orthogonal-Optimisers
+
+Tuddenham, M., Prügel-Bennett, A., & Hare, J. (2022).
+Orthogonalising gradients to speed up neural network optimisation. arXiv preprint arXiv:2202.07052.
+"""
+import logging
+from collections.abc import Iterable, Sequence
+
+import torch
+
+from ...core import OptimizerModule, _Targets
+
+@torch.no_grad()
+def _orthogonalize_update_(updates: Sequence[torch.Tensor], toggle = None, warn_fail=True) -> None:
+    """adapted from https://github.com/MarkTuddenham/Orthogonal-Optimisers"""
+    if toggle is None: toggle = [True] * len(updates)
+
+    # Orthogonalise the gradients using SVD
+    for grad, orth in zip(updates, toggle):
+        if orth and grad.ndim > 1:
+            G: torch.Tensor = grad.view(grad.shape[0], -1)
+            orth_G: torch.Tensor | None = None
+            try:
+                u, s, vt = torch.linalg.svd(G, full_matrices=False) # pylint:disable=not-callable
+                orth_G = u @ vt
+            except RuntimeError:
+                # if warn: logging.warning('Failed to perform SVD, adding some noise.')
+                try:
+                    u, s, v = torch.svd_lowrank(
+                        G,
+                        q=1,    # assume rank is at least 1
+                        M=1e-4 * G.mean() * torch.randn_like(G))
+                    orth_G = u @ v.T
+                except RuntimeError:
+                    if warn_fail: logging.error(('Failed to perform SVD with noise,'
+                                    ' skipping gradient orthogonalisation'))
+            if orth_G is not None:
+                grad.set_(orth_G.reshape_as(grad)) # type:ignore
+
+    return updates
+
+def orthogonalize_grad_(params: Iterable[torch.Tensor], warn_fail=False):
+    """orthogonalizes gradients of an iterable of parameters.
+
+    This updates gradients in-place.
+
+    The orthogonalization code is adapted from https://github.com/MarkTuddenham/Orthogonal-Optimisers
+    Args:
+        params (abc.Iterable[torch.Tensor]): parameters that hold gradients to orthogonalize.
+        warn_fail (bool, optional):
+            whether to print a warning when orthogonalization fails, and gradients are not
+            orthogonalized. Defaults to True.
+    """
+    grads = [p.grad for p in params if p.grad is not None]
+    _orthogonalize_update_(grads, warn_fail=warn_fail)
+
+class Orthogonalize(OptimizerModule):
+    """Orthogonalizes the update using SVD.
+
+    To disable orthogonalization for a parameter, put it into a parameter group with "orth" = False.
+
+    The orthogonalization code is adapted from https://github.com/MarkTuddenham/Orthogonal-Optimisers
+
+    Tip: :py:class:`tz.m.ZeropowerViaNewtonSchulz` is a significantly faster version of this.
+    Args:
+        warn_fail (bool, optional):
+            whether to print a warning when orthogonalization fails, and gradients are not
+            orthogonalized. Defaults to True.
+        target (str, optional):
+            determines what this module updates.
+
+            "ascent" - it updates the ascent
+
+            "grad" - it updates the gradient (and sets `.grad` attributes to updated gradient).
+
+            "closure" - it makes a new closure that sets the updated ascent to the .`grad` attributes.
+    """
+    def __init__(self, warn_fail=True, target: _Targets = 'ascent'):
+        defaults = dict(orth = True)
+        super().__init__(defaults, target = target)
+        self.warn_fail = warn_fail
+
+    def _update(self, state, ascent):
+        toggle = self.get_group_key('orth', cls=list)
+        _orthogonalize_update_(ascent, toggle, self.warn_fail)
+        return ascent

torchzero/modules/quasi_newton/__init__.py

@@ -0,0 +1,4 @@
+r"""
+This includes modules that compute a step direction via quasi-newton methods.
+"""
+# from .hv_inv_fdm import HvInvFDM

torchzero/modules/regularization/__init__.py

@@ -0,0 +1,22 @@
+r"""
+This includes regularization modules like weight decay.
+"""
+from .dropout import Dropout
+from .noise import AddNoise, Random, add_noise_
+from .normalization import (
+    Centralize,
+    ClipNorm,
+    ClipValue,
+    Normalize,
+    centralize_grad_,
+    clip_grad_norm_,
+    clip_grad_value_,
+    normalize_grad_,
+)
+from .weight_decay import (
+    WeightDecay,
+    l1_regularize_,
+    l2_regularize_,
+    weight_decay_penalty,
+)
+from .ortho_grad import OrthoGrad, orthograd_

torchzero/modules/regularization/dropout.py

@@ -0,0 +1,34 @@
+import typing as T
+from collections import abc
+
+import torch
+
+from ...tensorlist import Distributions, TensorList
+from ...core import OptimizerModule
+
+
+class Dropout(OptimizerModule):
+    """
+    Applies dropout to the update - sets random elements to 0.
+
+    This can be used to apply learning rate dropout, if put after other modules, or gradient dropout,
+    if put first.
+
+    Args:
+        p (float, optional): probability to replace update value with zero. Defaults to 0.5.
+
+    reference
+        *Lin, H., Zeng, W., Zhuang, Y., Ding, X., Huang, Y., & Paisley, J. (2022).
+        Learning rate dropout. IEEE Transactions on Neural Networks and Learning Systems,
+        34(11), 9029-9039.*
+    """
+    def __init__(self, p: float = 0.5):
+        defaults = dict(p = p)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        p = self.get_group_key('p')
+
+        ascent *= ascent.bernoulli_like(p)
+        return ascent