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

torchzero/modules/optimizers/ladagrad.py

@@ -0,0 +1,183 @@
+from collections import deque
+from typing import Literal, Any
+import warnings
+
+import torch
+from ...core import Chainable, TensorwiseTransform
+
+def lm_adagrad_update(history: deque[torch.Tensor], damping, rdamping):
+    M = torch.stack(tuple(history), dim=1)# / len(history)
+    MTM = M.T @ M
+    if damping != 0:
+        MTM.add_(torch.eye(MTM.size(0), device=MTM.device, dtype=MTM.dtype).mul_(damping))
+
+    try:
+        L, Q = torch.linalg.eigh(MTM) # pylint:disable=not-callable
+
+        tol = torch.finfo(M.dtype).eps * L.amax() # remove small eigenvalues
+        indices = L > tol
+        L = L[indices]
+        Q = Q[:, indices]
+
+        U = (M @ Q) * L.rsqrt()
+
+        if rdamping != 0:
+            rdamping *= torch.linalg.vector_norm(L) # pylint:disable=not-callable
+            L.add_(rdamping)
+
+        return U, L
+
+    except torch.linalg.LinAlgError:
+        return None, None
+
+def lm_adagrad_apply(g: torch.Tensor, U: torch.Tensor, L: torch.Tensor):
+    Z = U.T @ g
+    return (U * L.rsqrt()) @ Z
+
+def maybe_lerp_(state_: dict, beta: float | None, key, value: Any):
+    if (key not in state_) or (beta is None) or (not isinstance(value, torch.Tensor)): state_[key] = value
+    else:
+        if state_[key] is None or state_[key].shape != value.shape: state_[key] = value
+        else: state_[key].lerp_(value, 1-beta)
+
+class LMAdagrad(TensorwiseTransform):
+    """
+    Limited-memory full matrix Adagrad.
+
+    The update rule is to stack recent gradients into M, compute U, S <- SVD(M), then calculate update as U S^-1 Uᵀg.
+    But it uses eigendecomposition on MᵀM to get U and S^2 because that is faster when you don't neeed V.
+
+    This is equivalent to full-matrix Adagrad on recent gradients.
+
+    Args:
+        history_size (int, optional): number of past gradients to store. Defaults to 10.
+        update_freq (int, optional): frequency of updating the preconditioner (U and S). Defaults to 1.
+        damping (float, optional): damping value. Defaults to 1e-4.
+        rdamping (float, optional): value of damping relative to singular values norm. Defaults to 0.
+        order (int, optional):
+            order=2 means gradient differences are used in place of gradients. Higher order uses higher order differences. Defaults to 1.
+        true_damping (bool, optional):
+            If True, damping is added to squared singular values to mimic Adagrad. Defaults to True.
+        eigh (bool, optional): uses a more efficient way to calculate U and S. Defaults to True.
+        U_beta (float | None, optional): momentum for U (too unstable, don't use). Defaults to None.
+        S_beta (float | None, optional): momentum for S (too unstable, don't use). Defaults to None.
+        interval (int, optional): Interval between gradients that are added to history (2 means every second gradient is used). Defaults to 1.
+        concat_params (bool, optional): if True, treats all parameters as a single vector, meaning it will also whiten inter-parameters. Defaults to True.
+        inner (Chainable | None, optional): preconditioner will be applied to output of this module. Defaults to None.
+
+    Examples:
+        Limited-memory Adagrad
+
+        .. code-block:: python
+
+            optimizer = tz.Modular(
+                model.parameters(),
+                tz.m.LMAdagrad(),
+                tz.m.LR(0.1)
+            )
+
+        Adam with L-Adagrad preconditioner (for debiasing second beta is 0.999 arbitrarily)
+
+        .. code-block:: python
+
+            optimizer = tz.Modular(
+                model.parameters(),
+                tz.m.LMAdagrad(inner=tz.m.EMA()),
+                tz.m.Debias(0.9, 0.999),
+                tz.m.LR(0.01)
+            )
+
+        Stable Adam with L-Adagrad preconditioner (this is what I would recommend)
+
+        .. code-block:: python
+
+            optimizer = tz.Modular(
+                model.parameters(),
+                tz.m.LMAdagrad(inner=tz.m.EMA()),
+                tz.m.Debias(0.9, 0.999),
+                tz.m.ClipNormByEMA(max_ema_growth=1.2),
+                tz.m.LR(0.01)
+            )
+
+    Reference:
+        Agarwal N. et al. Efficient full-matrix adaptive regularization //International Conference on Machine Learning. – PMLR, 2019. – С. 102-110.
+    """
+
+    def __init__(
+        self,
+        history_size: int = 100,
+        update_freq: int = 1,
+        damping: float = 1e-4,
+        rdamping: float = 0,
+        order: int = 1,
+        true_damping: bool = True,
+        U_beta: float | None = None,
+        L_beta: float | None = None,
+        interval: int = 1,
+        concat_params: bool = True,
+        inner: Chainable | None = None,
+    ):
+        # history is still updated each step so Precondition's update_freq has different meaning
+        defaults = dict(history_size=history_size, update_freq=update_freq, damping=damping, rdamping=rdamping, true_damping=true_damping, order=order, U_beta=U_beta, L_beta=L_beta)
+        super().__init__(defaults, uses_grad=False, concat_params=concat_params, inner=inner, update_freq=interval)
+
+    @torch.no_grad
+    def update_tensor(self, tensor, param, grad, loss, state, setting):
+        order = setting['order']
+        history_size = setting['history_size']
+        update_freq = setting['update_freq']
+        damping = setting['damping']
+        rdamping = setting['rdamping']
+        U_beta = setting['U_beta']
+        L_beta = setting['L_beta']
+
+        if 'history' not in state: state['history'] = deque(maxlen=history_size)
+        history = state['history']
+
+        if order == 1:
+            t = tensor.clone().view(-1)
+            history.append(t)
+        else:
+
+            # if order=2, history is of gradient differences, order 3 is differences between differences, etc
+            # scaled by parameter differences
+            cur_p = param.clone()
+            cur_g = tensor.clone()
+            for i in range(1, order):
+                if f'prev_g_{i}' not in state:
+                    state[f'prev_p_{i}'] = cur_p
+                    state[f'prev_g_{i}'] = cur_g
+                    break
+
+                s = cur_p - state[f'prev_p_{i}']
+                y = cur_g - state[f'prev_g_{i}']
+                state[f'prev_p_{i}'] = cur_p
+                state[f'prev_g_{i}'] = cur_g
+                cur_p = s
+                cur_g = y
+
+                if i == order - 1:
+                    cur_g = cur_g / torch.linalg.norm(cur_p).clip(min=1e-8) # pylint:disable=not-callable
+                    history.append(cur_g.view(-1))
+
+        step = state.get('step', 0)
+        if step % update_freq == 0 and len(history) != 0:
+            U, L = lm_adagrad_update(history, damping=damping, rdamping=rdamping)
+            maybe_lerp_(state, U_beta, 'U', U)
+            maybe_lerp_(state, L_beta, 'L', L)
+
+        if len(history) != 0:
+            state['step'] = step + 1 # do not increment if no history (gathering s_ks and y_ks)
+
+    @torch.no_grad
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+        U = state.get('U', None)
+        if U is None:
+            # make a conservative step to avoid issues due to different GD scaling
+            return tensor.clip_(-0.1, 0.1) # pyright:ignore[reportArgumentType]
+
+        L = state['L']
+        update = lm_adagrad_apply(tensor.view(-1), U, L).view_as(tensor)
+
+        return update
+

