torchzero 0.3.9__py3-none-any.whl → 0.3.11__py3-none-any.whl

torchzero/modules/smoothing/gaussian.py

@@ -6,7 +6,7 @@ from typing import Literal
 
 import torch
 
-from ...core import Modular, Module, Vars
+from ...core import Modular, Module, Var
 from ...utils import NumberList, TensorList
 from ...utils.derivatives import jacobian_wrt
 from ..grad_approximation import GradApproximator, GradTarget
@@ -17,24 +17,24 @@ class Reformulation(Module, ABC):
         super().__init__(defaults)
 
     @abstractmethod
-    def closure(self, backward: bool, closure: Callable, params:list[torch.Tensor], vars: Vars) -> tuple[float | torch.Tensor, Sequence[torch.Tensor] | None]:
+    def closure(self, backward: bool, closure: Callable, params:list[torch.Tensor], var: Var) -> tuple[float | torch.Tensor, Sequence[torch.Tensor] | None]:
         """returns loss and gradient, if backward is False then gradient can be None"""
 
-    def pre_step(self, vars: Vars) -> Vars | None:
+    def pre_step(self, var: Var) -> Var | None:
         """This runs once before each step, whereas `closure` may run multiple times per step if further modules
         evaluate gradients at multiple points. This is useful for example to pre-generate new random perturbations."""
-        return vars
+        return var
 
-    def step(self, vars):
-        ret = self.pre_step(vars)
-        if isinstance(ret, Vars): vars = ret
+    def step(self, var):
+        ret = self.pre_step(var)
+        if isinstance(ret, Var): var = ret
 
-        if vars.closure is None: raise RuntimeError("Reformulation requires closure")
-        params, closure = vars.params, vars.closure
+        if var.closure is None: raise RuntimeError("Reformulation requires closure")
+        params, closure = var.params, var.closure
 
 
         def modified_closure(backward=True):
-            loss, grad = self.closure(backward, closure, params, vars)
+            loss, grad = self.closure(backward, closure, params, var)
 
             if grad is not None:
                 for p,g in zip(params, grad):
@@ -42,8 +42,8 @@ class Reformulation(Module, ABC):
 
             return loss
 
-        vars.closure = modified_closure
-        return vars
+        var.closure = modified_closure
+        return var
 
 
 def _decay_sigma_(self: Module, params):
