torchzero 0.3.10__py3-none-any.whl → 0.3.13__py3-none-any.whl

torchzero/modules/adaptive/adan.py

@@ -0,0 +1,96 @@
+import torch
+
+from ...core import Transform
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+
+def adan_(
+    g: TensorList,
+    g_prev_: TensorList,
+    m_: TensorList, # exponential moving average
+    v_: TensorList, # exponential moving average of gradient differences
+    n_: TensorList, # kinda like squared momentum
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    beta3: float | NumberList,
+    eps: float | NumberList,
+    step: int,
+):
+    """Returns new tensors"""
+    m_.lerp_(g, 1 - beta1)
+
+    if step == 1:
+        term = g
+    else:
+        diff = g - g_prev_
+        v_.lerp_(diff, 1 - beta2)
+        term = g + beta2 * diff
+
+    n_.mul_(beta3).addcmul_(term, term, value=(1 - beta3))
+
+    m = m_ / (1.0 - beta1**step)
+    v = v_ / (1.0 - beta2**step)
+    n = n_ / (1.0 - beta3**step)
+
+    denom = n.sqrt_().add_(eps)
+    num = m + beta2 * v
+
+    update = num.div_(denom)
+    g_prev_.copy_(g)
+
+    return update
+
+
+
+class Adan(Transform):
+    """Adaptive Nesterov Momentum Algorithm from https://arxiv.org/abs/2208.06677
+
+    Args:
+        beta1 (float, optional): momentum. Defaults to 0.98.
+        beta2 (float, optional): momentum for gradient differences. Defaults to 0.92.
+        beta3 (float, optional): thrid (squared) momentum. Defaults to 0.99.
+        eps (float, optional): epsilon. Defaults to 1e-8.
+        use_n_prev (bool, optional):
+            whether to use previous gradient differences momentum.
+
+    Example:
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.Adan(),
+        tz.m.LR(1e-3),
+    )
+    Reference:
+        Xie, X., Zhou, P., Li, H., Lin, Z., & Yan, S. (2024). Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. IEEE Transactions on Pattern Analysis and Machine Intelligence. https://arxiv.org/abs/2208.06677
+    """
+    def __init__(
+        self,
+        beta1: float = 0.98,
+        beta2: float = 0.92,
+        beta3: float = 0.99,
+        eps: float = 1e-8,
+    ):
+        defaults=dict(beta1=beta1,beta2=beta2,beta3=beta3,eps=eps)
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = TensorList(tensors)
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        beta1,beta2,beta3,eps=unpack_dicts(settings, 'beta1','beta2','beta3','eps', cls=NumberList)
+        g_prev, m, v, n = unpack_states(states, tensors, 'g_prev','m','v','n', cls=TensorList)
+
+        update = adan_(
+            g=tensors,
+            g_prev_=g_prev,
+            m_=m,
+            v_=v,
+            n_=n,
+            beta1=beta1,
+            beta2=beta2,
+            beta3=beta3,
+            eps=eps,
+            step=step,
+        )
+
+        return update

torchzero/modules/adaptive/adaptive_heavyball.py

@@ -0,0 +1,54 @@
+import torch
+from ...core import  Transform
+from ...utils import TensorList, unpack_dicts, unpack_states
+
+
+def adaptive_heavy_ball(f, f_star, f_prev, g: TensorList, g_prev: TensorList, p: TensorList, p_prev: TensorList):
+    if f - f_star <= torch.finfo(p[0].dtype).tiny * 2: return g
+
+    g_g = g.dot(g)
+    g_gp = g.dot(g_prev)
+    num = -(f - f_star) * g.dot(g_prev)
+    denom = (f_prev - f_star) * g_g + (f - f_star) * g_gp
+    m = num/denom
+
+    h = 2*(f - f_star) / g_g
+    return (1 + m) * h * g - m*(p-p_prev)
+
+
+class AdaptiveHeavyBall(Transform):
+    """Adaptive heavy ball from https://hal.science/hal-04832983v1/file/OJMO_2024__5__A7_0.pdf.
+
+    This is related to conjugate gradient methods, it may be very good for non-stochastic convex objectives, but won't work on stochastic ones.
+
+    note:
+        The step size is determined by the algorithm, so learning rate modules shouldn't be used.
+
+    Args:
+        f_star (int, optional):
+            (estimated) minimal possible value of the objective function (lowest possible loss). Defaults to 0.
+    """
+    def __init__(self, f_star: float = 0):
+        defaults = dict(f_star=f_star)
+        super().__init__(defaults, uses_grad=False, uses_loss=True)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        assert loss is not None
+        tensors = TensorList(tensors)
+        f_star = self.defaults['f_star']
+
+        f_prev = self.global_state.get('f_prev', None)
+        p_prev, g_prev = unpack_states(states, tensors, 'p_prev', 'g_prev', init=[params,tensors], cls=TensorList)
+
+        if f_prev is None:
+            self.global_state['f_prev'] = loss
+            h = 2*(loss - f_star) / tensors.dot(tensors)
+            return h * tensors
+
+        update = adaptive_heavy_ball(f=loss, f_star=f_star, f_prev=f_prev, g=tensors, g_prev=g_prev, p=TensorList(params), p_prev=p_prev)
+
+        self.global_state['f_prev'] = loss
+        p_prev.copy_(params)
+        g_prev.copy_(tensors)
+        return update