torchzero/modules/optimizers/lion.py

@@ -1,7 +1,7 @@
 import torch
 
 from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 
 
 def lion_(tensors: TensorList, exp_avg_: TensorList, beta1, beta2,):
@@ -28,8 +28,8 @@ class Lion(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        beta1, beta2 = self.get_settings('beta1', 'beta2', params = params, cls=NumberList)
-        exp_avg = self.get_state('ema', params=params, cls=TensorList)
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        beta1, beta2 = unpack_dicts(settings, 'beta1', 'beta2', cls=NumberList)
+        exp_avg = unpack_states(states, tensors, 'ema', cls=TensorList)
         return lion_(TensorList(tensors),exp_avg,beta1,beta2)
 

torchzero/modules/optimizers/mars.py

@@ -0,0 +1,91 @@
+from operator import itemgetter
+from functools import partial
+
+import torch
+
+from ...core import Module, Target, Transform, apply_transform, Chainable
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+from ..functional import (
+    debias, debiased_step_size,
+    ema_,
+    sqrt_ema_sq_,
+)
+from ..step_size.lr import lazy_lr
+from ..momentum.experimental import sqrt_nag_ema_sq_
+from ..momentum.momentum import nag_
+
+
+def mars_correction_(
+    tensors_: TensorList,
+    prev_: TensorList,
+    beta: float | NumberList,
+    scaling: float | NumberList,
+    max_norm: float | NumberList |  None,
+):
+    dg = (tensors_ - prev_).mul_(scaling * beta / (1-beta))
+    prev_.copy_(tensors_)
+
+    c = tensors_.add_(dg)
+    if max_norm is not None:
+        c.clip_norm_(max=max_norm, tensorwise=False)
+
+    return c
+
+class MARSCorrection(Transform):
+    """MARS variance reduction correction.
+
+    Place any other momentum-based optimizer after this,
+    make sure :code:`beta` parameter matches with momentum in the optimizer.
+
+    Args:
+        beta (float, optional): use the same beta as you use in the momentum module. Defaults to 0.9.
+        scaling (float, optional): controls the scale of gradient correction in variance reduction. Defaults to 0.025.
+        max_norm (float, optional): clips norm of corrected gradients, None to disable. Defaults to 1.
+
+    Examples:
+        Mars-AdamW
+
+        .. code-block:: python
+
+            optimizer = tz.Modular(
+                model.parameters(),
+                tz.m.MARSCorrection(beta=0.95),
+                tz.m.Adam(beta1=0.95, beta2=0.99),
+                tz.m.WeightDecay(1e-3),
+                tz.m.LR(0.1)
+            )
+
+        Mars-Lion
+
+        .. code-block:: python
+
+            optimizer = tz.Modular(
+                model.parameters(),
+                tz.m.MARSCorrection(beta=0.9),
+                tz.m.Lion(beta1=0.9),
+                tz.m.LR(0.1)
+            )
+
+    """
+    def __init__(
+        self,
+        beta: float = 0.9,
+        scaling: float = 0.025,
+        max_norm: float | None = 1,
+    ):
+        defaults=dict(beta=beta, scaling=scaling, max_norm=max_norm)
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        prev = unpack_states(states, tensors, 'prev', init=tensors, cls=TensorList)
+        beta, scaling = unpack_dicts(settings, 'beta', 'scaling', cls=NumberList)
+        max_norm = settings[0]['max_norm']
+
+        return mars_correction_(
+            tensors_=TensorList(tensors),
+            prev_=prev,
+            beta=beta,
+            scaling=scaling,
+            max_norm=max_norm,
+        )

torchzero/modules/optimizers/msam.py

@@ -0,0 +1,186 @@
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, Target, Transform, apply_transform
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states, generic_ne
+from ..functional import ema_
+from ..momentum.momentum import nag_
+
+
+def msam_(
+    tensors: TensorList,
+    params: TensorList,
+    velocity_: TensorList,
+    momentum: float | NumberList,
+    lr: NumberList | None,
+    rho: float | NumberList,
+    weight_decay: float | NumberList,
+    nesterov: bool = False,
+    lerp: bool = False,
+
+    # inner args
+    inner: Module | None = None,
+    grads: list[torch.Tensor] | None = None,
+):
+    # weights w and wh, momentum μ, perturbation strength ρ
+    # w = wh + rho * v / ||v||
+    # v1 = μv + g
+    # w1 = w - lr*v1
+    # wh1 = w1 - rho * v1 / ||v1||
+
+    # w1 = wh + rho * v / ||v|| - lr*v1
+    # vn = rho * v / ||v||
+    # v1n = rho * v1 / ||v1||
+    # wh1 = wh + vn - lr*v1 - v1n
+
+    # the update is
+    # vn - lr*v1 - v1n
+
+    # we track ascent direction so it becomes lr*v1 + v1n - vn
+
+    # can't really decouple it from lr
+    # but at least it is now expressed as function of g
+
+    denom = (velocity_.global_vector_norm() / rho).clip(min=1e-8)
+    vn = velocity_ / denom
+
+    mom_ = nag_ if nesterov else ema_
+    velocity_ = mom_(tensors, velocity_, momentum, dampening=0, lerp=lerp)
+
+    denom = (velocity_.global_vector_norm() / rho).clip(min=1e-8)
+    v1n = velocity_ / denom
+
+    if inner is not None:
+        assert params is not None
+        inner_update = TensorList(apply_transform(inner, tensors, params=params, grads=grads))
+
+    else:
+        assert lr is not None
+        inner_update = velocity_ * lr
+
+    update = inner_update.add_(v1n).sub_(vn)
+
+    if generic_ne(weight_decay, 0):
+        wd = (params + vn).mul_(weight_decay)
+        update.add_(wd)
+
+    return update
+
+class MSAM(Transform):
+    """Momentum-SAM from https://arxiv.org/pdf/2401.12033.
+
+    This implementation expresses the update rule as function of gradient. This way it can be used as a drop-in
+    replacement for momentum strategies in other optimizers.
+
+    To combine MSAM with other optimizers in the way done in the official implementation,
+    e.g. to make Adam_MSAM, use :code:`tz.m.MSAMObjective` module.
+
+    .. note::
+        MSAM has a learning rate hyperparameter that can't really be removed from the update rule.
+        To avoid compounding learning rate mofications, remove the :code:`tz.m.LR` module if you had it.
+
+    Args:
+        lr (float): learning rate. Adding this module adds support for learning rate schedulers.
+        momentum (float, optional): momentum (beta). Defaults to 0.9.
+        rho (float, optional): perturbation strength. Defaults to 0.3.
+        weight_decay (float, optional):
+            weight decay. It is applied to perturbed parameters, so it is differnet
+            from applying :code:`tz.m.WeightDecay` after MSAM. Defaults to 0.
+        nesterov (bool, optional): whether to use nesterov momentum formula. Defaults to False.
+        lerp (bool, optional):
+            whether to use linear interpolation, if True, this becomes similar to exponential moving average. Defaults to False.
+
+    Examples:
+        MSAM
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.MSAM(1e-3)
+            )
+
+        Adam with MSAM instead of exponential average. Note that this is different from Adam_MSAM.
+        To make Adam_MSAM and such, use the :code:`tz.m.MSAMObjective` module.
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.RMSprop(0.999, inner=tz.m.MSAM(1e-3)),
+                tz.m.Debias(0.9, 0.999),
+            )
+    """
+    USES_LR = True
+    def __init__(self, lr: float, momentum:float=0.9, rho:float=0.3,  weight_decay:float=0, nesterov=False, lerp=False,):
+        defaults = dict(momentum=momentum,rho=rho, nesterov=nesterov, lerp=lerp, weight_decay=weight_decay)
+        if self.USES_LR: defaults['lr'] = lr
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        velocity = unpack_states(states, tensors, 'velocity', cls=TensorList)
+        s = self.settings[params[0]]
+        lerp = s['lerp']
+        nesterov = s['nesterov']
+
+        if self.USES_LR:
+            lr, momentum, rho, weight_decay = unpack_dicts(settings, 'lr','momentum','rho','weight_decay', cls=NumberList)
+
+        else:
+            lr=None
+            momentum,rho,weight_decay = unpack_dicts(settings, 'momentum','rho','weight_decay', cls=NumberList)
+
+        return msam_(
+            TensorList(tensors),
+            params=TensorList(params),
+            velocity_=velocity,
+            momentum=momentum,
+            lr=lr,
+            rho=rho,
+            weight_decay=weight_decay,
+            nesterov=nesterov,
+            lerp=lerp,
+
+            # inner args
+            inner=self.children.get("modules", None),
+            grads=grads,
+        )
+
+
+class MSAMObjective(MSAM):
+    """Momentum-SAM from https://arxiv.org/pdf/2401.12033.
+
+    .. note::
+        Please make sure to place :code:`tz.m.LR` inside the :code:`modules` argument. For example,
+        :code:`tz.m.MSAMObjective([tz.m.Adam(), tz.m.LR(1e-3)])`. Putting LR after MSAM will lead
+        to an incorrect update rule.
+
+    Args:
+        modules (Chainable): modules that will optimizer the MSAM objective. Make sure :code:`tz.m.LR` is one of them.
+        momentum (float, optional): momentum (beta). Defaults to 0.9.
+        rho (float, optional): perturbation strength. Defaults to 0.3.
+        nesterov (bool, optional): whether to use nesterov momentum formula. Defaults to False.
+        lerp (bool, optional):
+            whether to use linear interpolation, if True, MSAM momentum becomes similar to exponential moving average.
+            Defaults to False.
+
+    Examples:
+        AdamW-MSAM
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                bench.parameters(),
+                tz.m.MSAMObjective(
+                    [tz.m.Adam(), tz.m.WeightDecay(1e-3), tz.m.LR(1e-3)],
+                    rho=1.
+                )
+            )
+    """
+    USES_LR = False
+    def __init__(self, modules: Chainable, momentum:float=0.9, rho:float=0.3, weight_decay:float=0, nesterov=False, lerp=False):
+        super().__init__(lr=0, momentum=momentum, rho=rho, weight_decay=weight_decay, nesterov=nesterov, lerp=lerp)
+        self.set_child('modules', modules)
+

torchzero/modules/optimizers/muon.py

@@ -19,6 +19,7 @@ def _is_at_least_2d(p: torch.Tensor):
 
 # stolen from:
 # https://github.com/KellerJordan/Muon/blob/master/muon.py
+# actually at this stage its a frankenstein
 @enable_compilation
 def zeropower_via_newtonschulz5(G: torch.Tensor, steps: int) -> torch.Tensor:
     """
@@ -152,7 +153,7 @@ class Orthogonalize(TensorwiseTransform):
     The Muon page says that embeddings and classifier heads should not be orthogonalized.
     Usually only matrix parameters that are directly used in matmuls should be orthogonalized.
 
-    To make Muon, use Split with Adam on 1d params: TODO code example.
+    To make Muon, use Split with Adam on 1d params
 
     Args:
         ns_steps (int, optional):
@@ -164,7 +165,31 @@ class Orthogonalize(TensorwiseTransform):
         method (str, optional):
             Newton-Schulz is very fast, SVD is extremely slow but can be slighly more precise.
         target (str, optional):
-            what to set on vars.
+            what to set on var.
+
+
+    Examples:
+        standard Muon with Adam fallback
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.head.parameters(),
+                tz.m.Split(
+                    # apply muon only to 2D+ parameters
+                    filter = lambda t: t.ndim >= 2,
+                    true = [
+                        tz.m.HeavyBall(),
+                        tz.m.Orthogonalize(),
+                        tz.m.LR(1e-2),
+                    ],
+                    false = tz.m.Adam()
+                ),
+                tz.m.LR(1e-2)
+            )
+
+    Reference:
+        Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cesista, Laker Newhouse, Jeremy Bernstein - Muon: An optimizer for hidden layers in neural networks (2024) https://github.com/KellerJordan/Muon
     """
     def __init__(self, ns_steps=5, adjust_lr=False, dual_norm_correction=False,
                  method: Literal['newton-schulz', 'svd'] = 'newton-schulz', target:Target='update'):
@@ -172,9 +197,9 @@ class Orthogonalize(TensorwiseTransform):
         super().__init__(uses_grad=False, defaults=defaults, target=target)
 
     @torch.no_grad
-    def transform(self, tensor, param, grad, vars):
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
         orthogonalize, ns_steps, dual_norm_correction, adjust_lr, method = itemgetter(
-            'orthogonalize', 'ns_steps', 'dual_norm_correction', 'adjust_lr', 'method')(self.settings[param])
+            'orthogonalize', 'ns_steps', 'dual_norm_correction', 'adjust_lr', 'method')(setting)
 
         if not orthogonalize: return tensor
 