@@ -58,12 +58,46 @@ def _generate_perturbations_to_state_(self: Module, params: TensorList, n_sample
     for param, prt in zip(params, zip(*perturbations)):
         self.state[param]['perturbations'] = prt
 
-def _clear_state_hook(optimizer: Modular, vars: Vars, self: Module):
+def _clear_state_hook(optimizer: Modular, var: Var, self: Module):
     for m in optimizer.unrolled_modules:
         if m is not self:
             m.reset()
 
 class GaussianHomotopy(Reformulation):
+    """Approximately smoothes the function with a gaussian kernel by sampling it at random perturbed points around current point. Both function values and gradients are averaged over all samples. The perturbed points are generated before each
+    step and remain the same throughout the step.
+
+    .. note::
+        This module reformulates the objective, it modifies the closure to evaluate value and gradients of a smoothed function. All modules after this will operate on the modified objective.
+
+    .. note::
+        This module requires the a closure passed to the optimizer step,
+        as it needs to re-evaluate the loss and gradients at perturbed points.
+
+    Args:
+        n_samples (int): number of points to sample, larger values lead to a more accurate smoothing.
+        init_sigma (float): initial scale of perturbations.
+        tol (float | None, optional):
+            if maximal parameters change value is smaller than this, sigma is reduced by :code:`decay`. Defaults to 1e-4.
+        decay (float, optional): multiplier to sigma when converged on a smoothed function. Defaults to 0.5.
+        max_steps (int | None, optional): maximum number of steps before decaying sigma. Defaults to None.
+        clear_state (bool, optional):
+            whether to clear all other module states when sigma is decayed, because the objective function changes. Defaults to True.
+        seed (int | None, optional): seed for random perturbationss. Defaults to None.
+
+    Examples:
+        Gaussian-smoothed NewtonCG
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.GaussianHomotopy(100),
+                tz.m.NewtonCG(maxiter=20),
+                tz.m.AdaptiveBacktracking(),
+            )
+
+    """
     def __init__(
         self,
         n_samples: int,
@@ -85,12 +119,12 @@ class GaussianHomotopy(Reformulation):
             else: self.global_state['generator'] = None
         return self.global_state['generator']
 
-    def pre_step(self, vars):
-        params = TensorList(vars.params)
+    def pre_step(self, var):
+        params = TensorList(var.params)
         settings = self.settings[params[0]]
         n_samples = settings['n_samples']
-        init_sigma = self.get_settings('init_sigma', params=params)
-        sigmas = self.get_state('sigma', params = params, init=init_sigma)
+        init_sigma = [self.settings[p]['init_sigma'] for p in params]
+        sigmas = self.get_state(params, 'sigma', init=init_sigma)
 
         if any('perturbations' not in self.state[p] for p in params):
             generator = self._get_generator(settings['seed'], params)
@@ -109,9 +143,9 @@ class GaussianHomotopy(Reformulation):
         tol = settings['tol']
         if tol is not None and not decayed:
             if not any('prev_params' in self.state[p] for p in params):
-                prev_params = self.get_state('prev_params', params=params, cls=TensorList, init='param')
+                prev_params = self.get_state(params, 'prev_params', cls=TensorList, init='param')
             else:
-                prev_params = self.get_state('prev_params', params=params, cls=TensorList, init='param')
+                prev_params = self.get_state(params, 'prev_params', cls=TensorList, init='param')
                 s = params - prev_params
 
                 if s.abs().global_max() <= tol:
@@ -124,10 +158,10 @@ class GaussianHomotopy(Reformulation):
             generator = self._get_generator(settings['seed'], params)
             _generate_perturbations_to_state_(self, params=params, n_samples=n_samples, sigmas=sigmas, generator=generator)
             if settings['clear_state']:
-                vars.post_step_hooks.append(partial(_clear_state_hook, self=self))
+                var.post_step_hooks.append(partial(_clear_state_hook, self=self))
 
     @torch.no_grad
-    def closure(self, backward, closure, params, vars):
+    def closure(self, backward, closure, params, var):
         params = TensorList(params)
 
         settings = self.settings[params[0]]

torchzero/modules/smoothing/laplacian.py

@@ -56,7 +56,7 @@ def _precompute_denominator(tensor: torch.Tensor, sigma) -> torch.Tensor:
     return 1 - sigma * torch.fft.fft(v) # pylint: disable = not-callable
 
 class LaplacianSmoothing(Transform):
-    """Applies laplacian smoothing via a fast Fourier transform solver.
+    """Applies laplacian smoothing via a fast Fourier transform solver which can improve generalization.
 
     Args:
         sigma (float, optional): controls the amount of smoothing. Defaults to 1.
@@ -67,11 +67,21 @@ class LaplacianSmoothing(Transform):
             minimum number of elements in a parameter to apply laplacian smoothing to.
             Only has effect if `layerwise` is True. Defaults to 4.
         target (str, optional):
-            what to set on vars.
+            what to set on var.
+
+    Examples:
+        Laplacian Smoothing Gradient Descent optimizer as in the paper
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.LaplacianSmoothing(),
+                tz.m.LR(1e-2),
+            )
 
     Reference:
-        *Osher, S., Wang, B., Yin, P., Luo, X., Barekat, F., Pham, M., & Lin, A. (2022).
-        Laplacian smoothing gradient descent. Research in the Mathematical Sciences, 9(3), 55.*
+        Osher, S., Wang, B., Yin, P., Luo, X., Barekat, F., Pham, M., & Lin, A. (2022). Laplacian smoothing gradient descent. Research in the Mathematical Sciences, 9(3), 55.
 
     """
     def __init__(self, sigma:float = 1, layerwise=True, min_numel = 4, target: Target = 'update'):
@@ -82,19 +92,17 @@ class LaplacianSmoothing(Transform):
 
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        layerwise = self.settings[params[0]]['layerwise']
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        layerwise = settings[0]['layerwise']
 
         # layerwise laplacian smoothing
         if layerwise:
 
             # precompute the denominator for each layer and store it in each parameters state
             smoothed_target = TensorList()
-            for p, t in zip(params, tensors):
-                settings = self.settings[p]
-                if p.numel() > settings['min_numel']:
-                    state = self.state[p]
-                    if 'denominator' not in state: state['denominator'] = _precompute_denominator(p, settings['sigma'])
+            for p, t, state, setting in zip(params, tensors, states, settings):
+                if p.numel() > setting['min_numel']:
+                    if 'denominator' not in state: state['denominator'] = _precompute_denominator(p, setting['sigma'])
                     smoothed_target.append(torch.fft.ifft(torch.fft.fft(t.view(-1)) / state['denominator']).real.view_as(t)) #pylint:disable=not-callable
                 else:
                     smoothed_target.append(t)
@@ -106,7 +114,7 @@ class LaplacianSmoothing(Transform):
         # precompute full denominator
         tensors = TensorList(tensors)
         if self.global_state.get('full_denominator', None) is None:
-            self.global_state['full_denominator'] = _precompute_denominator(tensors.to_vec(), self.settings[params[0]]['sigma'])
+            self.global_state['full_denominator'] = _precompute_denominator(tensors.to_vec(), settings[0]['sigma'])
 
         # apply the smoothing
         vec = tensors.to_vec()

torchzero/modules/step_size/__init__.py

@@ -0,0 +1,2 @@
+from .lr import LR, StepSize, Warmup, WarmupNormClip, RandomStepSize
+from .adaptive import PolyakStepSize, BarzilaiBorwein

torchzero/modules/step_size/adaptive.py

@@ -0,0 +1,122 @@
+"""Various step size strategies"""
+from typing import Any, Literal
+from operator import itemgetter
+import torch
+
+from ...core import Transform, Chainable
+from ...utils import TensorList, unpack_dicts, unpack_states, NumberList
+
+
+class PolyakStepSize(Transform):
+    """Polyak's subgradient method.
+
+    Args:
+        f_star (int, optional):
+            (estimated) minimal possible value of the objective function (lowest possible loss). Defaults to 0.
+        max (float | None, optional): maximum possible step size. Defaults to None.
+        use_grad (bool, optional):
+            if True, uses dot product of update and gradient to compute the step size.
+            Otherwise, dot product of update with itself is used, which has no geometric meaning so it probably won't work well.
+            Defaults to False.
+        alpha (float, optional): multiplier to Polyak step-size. Defaults to 1.
+    """
+    def __init__(self, f_star: float = 0, max: float | None = None, use_grad=False, alpha: float = 1, inner: Chainable | None = None):
+
+        defaults = dict(alpha=alpha, max=max, f_star=f_star, use_grad=use_grad)
+        super().__init__(defaults, uses_grad=use_grad, uses_loss=True, inner=inner)
+
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
+        assert grads is not None and loss is not None
+        tensors = TensorList(tensors)
+        grads = TensorList(grads)
+
+        use_grad, max, f_star = itemgetter('use_grad', 'max', 'f_star')(settings[0])
+
+        if use_grad: gg = tensors.dot(grads)
+        else: gg = tensors.dot(tensors)
+
+        if gg.abs() <= torch.finfo(gg.dtype).eps: step_size = 0 # converged
+        else: step_size = (loss - f_star) / gg
+
+        if max is not None:
+            if step_size > max: step_size = max
+
+        self.global_state['step_size'] = step_size
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        step_size = self.global_state.get('step_size', 1)
+        torch._foreach_mul_(tensors, step_size * unpack_dicts(settings, 'alpha', cls=NumberList))
+        return tensors
+
+
+
+def _bb_short(s: TensorList, y: TensorList, sy, eps, fallback):
+    yy = y.dot(y)
+    if yy < eps:
+        if sy < eps: return fallback # try to fallback on long
+        ss = s.dot(s)
+        return ss/sy
+    return sy/yy
+
+def _bb_long(s: TensorList, y: TensorList, sy, eps, fallback):
+    ss = s.dot(s)
+    if sy < eps:
+        yy = y.dot(y) # try to fallback on short
+        if yy < eps: return fallback
+        return sy/yy
+    return ss/sy
+
+def _bb_geom(s: TensorList, y: TensorList, sy, eps, fallback):
+    short = _bb_short(s, y, sy, eps, fallback)
+    long = _bb_long(s, y, sy, eps, fallback)
+    return (short * long) ** 0.5
+
+class BarzilaiBorwein(Transform):
+    """Barzilai-Borwein method.
+
+    Args:
+        type (str, optional):
+            one of "short" with formula sᵀy/yᵀy, "long" with formula sᵀs/sᵀy, or "geom" to use geometric mean of short and long.
+            Defaults to 'geom'.
+        scale_first (bool, optional):
+            whether to make first step very small when previous gradient is not available. Defaults to True.
+        fallback (float, optional): step size when denominator is less than 0 (will happen on negative curvature). Defaults to 1e-3.
+        inner (Chainable | None, optional):
+            step size will be applied to outputs of this module. Defaults to None.
+
+    """
+    def __init__(self, type: Literal['long', 'short', 'geom'] = 'geom', scale_first:bool=True, fallback:float=1e-3, inner:Chainable|None = None):
+        defaults = dict(type=type, fallback=fallback)
+        super().__init__(defaults, uses_grad=False, scale_first=scale_first, inner=inner)
+
+    def reset_for_online(self):
+        super().reset_for_online()
+        self.clear_state_keys('prev_p', 'prev_g')
+
+    @torch.no_grad
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
+        prev_p, prev_g = unpack_states(states, tensors, 'prev_p', 'prev_g', cls=TensorList)
+        fallback = unpack_dicts(settings, 'fallback', cls=NumberList)
+        type = settings[0]['type']
+
+        s = params-prev_p
+        y = tensors-prev_g
+        sy = s.dot(y)
+        eps = torch.finfo(sy.dtype).eps
+
+        if type == 'short': step_size = _bb_short(s, y, sy, eps, fallback)
+        elif type == 'long': step_size = _bb_long(s, y, sy, eps, fallback)
+        elif type == 'geom': step_size = _bb_geom(s, y, sy, eps, fallback)
+        else: raise ValueError(type)
+
+        self.global_state['step_size'] = step_size
+
+        prev_p.copy_(params)
+        prev_g.copy_(tensors)
+
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        step_size = self.global_state.get('step_size', 1)
+        torch._foreach_mul_(tensors, step_size)
+        return tensors
+

torchzero/modules/step_size/lr.py

@@ -0,0 +1,154 @@
+"""Learning rate"""
+import torch
+import random
+
+from ...core import Transform
+from ...utils import NumberList, TensorList, generic_ne, unpack_dicts
+
+def lazy_lr(tensors: TensorList, lr: float | list, inplace:bool):
+    """multiplies by lr if lr is not 1"""
+    if generic_ne(lr, 1):
+        if inplace: return tensors.mul_(lr)
+        return tensors * lr
+    return tensors
+
+class LR(Transform):
+    """Learning rate. Adding this module also adds support for LR schedulers."""
+    def __init__(self, lr: float):
+        defaults=dict(lr=lr)
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        return lazy_lr(TensorList(tensors), lr=[s['lr'] for s in settings], inplace=True)
+
+class StepSize(Transform):
+    """this is exactly the same as LR, except the `lr` parameter can be renamed to any other name to avoid clashes"""
+    def __init__(self, step_size: float, key = 'step_size'):
+        defaults={"key": key, key: step_size}
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        return lazy_lr(TensorList(tensors), lr=[s[s['key']] for s in settings], inplace=True)
+
+
+def _warmup_lr(step: int, start_lr: float | NumberList, end_lr: float | NumberList, steps: float):
+    """returns warm up lr scalar"""
+    if step > steps: return end_lr
+    return start_lr + (end_lr - start_lr) * (step / steps)
+
+class Warmup(Transform):
+    """Learning rate warmup, linearly increases learning rate multiplier from :code:`start_lr` to :code:`end_lr` over :code:`steps` steps.
+
+    Args:
+        steps (int, optional): number of steps to perform warmup for. Defaults to 100.
+        start_lr (_type_, optional): initial learning rate multiplier on first step. Defaults to 1e-5.
+        end_lr (float, optional): learning rate multiplier at the end and after warmup. Defaults to 1.
+
+    Example:
+        Adam with 1000 steps warmup
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.LR(1e-2),
+                tz.m.Warmup(steps=1000)
+            )
+
+    """
+    def __init__(self, steps = 100, start_lr = 1e-5, end_lr:float = 1):
+        defaults = dict(start_lr=start_lr,end_lr=end_lr, steps=steps)
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        start_lr, end_lr = unpack_dicts(settings, 'start_lr', 'end_lr', cls = NumberList)
+        num_steps = settings[0]['steps']
+        step = self.global_state.get('step', 0)
+
+        tensors = lazy_lr(
+            TensorList(tensors),
+            lr=_warmup_lr(step=step, start_lr=start_lr, end_lr=end_lr, steps=num_steps),
+            inplace=True
+        )
+        self.global_state['step'] = step + 1
+        return tensors
+
+class WarmupNormClip(Transform):
+    """Warmup via clipping of the update norm.
+
+    Args:
+        start_norm (_type_, optional): maximal norm on the first step. Defaults to 1e-5.
+        end_norm (float, optional): maximal norm on the last step. After that, norm clipping is disabled. Defaults to 1.
+        steps (int, optional): number of steps to perform warmup for. Defaults to 100.
+
+    Example:
+        Adam with 1000 steps norm clip warmup
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.WarmupNormClip(steps=1000)
+                tz.m.LR(1e-2),
+            )
+    """
+    def __init__(self, steps = 100, start_norm = 1e-5, end_norm:float = 1):
+        defaults = dict(start_norm=start_norm,end_norm=end_norm, steps=steps)
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        start_norm, end_norm = unpack_dicts(settings, 'start_norm', 'end_norm', cls = NumberList)
+        num_steps = settings[0]['steps']
+        step = self.global_state.get('step', 0)
+        if step > num_steps: return tensors
+
+        tensors = TensorList(tensors)
+        norm = tensors.global_vector_norm()
+        current_max_norm = _warmup_lr(step, start_norm[0], end_norm[0], num_steps)
+        if norm > current_max_norm:
+            tensors.mul_(current_max_norm / norm)
+
+        self.global_state['step'] = step + 1
+        return tensors
+
+
+class RandomStepSize(Transform):
+    """Uses random global or layer-wise step size from `low` to `high`.
+
+    Args:
+        low (float, optional): minimum learning rate. Defaults to 0.
+        high (float, optional): maximum learning rate. Defaults to 1.
+        parameterwise (bool, optional):
+            if True, generate random step size for each parameter separately,
+            if False generate one global random step size. Defaults to False.
+    """
+    def __init__(self, low: float = 0, high: float = 1, parameterwise=False, seed:int|None=None):
+        defaults = dict(low=low, high=high, parameterwise=parameterwise,seed=seed)
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        s = settings[0]
+        parameterwise = s['parameterwise']
+
+        seed = s['seed']
+        if 'generator' not in self.global_state:
+            self.global_state['generator'] = random.Random(seed)
+        generator: random.Random = self.global_state['generator']
+
+        if parameterwise:
+            low, high = unpack_dicts(settings, 'low', 'high')
+            lr = [generator.uniform(l, h) for l, h in zip(low, high)]
+        else:
+            low = s['low']
+            high = s['high']
+            lr = generator.uniform(low, high)
+
+        torch._foreach_mul_(tensors, lr)
+        return tensors

torchzero/modules/weight_decay/__init__.py

@@ -1 +1 @@
-from .weight_decay import WeightDecay, DirectWeightDecay, decay_weights_
+from .weight_decay import WeightDecay, DirectWeightDecay, decay_weights_, RelativeWeightDecay

torchzero/modules/weight_decay/weight_decay.py

@@ -1,9 +1,11 @@
 from collections.abc import Iterable, Sequence
+from typing import Literal
 
 import torch
 
 from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList, as_tensorlist
+from ...utils import NumberList, TensorList, as_tensorlist, unpack_dicts, unpack_states
+
 
 @torch.no_grad
 def weight_decay_(
@@ -20,17 +22,126 @@ def weight_decay_(
 
 
 class WeightDecay(Transform):
+    """Weight decay.
+
+    Args:
+        weight_decay (float): weight decay scale.
+        ord (int, optional): order of the penalty, e.g. 1 for L1 and 2 for L2. Defaults to 2.
+        target (Target, optional): what to set on var. Defaults to 'update'.
+
+    Examples:
+        Adam with non-decoupled weight decay
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.WeightDecay(1e-3),
+                tz.m.Adam(),
+                tz.m.LR(1e-3)
+            )
+
+        Adam with decoupled weight decay that still scales with learning rate
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.WeightDecay(1e-3),
+                tz.m.LR(1e-3)
+            )
+
+        Adam with fully decoupled weight decay that doesn't scale with learning rate
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.LR(1e-3),
+                tz.m.WeightDecay(1e-6)
+            )
+
+    """
     def __init__(self, weight_decay: float, ord: int = 2, target: Target = 'update'):
+
         defaults = dict(weight_decay=weight_decay, ord=ord)
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        weight_decay = self.get_settings('weight_decay', params=params, cls=NumberList)
-        ord = self.settings[params[0]]['ord']
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        weight_decay = NumberList(s['weight_decay'] for s in settings)
+        ord = settings[0]['ord']
 
         return weight_decay_(as_tensorlist(tensors), as_tensorlist(params), weight_decay, ord)
 
+class RelativeWeightDecay(Transform):
+    """Weight decay relative to the mean absolute value of update, gradient or parameters depending on value of :code:`norm_input` argument.
+
+    Args:
+        weight_decay (float): relative weight decay scale.
+        ord (int, optional): order of the penalty, e.g. 1 for L1 and 2 for L2. Defaults to 2.
+        norm_input (str, optional):
+            determines what should weight decay be relative to. "update", "grad" or "params".
+            Defaults to "update".
+        target (Target, optional): what to set on var. Defaults to 'update'.
+
+    Examples:
+        Adam with non-decoupled relative weight decay
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.RelativeWeightDecay(1e-3),
+                tz.m.Adam(),
+                tz.m.LR(1e-3)
+            )
+
+        Adam with decoupled relative weight decay
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.RelativeWeightDecay(1e-3),
+                tz.m.LR(1e-3)
+            )
+
+    """
+    def __init__(
+        self,
+        weight_decay: float = 0.1,
+        ord: int  = 2,
+        norm_input: Literal["update", "grad", "params"] = "update",
+        target: Target = "update",
+    ):
+        defaults = dict(weight_decay=weight_decay, ord=ord, norm_input=norm_input)
+        super().__init__(defaults, uses_grad=norm_input == 'grad', target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        weight_decay = NumberList(s['weight_decay'] for s in settings)
+
+        ord = settings[0]['ord']
+        norm_input = settings[0]['norm_input']
+
+        if norm_input == 'update': src = TensorList(tensors)
+        elif norm_input == 'grad':
+            assert grads is not None
+            src = TensorList(grads)
+        elif norm_input == 'params':
+            src = TensorList(params)
+        else:
+            raise ValueError(norm_input)
+
+        mean_abs = src.abs().global_mean()
+
+        return weight_decay_(as_tensorlist(tensors), as_tensorlist(params), weight_decay * mean_abs, ord)
+
+
 @torch.no_grad
 def decay_weights_(params: Iterable[torch.Tensor], weight_decay: float | NumberList, ord:int=2):
     """directly decays weights in-place"""
@@ -38,15 +149,20 @@ def decay_weights_(params: Iterable[torch.Tensor], weight_decay: float | NumberL
     weight_decay_(params, params, -weight_decay, ord)
 
 class DirectWeightDecay(Module):
-    """directly decays weights in-place"""
+    """Directly applies weight decay to parameters.
+
+    Args:
+        weight_decay (float): weight decay scale.
+        ord (int, optional): order of the penalty, e.g. 1 for L1 and 2 for L2. Defaults to 2.
+    """
     def __init__(self, weight_decay: float, ord: int = 2,):
         defaults = dict(weight_decay=weight_decay, ord=ord)
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, vars):
-        weight_decay = self.get_settings('weight_decay', params=vars.params, cls=NumberList)
-        ord = self.settings[vars.params[0]]['ord']
+    def step(self, var):
+        weight_decay = self.get_settings(var.params, 'weight_decay', cls=NumberList)
+        ord = self.settings[var.params[0]]['ord']
 
-        decay_weights_(vars.params, weight_decay, ord)
-        return vars
+        decay_weights_(var.params, weight_decay, ord)
+        return var