torchzero/modules/adaptive/aegd.py

@@ -0,0 +1,54 @@
+import math
+
+import torch
+
+from ...core import Transform
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+
+# i've verified, it is identical to official
+# https://github.com/txping/AEGD/blob/master/aegd.py
+def aegd_(f: torch.Tensor | float, g: TensorList, r_: TensorList, c:float|NumberList=1, eta:float|NumberList=0.1) -> TensorList:
+    v = g / (2 * (f + c)**0.5)
+    r_ /= 1 + (v ** 2).mul_(2*eta) # update energy
+    return 2*eta * r_*v # pyright:ignore[reportReturnType]
+
+class AEGD(Transform):
+    """AEGD (Adaptive gradient descent with energy) from https://arxiv.org/abs/2010.05109#page=10.26.
+
+    Note:
+        AEGD has a learning rate hyperparameter that can't really be removed from the update rule.
+        To avoid compounding learning rate mofications, remove the ``tz.m.LR`` module if you had it.
+
+    Args:
+        eta (float, optional): step size. Defaults to 0.1.
+        c (float, optional): c. Defaults to 1.
+        beta3 (float, optional): thrid (squared) momentum. Defaults to 0.1.
+        eps (float, optional): epsilon. Defaults to 1e-8.
+        use_n_prev (bool, optional):
+            whether to use previous gradient differences momentum.
+    """
+    def __init__(
+        self,
+        lr: float = 0.1,
+        c: float = 1,
+    ):
+        defaults=dict(c=c,lr=lr)
+        super().__init__(defaults, uses_loss=True)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        assert loss is not None
+        tensors = TensorList(tensors)
+
+        c,lr=unpack_dicts(settings, 'c','lr', cls=NumberList)
+        r = unpack_states(states, tensors, 'r', init=lambda t: torch.full_like(t, float(loss+c[0])**0.5), cls=TensorList)
+
+        update = aegd_(
+            f=loss,
+            g=tensors,
+            r_=r,
+            c=c,
+            eta=lr,
+        )
+
+        return update

torchzero/modules/adaptive/esgd.py