@@ -199,7 +224,7 @@ class DualNormCorrection(TensorwiseTransform):
     def __init__(self, target: Target='update'):
         super().__init__({}, uses_grad=True, target=target)
 
-    def transform(self, tensor, param, grad, vars):
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
         assert grad is not None
         if (tensor.ndim >= 2) and (tensor.size(0) > 1) and (tensor.size(1) > 1):
             return _dual_norm_correction(tensor, grad, batch_first=False)
@@ -213,8 +238,8 @@ class MuonAdjustLR(Transform):
         defaults = dict(alpha=alpha)
         super().__init__(defaults=defaults, uses_grad=False, target=target)
 
-    def transform(self, tensors, params, grads, vars):
-        alphas = self.get_settings('alpha', params=params)
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        alphas = [s['alpha'] for s in settings]
         tensors_alphas = [(t, adjust_lr_for_muon(a, t.shape)) for t, a in zip(tensors, alphas) if _is_at_least_2d(t)]
         tensors = [i[0] for i in tensors_alphas]
         a = [i[1] for i in alphas]

torchzero/modules/optimizers/orthograd.py

@@ -30,16 +30,15 @@ class OrthoGrad(Transform):
     Args:
         eps (float, optional): epsilon added to the denominator for numerical stability (default: 1e-30)
         renormalize (bool, optional): whether to graft projected gradient to original gradient norm. Defaults to True.
