torchzero 0.1.7__py3-none-any.whl → 0.3.1__py3-none-any.whl

torchzero/modules/momentum/__init__.py

@@ -1,4 +1,14 @@
-"""
-Modules that implement momentum.
-"""
-from .momentum import HeavyBall, NesterovMomentum, RandomCoordinateMomentum, GradientAveraging
+from .averaging import Averaging, MedianAveraging, WeightedAveraging
+from .cautious import (
+    Cautious,
+    IntermoduleCautious,
+    ScaleByGradCosineSimilarity,
+    ScaleModulesByCosineSimilarity,
+    UpdateGradientSignConsistency,
+)
+from .ema import EMA, Debias, Debias2, EMASquared, SqrtEMASquared, CenteredEMASquared, CenteredSqrtEMASquared
+from .experimental import CoordinateMomentum
+# from .matrix_momentum import MatrixMomentum
+
+from .momentum import NAG, HeavyBall
+from .matrix_momentum import MatrixMomentum, AdaptiveMatrixMomentum

torchzero/modules/momentum/averaging.py

@@ -0,0 +1,78 @@
+from collections import deque
+from collections.abc import Sequence
+from typing import Any, Literal, cast
+
+import torch
+
+from ...core import TensorwiseTransform, Target
+from ...utils import tolist
+
+
+class Averaging(TensorwiseTransform):
+    def __init__(self, history_size: int, target: Target = 'update'):
+        defaults = dict(history_size=history_size)
+        super().__init__(uses_grad=False, defaults=defaults, target=target)
+
+    @torch.no_grad
+    def transform(self, tensor, param, grad, vars):
+        history_size = self.settings[param]['history_size']
+        state = self.state[param]
+        if 'history' not in state:
+            state['history'] = deque(maxlen=history_size)
+            state['average'] = torch.zeros_like(tensor)
+
+        history = state['history']; average = state['average']
+        if len(history) == history_size: average -= history[0]
+        history.append(tensor)
+        average += tensor
+
+        return average / len(history)
+
+class WeightedAveraging(TensorwiseTransform):
+    """weights are oldest to newest"""
+    def __init__(self, weights: Sequence[float] | torch.Tensor | Any, target: Target = 'update'):
+        defaults = dict(weights = tolist(weights))
+        super().__init__(uses_grad=False, defaults=defaults, target=target)
+
+    @torch.no_grad
+    def transform(self, tensor, param, grad, vars):
+        weights = self.settings[param]['weights']
+        state = self.state[param]
+
+        if 'history' not in state:
+            state['history'] = deque(maxlen=len(weights))
+
+        history = state['history']
+        history.append(tensor)
+        if len(history) != len(weights):
+            weights = weights[-len(history):]
+
+        average = None
+        for i, (h, w) in enumerate(zip(history, weights)):
+            if average is None: average = h * (w / len(history))
+            else:
+                if w == 0: continue
+                average += h * (w / len(history))
+
+        assert average is not None
+        return average
+
+
+class MedianAveraging(TensorwiseTransform):
+    def __init__(self, history_size: int, target: Target = 'update'):
+        defaults = dict(history_size = history_size)
+        super().__init__(uses_grad=False, defaults=defaults, target=target)
+
+    @torch.no_grad
+    def transform(self, tensor, param, grad, vars):
+        history_size = self.settings[param]['history_size']
+        state = self.state[param]
+
+        if 'history' not in state:
+            state['history'] = deque(maxlen=history_size)
+
+        history = state['history']
+        history.append(tensor)
+
+        stacked = torch.stack(tuple(history), 0)
+        return torch.quantile(stacked, 0.5, dim = 0)

torchzero/modules/momentum/cautious.py