@@ -0,0 +1,171 @@
+import math
+from collections.abc import Callable
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, Target, Transform, apply_transform
+from ...utils import NumberList, TensorList, as_tensorlist
+
+
+def esgd_(
+    tensors_: TensorList,
+    D: TensorList | None,
+    D_sq_acc_: TensorList,
+    damping: float | NumberList,
+    update_freq: int,
+    step: int,
+    i: int,
+):
+    # update preconditioner
+    if step % update_freq == 0:
+        assert D is not None
+        D_sq_acc_.addcmul_(D, D)
+        i += 1
+    else:
+        assert D is None
+
+    denom = (D_sq_acc_ / max(i, 1)).sqrt_().add_(damping)
+    return tensors_.div_(denom), i
+
+
+class ESGD(Module):
+    """Equilibrated Gradient Descent (https://arxiv.org/abs/1502.04390)
+
+    This is similar to Adagrad, but the accumulates squared randomized hessian diagonal estimates instead of squared gradients.
+
+    .. note::
+        In most cases Adagrad should be the first module in the chain because it relies on autograd. Use the :code:`inner` argument if you wish to apply Adagrad preconditioning to another module's output.
+
+    .. note::
+        If you are using gradient estimators or reformulations, set :code:`hvp_method` to "forward" or "central".
+
+    .. note::
+        This module requires a closure passed to the optimizer step,
+        as it needs to re-evaluate the loss and gradients for calculating HVPs.
+        The closure must accept a ``backward`` argument (refer to documentation).
+
+    Args:
+        damping (float, optional): added to denominator for stability. Defaults to 1e-4.
+        update_freq (int, optional):
+            frequency of updating hessian diagonal estimate via a hessian-vector product.
+            This value can be increased to reduce computational cost. Defaults to 20.
+        hvp_method (str, optional):
+            Determines how Hessian-vector products are evaluated.
+
+            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
+              This requires creating a graph for the gradient.
+            - ``"forward"``: Use a forward finite difference formula to
+              approximate the HVP. This requires one extra gradient evaluation.
+            - ``"central"``: Use a central finite difference formula for a
+              more accurate HVP approximation. This requires two extra
+              gradient evaluations.
+            Defaults to "autograd".
+        fd_h (float, optional): finite difference step size if :code:`hvp_method` is "forward" or "central". Defaults to 1e-3.
+        n_samples (int, optional):
+            number of hessian-vector products with random vectors to evaluate each time when updating
+            the preconditioner. Larger values may lead to better hessian diagonal estimate. Defaults to 1.
+        seed (int | None, optional): seed for random vectors. Defaults to None.
+        inner (Chainable | None, optional):
+            Inner module. If this is specified, operations are performed in the following order.
+            1. compute hessian diagonal estimate.
+            2. pass inputs to :code:`inner`.
+            3. momentum and preconditioning are applied to the ouputs of :code:`inner`.
+
+    Examples:
+        Using ESGD:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.ESGD(),
+                tz.m.LR(0.1)
+            )
+
+        ESGD preconditioner can be applied to any other module by passing it to the :code:`inner` argument. Here is an example of applying
+        ESGD preconditioning to nesterov momentum (:code:`tz.m.NAG`):
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.ESGD(beta1=0, inner=tz.m.NAG(0.9)),
+                tz.m.LR(0.1)
+            )
+
+    """
+    def __init__(
+        self,
+        damping: float = 1e-4,
+        update_freq: int = 20,
+        hvp_method: Literal['autograd', 'forward', 'central'] = 'autograd',
+        fd_h: float = 1e-3,
+        n_samples = 1,
+        seed: int | None = None,
+        inner: Chainable | None = None
+    ):
+        defaults = dict(damping=damping, update_freq=update_freq, hvp_method=hvp_method, n_samples=n_samples, fd_h=fd_h, seed=seed)
+        super().__init__(defaults)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, var):
+        params = var.params
+        settings = self.settings[params[0]]
+        hvp_method = settings['hvp_method']
+        fd_h = settings['fd_h']
+        update_freq = settings['update_freq']
+        n_samples = settings['n_samples']
+
+        seed = settings['seed']
+        generator = None
+        if seed is not None:
+            if 'generator' not in self.global_state:
+                self.global_state['generator'] = torch.Generator(params[0].device).manual_seed(seed)
+            generator = self.global_state['generator']
+
+        damping = self.get_settings(params, 'damping', cls=NumberList)
+        D_sq_acc = self.get_state(params, 'D_sq_acc', cls=TensorList)
+        i = self.global_state.get('i', 0)
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        closure = var.closure
+        assert closure is not None
+
+        D = None
+        if step % update_freq == 0:
+
+            rgrad=None
+            for j in range(n_samples):
+                u = [torch.randn(p.size(), generator=generator, device=p.device, dtype=p.dtype) for p in params]
+
+                Hvp, rgrad = self.Hvp(u, at_x0=True, var=var, rgrad=rgrad, hvp_method=hvp_method,
+                                     h=fd_h, normalize=True, retain_grad=j < n_samples-1)
+
+                if D is None: D = Hvp
+                else: torch._foreach_add_(D, Hvp)
+
+            assert D is not None
+            if n_samples > 1: torch._foreach_div_(D, n_samples)
+
+            D = TensorList(D)
+
+        update = var.get_update()
+        if 'inner' in self.children:
+            update = apply_transform(self.children['inner'], tensors=update, params=params, grads=var.grad, var=var)
+
+        var.update, self.global_state['i'] = esgd_(
+            tensors_=TensorList(update),
+            D=TensorList(D) if D is not None else None,
+            D_sq_acc_=D_sq_acc,
+            damping=damping,
+            update_freq=update_freq,
+            step=step,
+            i=i,
+        )
+        return var

torchzero/modules/{optimizers → adaptive}/lion.py

@@ -28,7 +28,7 @@ class Lion(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    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/{experimental/spectral.py → adaptive/lmadagrad.py}

@@ -1,55 +1,57 @@
-from abc import ABC, abstractmethod
-import math
 from collections import deque
 from typing import Literal, Any
-import itertools
+import warnings
 
 import torch
 from ...core import Chainable, TensorwiseTransform
-from ...utils.linalg.matrix_funcs import matrix_power_eigh
-from ...utils.linalg.svd import randomized_svd
-from ...utils.linalg.qr import qr_householder
 
-def spectral_update(history, damping, rdamping, true_damping: bool):
-    M_hist = torch.stack(tuple(history), dim=1)
-    device = M_hist.device
-    M_hist = M_hist.cuda()
+def lm_adagrad_update(history: deque[torch.Tensor] | torch.Tensor, damping, rdamping):
+    if isinstance(history, torch.Tensor):
+        M = history
+    else:
+        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:
-        U, S, _ = torch.linalg.svd(M_hist, full_matrices=False, driver='gesvda') # pylint:disable=not-callable
-        U = U.to(device); S = S.to(device)
+        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]
 