-        target (Target, optional): what to set on vars. Defaults to 'update'.
+        target (Target, optional): what to set on var. Defaults to 'update'.
     """
     def __init__(self, eps: float = 1e-8, renormalize=True, target: Target = 'update'):
         defaults = dict(eps=eps, renormalize=renormalize)
         super().__init__(defaults, uses_grad=False, target=target)
 
-    def transform(self, tensors, params, grads, vars):
-        settings = self.settings[params[0]]
-        eps = settings['eps']
-        renormalize = settings['renormalize']
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        eps = settings[0]['eps']
+        renormalize = settings[0]['renormalize']
 
         params = as_tensorlist(params)
         target = as_tensorlist(tensors)

torchzero/modules/optimizers/rmsprop.py

@@ -3,8 +3,8 @@ from typing import Literal
 
 import torch
 
-from ...core import Module, Target, Transform, Chainable, Vars, apply
-from ...utils import NumberList, TensorList
+from ...core import Module, Target, Transform, Chainable, Var, apply_transform
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 from ..functional import sqrt_centered_ema_sq_, sqrt_ema_sq_
 
 
@@ -23,7 +23,6 @@ def rmsprop_(
     inner: Module | None = None,
     params: list[torch.Tensor] | None = None,
     grads: list[torch.Tensor] | None = None,
-    vars: Vars | None = None,
 ):
     """returns `tensors_`"""
     if exp_avg_ is not None:
@@ -36,12 +35,14 @@ def rmsprop_(
 
     if inner is not None:
         assert params is not None
-        tensors_ = TensorList(apply(inner, tensors_, params=params, grads=grads, vars=vars))
+        tensors_ = TensorList(apply_transform(inner, tensors_, params=params, grads=grads))
 
     return tensors_.div_(sqrt_exp_avg_sq.add_(eps))
 
 class RMSprop(Transform):
-    """Divides graient by EMA of gradient squares. Matches pytorch RMSprop if "init" is set to "zeros".
+    """Divides graient by EMA of gradient squares.
+
+    This implementation is identical to :code:`torch.optim.RMSprop`.
 
     Args:
         smoothing (float, optional): beta for exponential moving average of gradient squares. Defaults to 0.99.
@@ -51,7 +52,8 @@ class RMSprop(Transform):
         amsgrad (bool, optional): Whether to divide by maximum of EMA of gradient squares instead. Defaults to False.
         pow (float, optional): power used in second momentum power and root. Defaults to 2.
         init (str, optional): how to initialize EMA, either "update" to use first update or "zeros". Defaults to "update".
-        inner (Chainable | None, optional): Inner modules that are applied after updating EMA and before preconditioning. Defaults to None.
+        inner (Chainable | None, optional):
+            Inner modules that are applied after updating EMA and before preconditioning. Defaults to None.
     """
     def __init__(
         self,
@@ -61,26 +63,25 @@ class RMSprop(Transform):
         debiased: bool = False,
         amsgrad: bool = False,
         pow: float = 2,
-        init: Literal["zeros", "update"] = "update",
+        init: Literal["zeros", "update"] = "zeros",
         inner: Chainable | None = None,
     ):
         defaults = dict(smoothing=smoothing,eps=eps,centered=centered,debiased=debiased,amsgrad=amsgrad,pow=pow,init=init)
         super().__init__(defaults=defaults, uses_grad=False)