@@ -0,0 +1,181 @@
+from collections import deque
+from operator import itemgetter
+from typing import Literal
+
+import torch
+
+from ...core import Target, Transform, Module, Chainable
+from ...utils import NumberList, TensorList
+
+
+def cautious_(
+    tensors_: TensorList,
+    grads: TensorList,
+    normalize: bool,
+    eps: float,
+    mode: Literal['zero', 'grad', 'backtrack']
+):
+    # mask will be > 0 for parameters where both signs are the same
+    mask = (tensors_ * grads) > 0
+    if mode in ('zero', 'grad'):
+        if normalize and mode == 'zero':
+            fmask = mask.to(tensors_[0].dtype)
+            fmask /= fmask.global_mean().clip(min=eps) # type:ignore
+        else:
+            fmask = mask
+
+        tensors_ *= fmask
+
+        if mode == 'grad':
+            tensors_ += grads * mask.logical_not_()
+
+        return tensors_
+
+    # mode = 'backtrack'
+    tensors_ -= tensors_.mul(2).mul_(mask.logical_not_())
+    return tensors_
+
+class Cautious(Transform):
+    """Negates update for parameters where update and gradient sign is inconsistent.
+    Optionally normalizes the update by the number of parameters that are not masked.
+    This is meant to be used after any momentum-based modules.
+
+    Args:
+        normalize (bool, optional):
+            renormalize update after masking.
+            only has effect when mode is 'zero'. Defaults to False.
+        eps (float, optional): epsilon for normalization. Defaults to 1e-6.
+        mode (str, optional):
+            what to do with updates with inconsistent signs.
+
+            "zero" - set them to zero (as in paper)
+
+            "grad" - set them to the gradient
+
+            "backtrack" - negate them (same as using update magnitude and gradient sign)
+
+    reference
+        *Cautious Optimizers: Improving Training with One Line of Code.
+        Kaizhao Liang, Lizhang Chen, Bo Liu, Qiang Liu*
+    """
+
+    def __init__(
+        self,
+        normalize=False,
+        eps=1e-6,
+        mode: Literal["zero", "grad", "backtrack"] = "zero",
+        target: Target = "update",
+    ):
+        defaults = dict(normalize=normalize, eps=eps, mode=mode)
+        super().__init__(defaults, uses_grad=True, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        assert grads is not None
+        mode, normalize, eps = itemgetter('mode', 'normalize', 'eps')(self.settings[params[0]])
+        return cautious_(TensorList(tensors), TensorList(grads), normalize=normalize, eps=eps, mode=mode)
+
+class UpdateGradientSignConsistency(Transform):
+    """1 where signs match 0 otherwise"""
+    def __init__(self, normalize = False, eps=1e-6, target: Target = 'update'):
+        defaults = dict(normalize=normalize, eps=eps)
+        super().__init__(defaults, uses_grad=True, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        assert grads is not None
+        normalize, eps = itemgetter('normalize', 'eps')(self.settings[params[0]])
+
+        mask = (TensorList(tensors).mul_(grads)).gt_(0)
+        if normalize: mask = mask / mask.global_mean().clip(min = eps) # pyright: ignore[reportOperatorIssue]
+
+        return mask
+
+class IntermoduleCautious(Module):
+    def __init__(
+        self,
+        main: Chainable,
+        compare: Chainable,
+        normalize=False,
+        eps=1e-6,
+        mode: Literal["zero", "grad", "backtrack"] = "zero",
+    ):
+        defaults = dict(normalize=normalize, eps=eps, mode=mode)
+        super().__init__(defaults)
+
+        self.set_child('main', main)
+        self.set_child('compare', compare)
+
+    @torch.no_grad
+    def step(self, vars):
+        main = self.children['main']
+        compare = self.children['compare']
+
+        main_vars = main.step(vars.clone(clone_update=True))
+        vars.update_attrs_from_clone_(main_vars)
+
+        compare_vars = compare.step(vars.clone(clone_update=True))
+        vars.update_attrs_from_clone_(compare_vars)
+
+        mode, normalize, eps = itemgetter('mode', 'normalize', 'eps')(self.settings[vars.params[0]])
+        vars.update = cautious_(
+            TensorList(main_vars.get_update()),
+            TensorList(compare_vars.get_update()),
+            normalize=normalize,
+            mode=mode,
+            eps=eps,
+        )
+
+        return vars
+
+class ScaleByGradCosineSimilarity(Transform):
+    def __init__(
+        self,
+        eps=1e-6,
+        target: Target = "update",
+    ):
+        defaults = dict(eps=eps)
+        super().__init__(defaults, uses_grad=True, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        assert grads is not None
+        eps = self.settings[params[0]]['eps']
+        tensors = TensorList(tensors)
+        grads = TensorList(grads)
+        cos_sim = (tensors.dot(grads)) / (tensors.global_vector_norm() * grads.global_vector_norm()).clip(min=eps)
+
+        return tensors.mul_(cos_sim)
+
+class ScaleModulesByCosineSimilarity(Module):
+    def __init__(
+        self,
+        main: Chainable,
+        compare: Chainable,
+        eps=1e-6,
+    ):
+        defaults = dict(eps=eps)
+        super().__init__(defaults)
+
+        self.set_child('main', main)
+        self.set_child('compare', compare)
+
+    @torch.no_grad
+    def step(self, vars):
+        main = self.children['main']
+        compare = self.children['compare']
+
+        main_vars = main.step(vars.clone(clone_update=True))
+        vars.update_attrs_from_clone_(main_vars)
+
+        compare_vars = compare.step(vars.clone(clone_update=True))
+        vars.update_attrs_from_clone_(compare_vars)
+
+        m = TensorList(main_vars.get_update())
+        c = TensorList(compare_vars.get_update())
+        eps = self.settings[vars.params[0]]['eps']
+
+        cos_sim = (m.dot(c)) / (m.global_vector_norm() * c.global_vector_norm()).clip(min=eps)
+
+        vars.update = m.mul_(cos_sim)
+        return vars

torchzero/modules/momentum/ema.py

@@ -0,0 +1,173 @@
+from collections import deque
+from operator import itemgetter
+from typing import Literal
+
+import torch
+
+from ...core import Target, Transform
+from ...utils import TensorList, NumberList
+from ..functional import debias, ema_, ema_sq_, sqrt_ema_sq_, centered_ema_sq_, sqrt_centered_ema_sq_, debias_second_momentum
+
+
+class EMA(Transform):
+    """Maintains EMA of update.
+
+    Args:
+        momentum (float, optional): momentum (beta). Defaults to 0.9.
+        dampening (float, optional): momentum dampening. Defaults to 0.
+        debiased (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
+        lerp (bool, optional): whether to use linear interpolation. Defaults to True.
+        target (Target, optional): target to apply EMA to. Defaults to 'update'.
+    """
+    def __init__(self, momentum:float=0.9, dampening:float=0, debiased: bool = False, lerp=True, ema_init: Literal['zeros', 'update'] = 'zeros', target: Target = 'update'):
+        defaults = dict(momentum=momentum,dampening=dampening,debiased=debiased,lerp=lerp,ema_init=ema_init)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        debiased, lerp, ema_init = itemgetter('debiased','lerp','ema_init')(self.settings[params[0]])
+
+        exp_avg = self.get_state('exp_avg', params=params, init=torch.zeros_like if ema_init=='zeros' else tensors, cls=TensorList)
+        momentum, dampening = self.get_settings('momentum','dampening', params=params, cls=NumberList)
+
+        exp_avg = ema_(TensorList(tensors), exp_avg_=exp_avg,beta=momentum,dampening=dampening,lerp=lerp)
+
+        if debiased: return debias(exp_avg, step=step, beta1=momentum, alpha=1, inplace=False)
+        else: return exp_avg.clone() # this has exp_avg storage so needs to be cloned
+
+
+class EMASquared(Transform):
+    EMA_SQ_FN: staticmethod = staticmethod(ema_sq_)
+
+    def __init__(self, beta:float=0.999, amsgrad=False, pow:float=2, target: Target = 'update'):
+        defaults = dict(beta=beta,pow=pow,amsgrad=amsgrad)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        amsgrad, pow = itemgetter('amsgrad', 'pow')(self.settings[params[0]])
+        beta = self.get_settings('beta', params=params, cls=NumberList)
+
+        if amsgrad:
+            exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg_sq', 'max_exp_avg_sq', params=params, cls=TensorList)
+        else:
+            exp_avg_sq = self.get_state('exp_avg_sq', params=params, cls=TensorList)
+            max_exp_avg_sq = None
+
+        return self.EMA_SQ_FN(TensorList(tensors), exp_avg_sq_=exp_avg_sq, beta=beta, max_exp_avg_sq_=max_exp_avg_sq, pow=pow).clone()
+
+class SqrtEMASquared(Transform):
+    SQRT_EMA_SQ_FN: staticmethod = staticmethod(sqrt_ema_sq_)
+
+    def __init__(self, beta:float=0.999, amsgrad=False, debiased: bool = False, pow:float=2, target: Target = 'update',):
+        defaults = dict(beta=beta,pow=pow,amsgrad=amsgrad,debiased=debiased)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        amsgrad, pow, debiased = itemgetter('amsgrad', 'pow', 'debiased')(self.settings[params[0]])
+        beta = self.get_settings('beta', params=params, cls=NumberList)
+
+        if amsgrad:
+            exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg_sq', 'max_exp_avg_sq', params=params, cls=TensorList)
+        else:
+            exp_avg_sq = self.get_state('exp_avg_sq', params=params, cls=TensorList)
+            max_exp_avg_sq = None
+
+        return self.SQRT_EMA_SQ_FN(
+            TensorList(tensors),
+            exp_avg_sq_=exp_avg_sq,
+            beta=beta,
+            max_exp_avg_sq_=max_exp_avg_sq,
+            debiased=debiased,
+            step=step,
+            pow=pow,
+        )
+
+
+class Debias(Transform):
+    def __init__(self, beta1: float | None = None, beta2: float | None = None, alpha: float = 1, pow:float=2, target: Target = 'update',):
+        defaults = dict(beta1=beta1, beta2=beta2, alpha=alpha, pow=pow)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        settings = self.settings[params[0]]
+        pow = settings['pow']
+        alpha, beta1, beta2 = self.get_settings('alpha', 'beta1', 'beta2', params=params, cls=NumberList)
+
+        return debias(TensorList(tensors), step=step, beta1=beta1, beta2=beta2, alpha=alpha, pow=pow, inplace=True)
+
+class Debias2(Transform):
+    def __init__(self, beta: float = 0.999, pow: float = 2, target: Target = 'update',):
+        defaults = dict(beta=beta, pow=pow)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        pow = self.settings[params[0]]['pow']
+        beta = self.get_settings('beta', params=params, cls=NumberList)
+        return debias_second_momentum(TensorList(tensors), step=step, beta=beta, pow=pow, inplace=True)
+
+class CenteredEMASquared(Transform):
+    def __init__(self, beta: float = 0.99, amsgrad=False, pow:float=2, target: Target = 'update'):
+        defaults = dict(beta=beta, amsgrad=amsgrad, pow=pow)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        amsgrad, pow = itemgetter('amsgrad', 'pow')(self.settings[params[0]])
+        beta = self.get_settings('beta', params=params, cls=NumberList)
+
+        if amsgrad:
+            exp_avg, exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg', 'exp_avg_sq', 'max_exp_avg_sq', params=params, cls=TensorList)
+        else:
+            exp_avg, exp_avg_sq = self.get_state('exp_avg', 'exp_avg_sq', params=params, cls=TensorList)
+            max_exp_avg_sq = None
+
+        return centered_ema_sq_(
+            TensorList(tensors),
+            exp_avg_=exp_avg,
+            exp_avg_sq_=exp_avg_sq,
+            beta=beta,
+            max_exp_avg_sq_=max_exp_avg_sq,
+            pow=pow,
+        ).clone()
+
+class CenteredSqrtEMASquared(Transform):
+    def __init__(self, beta: float = 0.99, amsgrad=False, debiased: bool = False, pow:float=2, target: Target = 'update'):
+        defaults = dict(beta=beta, amsgrad=amsgrad, debiased=debiased, pow=pow)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        amsgrad, pow, debiased = itemgetter('amsgrad', 'pow', 'debiased')(self.settings[params[0]])
+        beta = self.get_settings('beta', params=params, cls=NumberList)
+
+        if amsgrad:
+            exp_avg, exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg', 'exp_avg_sq', 'max_exp_avg_sq', params=params, cls=TensorList)
+        else:
+            exp_avg, exp_avg_sq = self.get_state('exp_avg', 'exp_avg_sq', params=params, cls=TensorList)
+            max_exp_avg_sq = None
+
+        return sqrt_centered_ema_sq_(
+            TensorList(tensors),
+            exp_avg_=exp_avg,
+            exp_avg_sq_=exp_avg_sq,
+            beta=beta,
+            debiased=debiased,
+            step=step,
+            max_exp_avg_sq_=max_exp_avg_sq,
+            pow=pow,
+        )

torchzero/modules/momentum/experimental.py

@@ -0,0 +1,189 @@
+from collections.abc import Sequence
+from functools import partial
+from operator import itemgetter
+from typing import Literal
+
+import torch
+
+from ...core import Target, Transform
+from ...utils import NumberList, TensorList
+from ..functional import ema_, ema_sq_, sqrt_ema_sq_
+from .ema import EMASquared, SqrtEMASquared
+from .momentum import nag_
+
+
+def precentered_ema_sq_(
+    tensors: TensorList,
+    exp_avg_: TensorList,
+    exp_avg_sq_: TensorList,
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    step: int,
+    min_step: int,
+    pow: float,
+    max_exp_avg_sq_: TensorList | None,
+):
+    """
+    Squared EMA of (update - 1st EMA). Starts taking effect after `min_step` to avoid division by epsilon.
+
+    returns `exp_avg_sq_` or `max_exp_avg_sq_`.
+    """
+    exp_avg_ = ema_(tensors, exp_avg_=exp_avg_, beta=beta1, dampening=0, lerp=False)
+
+    if step < min_step: centered_update = tensors
+    else: centered_update = tensors - exp_avg_
+
+    exp_avg_sq_=ema_sq_(
+        centered_update,
+        exp_avg_sq_=exp_avg_sq_,
+        beta=beta2,
+        pow=pow,
+        max_exp_avg_sq_=max_exp_avg_sq_,
+    )
+    return exp_avg_sq_
+
+class PrecenteredEMASquared(Transform):
+    def __init__(self, beta1:float=0.99, beta2=0.99, min_step: int = 2, amsgrad=False, pow:float=2, target: Target = 'update'):
+        defaults = dict(beta1=beta1,beta2=beta2,pow=pow,amsgrad=amsgrad, min_step=min_step)
+        super().__init__(defaults, uses_grad=False, target=target)
+        self.current_step = 0
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        self.current_step += 1
+
+        beta1, beta2 = self.get_settings('beta1','beta2', params=params, cls=NumberList)
+        amsgrad, pow, min_step = itemgetter('amsgrad', 'pow', 'min_step')(self.settings[params[0]])
+
+        if amsgrad:
+            exp_avg, exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg', 'exp_avg_sq', 'max_exp_avg_sq', params=params, cls=TensorList)
+        else:
+            exp_avg, exp_avg_sq = self.get_state('exp_avg', 'exp_avg_sq', params=params, cls=TensorList)
+            max_exp_avg_sq = None
+
+        return precentered_ema_sq_(
+            TensorList(tensors),
+            exp_avg_ = exp_avg,
+            exp_avg_sq_=exp_avg_sq,
+            beta1=beta1,
+            beta2=beta2,
+            step = self.current_step,
+            min_step=min_step,
+            pow=pow,
+            max_exp_avg_sq_=max_exp_avg_sq,
+        ).clone()
+
+
+def nag_ema_sq_(
+    tensors: TensorList,
+    exp_avg_sq_: TensorList,
+    beta: float | NumberList,
+    max_exp_avg_sq_: TensorList | None,
+    pow: float,
+    lerp:bool=True,
+):
+    """
+    Nesterov EMA of squared tensors.
+
+    Returns `exp_avg_sq_` or `max_exp_avg_sq_`.
+    """
+    if pow == 1: tensors = tensors.abs()
+    elif pow%2 == 0: tensors = tensors.pow(pow)
+    else: tensors = tensors.pow(pow).abs()
+
+    exp_avg_sq_=nag_(tensors,velocity_=exp_avg_sq_,momentum=beta,dampening=0,lerp=lerp,)
+
+    # AMSGrad
+    if max_exp_avg_sq_ is not None:
+        max_exp_avg_sq_.maximum_(exp_avg_sq_)
+        exp_avg_sq_ = max_exp_avg_sq_
+
+    return exp_avg_sq_
+
+def sqrt_nag_ema_sq_(
+    tensors: TensorList,
+    exp_avg_sq_: TensorList,
+    beta: float | NumberList,
+    max_exp_avg_sq_: TensorList | None,
+    debiased: bool,
+    step: int,
+    pow: float,
+    lerp:bool=False,
+):
+    """
+    Square root of nesterov EMA of squared tensors.
+
+    Returns new tensors.
+    """
+    return sqrt_ema_sq_(tensors=tensors,exp_avg_sq_=exp_avg_sq_,beta=beta,max_exp_avg_sq_=max_exp_avg_sq_,
+                        pow=pow,debiased=debiased,step=step,ema_sq_fn=partial(nag_ema_sq_,lerp=lerp))
+
+class NesterovEMASquared(EMASquared):
+    EMA_SQ_FN = staticmethod(nag_ema_sq_)
+
+class SqrtNesterovEMASquared(SqrtEMASquared):
+    SQRT_EMA_SQ_FN = staticmethod(sqrt_nag_ema_sq_)
+
+
+def coordinate_momentum_(
+    tensors: TensorList,
+    velocity_: TensorList,
+    p: float | NumberList,
+):
+    """
+    sets `velocity_` to p% random values from `tensors`.
+
+    Returns `velocity_`
+    """
+    mask = tensors.bernoulli_like(p).as_bool()
+    velocity_.masked_set_(mask, tensors)
+    return velocity_
+
+
+class CoordinateMomentum(Transform):
+    def __init__(self, p: float = 0.1, target: Target = 'update'):
+        defaults = dict(p=p)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        p = self.get_settings('p', params=params, cls=NumberList)
+        velocity = self.get_state('velocity', params=params, cls=TensorList)
+        return coordinate_momentum_(TensorList(tensors), velocity_=velocity, p=p).clone()
+
+
+# def multiplicative_momentum_(
+#     tensors_: TensorList,
+#     velocity_: TensorList,
+#     momentum: float | NumberList,
+#     dampening: float | NumberList,
+#     normalize_velocity: bool = True,
+#     abs: bool = False,
+#     lerp: bool = False,
+# ):
+#     """
+#     abs: if True, tracks momentum of absolute magnitudes.
+
+#     returns `tensors_`.
+#     """
+#     tensors_into_velocity = tensors_.abs() if abs else tensors_
+#     ema_(tensors_into_velocity, exp_avg_=velocity_, beta=momentum, dampening=0, lerp=lerp)
+
+#     if normalize_velocity: velocity_ = velocity_ / velocity_.std().add_(1e-8)
+#     return tensors_.mul_(velocity_.lazy_mul(1-dampening) if abs else velocity_.abs().lazy_mul_(1-dampening))
+
+
+# class MultiplicativeMomentum(Transform):
+#     """sucks"""
+#     def __init__(self, momentum: float = 0.9, dampening: float = 0,normalize_velocity: bool = True, abs: bool = False, lerp: bool = False):
+#         defaults = dict(momentum=momentum, dampening=dampening, normalize_velocity=normalize_velocity,abs=abs, lerp=lerp)
+#         super().__init__(defaults, uses_grad=False)
+
+#     @torch.no_grad
+#     def transform(self, tensors, params, grads, vars):
+#         momentum,dampening = self.get_settings('momentum','dampening', params=params, cls=NumberList)
+#         abs,lerp,normalize_velocity = self.first_setting('abs','lerp','normalize_velocity', params=params)
+#         velocity = self.get_state('velocity', params=params, cls=TensorList)
+#         return multiplicative_momentum_(TensorList(target), velocity_=velocity, momentum=momentum, dampening=dampening,
+#                                         normalize_velocity=normalize_velocity,abs=abs,lerp=lerp)
+