-        if damping != 0 or rdamping != 0:
-            if rdamping != 0: rdamping *= torch.linalg.vector_norm(S) # pylint:disable=not-callable
-            Iu = damping + rdamping
-            if true_damping:
-                S.pow_(2)
-                Iu **= 2
-            S.add_(Iu)
-            if true_damping: S.sqrt_()
+        U = (M @ Q) * L.rsqrt()
 
-        return U, 1/S
+        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 spectral_apply(g: torch.Tensor, U: torch.Tensor, S_inv: torch.Tensor):
-    Utg = (U.T @ g)*S_inv
-    return U @ Utg
-
+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].shape != value.shape: state_[key] = value
+        if state_[key] is None or state_[key].shape != value.shape: state_[key] = value
         else: state_[key].lerp_(value, 1-beta)
 
-class SpectralPreconditioner(TensorwiseTransform):
+class LMAdagrad(TensorwiseTransform):
     """
-    The update rule is to stack recent gradients into M, compute U, S <- SVD(M), then calculate U (Uᵀg)/S.
-    This is equivalent to full matrix Adagrad with accumulator initialized to zeros,
-    except only recent :code:`history_size` gradients are used.
-    However this doesn't require N^2 memory and is computationally less expensive than Shampoo.
+    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.
@@ -61,54 +63,81 @@ class SpectralPreconditioner(TensorwiseTransform):
         true_damping (bool, optional):
             If True, damping is added to squared singular values to mimic Adagrad. 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 1/S (too unstable, don't use). Defaults to None.
+        L_beta (float | None, optional): momentum for L (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 False.
-        normalize (bool, optional): whether to normalize gradients, this doesn't work well so don't use it. Defaults to False.
-        centralize (bool, optional): whether to centralize gradients, this doesn't work well so don't use it. Defaults to False.
+        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
+
+    ```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)
+
+    ```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)
+
+    ```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 = 10,
+        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,
-        S_beta: float | None = None,
+        L_beta: float | None = None,
         interval: int = 1,
-        concat_params: bool = False,
-        normalize: bool=False,
-        centralize:bool = False,
+        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, S_beta=S_beta, normalize=normalize, centralize=centralize)
+        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, settings):
-        order = settings['order']
-        history_size = settings['history_size']
-        update_freq = settings['update_freq']
-        damping = settings['damping']
-        rdamping = settings['rdamping']
-        true_damping = settings['true_damping']
-        U_beta = settings['U_beta']
-        S_beta = settings['S_beta']
-        normalize = settings['normalize']
-        centralize = settings['centralize']
+    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)
-            if centralize: t -= t.mean()
-            if normalize: t /= torch.linalg.vector_norm(t).clip(min=1e-8) # pylint:disable=not-callable
             history.append(t)
         else:
 
@@ -116,48 +145,42 @@ class SpectralPreconditioner(TensorwiseTransform):
             # scaled by parameter differences
             cur_p = param.clone()
             cur_g = tensor.clone()
+            eps = torch.finfo(cur_p.dtype).tiny * 2
             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_k = cur_p - state[f'prev_p_{i}']
-                y_k = cur_g - state[f'prev_g_{i}']
+                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_k
-                cur_g = y_k
+                cur_p = s
+                cur_g = y
 
                 if i == order - 1:
-                    if centralize: cur_g = cur_g - cur_g.mean()
-                    if normalize: cur_g = cur_g / torch.linalg.vector_norm(cur_g).clip(min=1e-8) # pylint:disable=not-callable
-                    else: cur_g = cur_g / torch.linalg.norm(cur_p).clip(min=1e-8) # pylint:disable=not-callable
+                    cur_g = cur_g / torch.linalg.norm(cur_p).clip(min=eps) # 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, S_inv = spectral_update(history, damping=damping, rdamping=rdamping, true_damping=true_damping)
+            U, L = lm_adagrad_update(history, damping=damping, rdamping=rdamping)
             maybe_lerp_(state, U_beta, 'U', U)
-            maybe_lerp_(state, S_beta, 'S_inv', S_inv)
+            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, settings):
-        history_size = settings['history_size']
-
+    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]
 
-        S_inv = state['S_inv']
-        update = spectral_apply(tensor.view(-1), U, S_inv).view_as(tensor)
+        L = state['L']
+        update = lm_adagrad_apply(tensor.view(-1), U, L).view_as(tensor)
 
-        n = len(state['history'])
-        mh = min(history_size, 10)
-        if n <= mh: update.mul_(n/mh)
         return update