-        self.current_step = 0
+
         if inner is not None:
             self.set_child('inner', inner)
 
-    def transform(self, tensors, params, grads, vars):
-        self.current_step += 1
-
-        smoothing,eps = self.get_settings('smoothing', 'eps', params=params, cls=NumberList)
-        centered,debiased,amsgrad,pow,init = itemgetter('centered','debiased','amsgrad','pow','init')(self.settings[params[0]])
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+        smoothing, eps = unpack_dicts(settings, 'smoothing', 'eps', cls=NumberList)
+        centered, debiased, amsgrad, pow, init = itemgetter('centered','debiased','amsgrad','pow','init')(settings[0])
 
-        exp_avg_sq = self.get_state('exp_avg_sq', params=params, cls=TensorList)
-        exp_avg = self.get_state('exp_avg', params=params, cls=TensorList) if centered else None
-        max_exp_avg_sq = self.get_state('max_exp_avg_sq', params=params, cls=TensorList) if amsgrad else None
+        exp_avg_sq = unpack_states(states, tensors, 'exp_avg_sq', cls=TensorList)
+        exp_avg = unpack_states(states, tensors, 'exp_avg', cls=TensorList) if centered else None
+        max_exp_avg_sq = unpack_states(states, tensors, 'max_exp_avg_sq', cls=TensorList) if amsgrad else None
 
-        if init == 'update' and self.current_step == 1:
+        if init == 'update' and step == 1:
             exp_avg_sq.set_([t**2 for t in tensors])
             if exp_avg is not None: exp_avg.set_([t.clone() for t in tensors])
 
@@ -90,7 +91,7 @@ class RMSprop(Transform):
             smoothing=smoothing,
             eps=eps,
             debiased=debiased,
-            step=self.current_step,
+            step=step,
             exp_avg_=exp_avg,
             max_exp_avg_sq_=max_exp_avg_sq,
             pow=pow,
@@ -99,5 +100,4 @@ class RMSprop(Transform):
             inner=self.children.get("inner", None),
             params=params,
             grads=grads,
-            vars=vars,
         )