torchzero 0.1.8__py3-none-any.whl → 0.3.2__py3-none-any.whl

torchzero/modules/experimental/dsoap.py

@@ -0,0 +1,290 @@
+from operator import itemgetter
+
+import torch
+
+from ...core import Chainable, Transform, apply
+from ...modules.optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
+
+@torch.no_grad
+def update_soap_covariances_(
+    grad: torch.Tensor,
+    GGs_: list[torch.Tensor | None],
+    beta: float | None,
+):
+    for i, GG in enumerate(GGs_):
+        if GG is None: continue
+
+        axes = list(range(i)) + list(range(i + 1, grad.ndim)) # this works fine with 1d params
+        if beta is None: GG.add_(torch.tensordot(grad, grad, (axes, axes))) # pyright:ignore[reportArgumentType]
+        else: GG.lerp_(torch.tensordot(grad, grad, (axes, axes)), 1-beta) # pyright:ignore[reportArgumentType]
+
+@torch.no_grad
+def project(tensors: torch.Tensor, Q: list[torch.Tensor | None]):
+    """
+    Projects the gradient to the eigenbases of the preconditioner.
+    """
+    for mat in Q:
+        if mat is None: continue
+        if len(mat) > 0:
+            tensors = torch.tensordot(tensors, mat, dims=[[0], [0]]) # pyright:ignore[reportArgumentType]
+        else:
+            # I don't understand this part but it is in https://github.com/nikhilvyas/SOAP/blob/main/soap.py
+            permute_order = list(range(1, len(tensors.shape))) + [0]
+            tensors = tensors.permute(permute_order)
+
+    return tensors
+
+@torch.no_grad
+def project_back(tensors: torch.Tensor, Q: list[torch.Tensor| None]):
+    """
+    Projects the gradient back to the original space.
+    """
+    for mat in Q:
+        if mat is None: continue
+        if len(mat) > 0:
+            tensors = torch.tensordot(tensors, mat,dims=[[0], [1]]) # pyright:ignore[reportArgumentType]
+        else:
+            permute_order = list(range(1, len(tensors.shape))) + [0]
+            tensors = tensors.permute(permute_order)
+
+    return tensors
+
+# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py
+@torch.no_grad
+def get_orthogonal_matrix(mat: list[torch.Tensor | None]):
+    """
+    Computes the eigenbases of the preconditioner using torch.linalg.eigh decomposition.
+    """
+    matrix = []
+    float_data = False
+    original_type = original_device = None
+    for m in mat:
+        if m is None: continue
+        if len(m) == 0:
+            matrix.append([])
+            continue
+        if m.dtype != torch.float:
+            original_type = m.dtype
+            original_device = m.device
+            matrix.append(m.float())
+        else:
+            float_data = True
+            matrix.append(m)
+
+    final = []
+    for m in matrix:
+        if len(m) == 0:
+            final.append([])
+            continue
+        try:
+            _, Q = torch.linalg.eigh(m+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
+        except Exception:
+            _, Q = torch.linalg.eigh(m.to(torch.float64)+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
+            Q = Q.to(m.dtype)
+        Q = torch.flip(Q, [1])
+
+        if not float_data:
+            Q = Q.to(original_device).type(original_type)
+        final.append(Q)
+    return final
+
+# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py#L240
+@torch.no_grad
+def get_orthogonal_matrix_QR(exp_avg_sq: torch.Tensor, GG: list[torch.Tensor | None], Q_list: list[torch.Tensor | None]):
+    """
+    Computes the eigenbases of the preconditioner using one round of power iteration
+    followed by torch.linalg.qr decomposition.
+    """
+    matrix = []
+    orth_matrix = []
+    float_data = False
+    original_type = original_device = None
+    for m,o in zip(GG, Q_list):
+        if m is None: continue
+        assert o is not None
+
+        if len(m) == 0:
+            matrix.append([])
+            orth_matrix.append([])
+            continue
+        if m.data.dtype != torch.float:
+            original_type = m.data.dtype
+            original_device = m.data.device
+            matrix.append(m.data.float())
+            orth_matrix.append(o.data.float())
+        else:
+            float_data = True
+            matrix.append(m.data.float())
+            orth_matrix.append(o.data.float())
+
+    final = []
+    for ind, (m,o) in enumerate(zip(matrix, orth_matrix)):
+        if len(m)==0:
+            final.append([])
+            continue
+        est_eig = torch.diag(o.T @ m @ o)
+        sort_idx = torch.argsort(est_eig, descending=True)
+        exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
+        o = o[:,sort_idx]
+        power_iter = m @ o
+        Q, _ = torch.linalg.qr(power_iter) # pylint:disable=not-callable
+
+        if not float_data:
+            Q = Q.to(original_device).type(original_type)
+        final.append(Q)
+
+    return final, exp_avg_sq
+
+class DSOAP(Transform):
+    """SOAP but uses scaled gradient differences
+
+    new args
+
+    scale by s whether to scale gradient differences by parameter differences
+
+    y_to_ema2 whether to use gradient differences for exponential moving average too
+    """
+    def __init__(
+        self,
+        beta1: float = 0.95,
+        beta2: float = 0.95,
+        shampoo_beta: float | None = 0.95,
+        precond_freq: int = 10,
+        merge_small: bool = True,
+        max_dim: int = 2_000,
+        precondition_1d: bool = True,
+        eps: float = 1e-8,
+        decay: float | None = None,
+        alpha: float = 1,
+        bias_correction: bool = True,
+        scale_by_s: bool = True,
+        y_to_ema2: bool = False,
+    ):
+        defaults = dict(
+            beta1=beta1,
+            beta2=beta2,
+            shampoo_beta=shampoo_beta,
+            precond_freq=precond_freq,
+            merge_small=merge_small,
+            max_dim=max_dim,
+            precondition_1d=precondition_1d,
+            eps=eps,
+            decay=decay,
+            bias_correction=bias_correction,
+            alpha=alpha,
+            scale_by_s=scale_by_s,
+            y_to_ema2=y_to_ema2,
+        )
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        updates = []
+        # update preconditioners
+        for i,(p,t) in enumerate(zip(params, tensors)):
+            state = self.state[p]
+            settings = self.settings[p]
+            beta1, beta2, shampoo_beta, merge_small, max_dim, precondition_1d, eps, alpha = itemgetter(
+                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'alpha')(settings)
+            scale_by_s = settings['scale_by_s']
+            y_to_ema2 = settings['y_to_ema2']
+
+            if merge_small:
+                t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
+
+            if 'g_prev' not in state:
+                state['p_prev'] = p.clone()
+                state['g_prev'] = t.clone()
+                updates.append(tensors[i].sign())
+                continue
+
+            p_prev = state['p_prev']
+            g_prev = state['g_prev']
+            s = p - p_prev
+            y = t - g_prev
+            if scale_by_s: y /= torch.linalg.norm(s).clip(min=1e-8) # pylint:disable=not-callable
+
+            state['p_prev'].copy_(p)
+            state['g_prev'].copy_(t)
+
+            # initialize state on 1st step
+            if 'GG' not in state:
+                state["exp_avg"] = torch.zeros_like(t)
+                if y_to_ema2: state["exp_avg_sq"] = torch.ones_like(t)
+                else: state["exp_avg_sq"] = torch.zeros_like(t)
+
+                if not precondition_1d and t.ndim <= 1:
+                    state['GG'] = []
+
+                else:
+                    state['GG'] = [torch.zeros(sh, sh, dtype=t.dtype, device=t.device) if 1<sh<max_dim else None for sh in t.shape]
+
+                # either scalar parameter, 1d with precondition_1d=False, or all dims are too big.
+                if len([i is not None for i in state['GG']]) == 0:
+                    state['GG'] = None
+
+                if state['GG'] is not None:
+                    update_soap_covariances_(y, GGs_=state['GG'], beta=shampoo_beta)
+                    state['Q'] = get_orthogonal_matrix(state['GG'])
+
+                state['step'] = 0
+                updates.append(tensors[i].sign())
+                continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
+                # I use sign instead as to not mess up with next modules. 1st Adam step is always sign anyway.
+
+            # Projecting gradients to the eigenbases of Shampoo's preconditioner
+            # i.e. projecting to the eigenbases of matrices in state['GG']
+            z_projected = None
+            if state['GG'] is not None:
+                if y_to_ema2: z_projected = project(y, state['Q'])
+                else: z_projected = project(t, state['Q'])
+
+            # exponential moving averages
+            # this part could be foreached but I will do that at some point its not a big difference compared to preconditioning
+            exp_avg: torch.Tensor = state["exp_avg"]
+            exp_avg_sq: torch.Tensor = state["exp_avg_sq"]
+
+            exp_avg.lerp_(t, 1-beta1)
+
+            if z_projected is None:
+                if y_to_ema2: exp_avg_sq.mul_(beta2).addcmul_(y, y, value=1-beta2)
+                else: exp_avg_sq.mul_(beta2).addcmul_(t, t, value=1-beta2)
+            else:
+                exp_avg_sq.mul_(beta2).addcmul_(z_projected, z_projected, value=1-beta2)
+
+            # project exponential moving averages if they are accumulated unprojected
+            exp_avg_projected = exp_avg
+            if z_projected is not None:
+                exp_avg_projected = project(exp_avg, state['Q'])
+
+            exp_avg_sq_projected = exp_avg_sq
+
+            denom = exp_avg_sq_projected.sqrt().add_(eps)
+            # print(f'{t_projected = }, {exp_avg = }, {exp_avg_projected = }, {exp_avg_sq = }, {exp_avg_sq_projected = }, {denom = }')
+
+            # Projecting back the preconditioned (by Adam) exponential moving average of gradients
+            # to the original space
+            update = exp_avg_projected / denom
+            if z_projected is not None:
+                update = project_back(update, state["Q"])
+
+            if settings['bias_correction']:
+                bias_correction1 = 1.0 - beta1 ** (state["step"]+1)
+                bias_correction2 = 1.0 - beta2 ** (state["step"]+1)
+                update *= ((bias_correction2 ** .5) / bias_correction1) * alpha
+            elif alpha is not None:
+                update *= alpha
+
+            if merge_small:
+                update = _unmerge_small_dims(update, state['flat_sizes'], state['sort_idxs'])
+
+            updates.append(update)
+            state["step"] += 1
+
+            # Update is done after the gradient step to avoid using current gradients in the projection.
+            if state['GG'] is not None:
+                update_soap_covariances_(y, state['GG'], shampoo_beta)
+                if state['step'] % settings['precond_freq'] == 0:
+                    state['Q'], state['exp_avg_sq'] = get_orthogonal_matrix_QR(exp_avg_sq, state['GG'], state['Q'])
+
+        return updates

torchzero/modules/experimental/gradmin.py

@@ -0,0 +1,85 @@
+import warnings
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Sequence
+from typing import Literal
+
+import torch
+
+from ...core import Module, Vars
+from ...utils import NumberList, TensorList
+from ...utils.derivatives import jacobian_wrt
+from ..grad_approximation import GradApproximator, GradTarget
+from ..smoothing.gaussian import Reformulation
+
+
+
+class GradMin(Reformulation):
+    """Reformulates the objective to minimize sum of gradient magnitudes via autograd.
+
+    Args:
+        loss_term (float, optional): adds loss value times this to sum of gradient magnitudes. Defaults to 1.
+        relative (bool, optional): whether to make loss_term relative to gradient magnitude. Defaults to False.
+        graft (bool, optional): whether to make loss term same as gradient magnitude. Defaults to False.
+        square (bool, optional): whether to use sum of squared gradient magnitudes, if False uses absolute values. Defaults to False.
+        mean (bool, optional): whether to use mean, if False uses sum. Defaults to True.
+        maximize_grad (bool, optional): whether to maximize gradient magnitudes instead of minimizing. Defaults to False.
+        create_graph (bool, optional): whether to create graph. Defaults to False.
+        modify_loss (bool, optional): whether to modify the loss value to make line searches minimize new objective. Defaults to True.
+    """
+    def __init__(
+        self,
+        loss_term: float | None = 0,
+        relative: Literal['loss_to_grad', 'grad_to_loss'] | None = None,
+        graft: Literal['loss_to_grad', 'grad_to_loss'] | None = None,
+        square=False,
+        mean=True,
+        maximize_grad=False,
+        create_graph=False,
+        modify_loss: bool = True,
+    ):
+        if (relative is not None) and (graft is not None): warnings.warn('both relative and graft loss are True, they will clash with each other')
+        defaults = dict(loss_term=loss_term, relative=relative, graft=graft, square=square, mean=mean, maximize_grad=maximize_grad, create_graph=create_graph, modify_loss=modify_loss)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def closure(self, backward, closure, params, vars):
+        settings = self.settings[params[0]]
+        loss_term = settings['loss_term']
+        relative = settings['relative']
+        graft = settings['graft']
+        square = settings['square']
+        maximize_grad = settings['maximize_grad']
+        create_graph = settings['create_graph']
+        modify_loss = settings['modify_loss']
+        mean = settings['mean']
+
+        with torch.enable_grad():
+            for p in params: p.grad = None
+            loss = closure(False)
+            grads = TensorList(torch.autograd.grad(loss, params, create_graph=True))
+
+            if square: grads = grads ** 2
+            else: grads = grads.abs()
+
+            if mean: f = grads.global_mean()
+            else: f = grads.global_sum()
+
+
+            if graft == 'grad_to_loss': f = f * (loss.detach()/f.detach()).detach()
+            if relative == 'grad_to_loss': f = f * loss
+
+            if loss_term is not None and loss_term != 0:
+                if relative == 'loss_to_grad': loss_term = loss_term * f
+                l = loss
+                if graft == 'loss_to_grad': l = loss * (f.detach()/loss.detach()).detach()
+                f = f + l*loss_term
+
+            if maximize_grad: f = -f
+            if modify_loss: loss = f
+
+            grad = None
+            if backward:
+                for p in params: p.grad = None
+                grad = TensorList(torch.autograd.grad(f, params, create_graph=create_graph))
+
+        return loss, grad

torchzero/modules/experimental/reduce_outward_lr.py

@@ -0,0 +1,35 @@
+import torch
+
+from ...core import Target, Transform
+from ...utils import TensorList
+
+class ReduceOutwardLR(Transform):
+    """
+    When update sign matches weight sign, the learning rate for that weight is multiplied by `mul`.
+
+    This means updates that move weights towards zero have higher learning rates.
+    """
+    def __init__(self, mul = 0.5, use_grad=False, invert=False, target: Target = 'update'):
+        defaults = dict(mul=mul, use_grad=use_grad, invert=invert)
+        super().__init__(defaults, uses_grad=use_grad, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        params = TensorList(params)
+        tensors = TensorList(tensors)
+
+        mul = self.get_settings('mul', params=params)
+        s = self.settings[params[0]]
+        use_grad = s['use_grad']
+        invert = s['invert']
+
+        if use_grad: cur = vars.get_grad()
+        else: cur = tensors
+
+        # mask of weights where sign matches with update sign (minus ascent sign), multiplied by `mul`.
+        if invert: mask = (params * cur) > 0
+        else: mask = (params * cur) < 0
+
+        tensors.masked_set_(mask, tensors*mul)
+
+        return tensors

torchzero/modules/experimental/spectral.py

@@ -0,0 +1,286 @@
+from abc import ABC, abstractmethod
+import math
+from collections import deque
+from typing import Literal, Any
+
+import torch
+from ...core import Chainable, TensorwisePreconditioner
+from ...utils.linalg.matrix_funcs import matrix_power_eigh
+from ...utils.linalg.svd import randomized_svd
+from ...utils.linalg.qr import qr_householder
+
+
+class _Solver:
+    @abstractmethod
+    def update(self, history: deque[torch.Tensor], damping: float | None) -> tuple[Any, Any]:
+        """returns stuff for apply"""
+    @abstractmethod
+    def apply(self, __g: torch.Tensor, __A:torch.Tensor, __B:torch.Tensor) -> torch.Tensor:
+        """apply preconditioning to tensor"""
+
+class _SVDSolver(_Solver):
+    def __init__(self, driver=None): self.driver=driver
+    def update(self, history, damping):
+        M_hist = torch.stack(tuple(history), dim=1)
+        device = None # driver is CUDA only
+        if self.driver is not None:
+            device = M_hist.device
+            M_hist = M_hist.cuda()
+
+        try:
+            U, S, _ = torch.linalg.svd(M_hist, full_matrices=False, driver=self.driver) # pylint:disable=not-callable
+
+            if self.driver is not None:
+                U = U.to(device); S = S.to(device)
+
+            if damping is not None and damping != 0: S.add_(damping)
+            return U, S
+
+        except torch.linalg.LinAlgError:
+            return None, None
+
+    def apply(self, g: torch.Tensor, U: torch.Tensor, S: torch.Tensor):
+        Utg = (U.T @ g).div_(S)
+        return U @ Utg
+
+class _SVDLowRankSolver(_Solver):
+    def __init__(self, q: int = 6, niter: int = 2): self.q, self.niter = q, niter
+    def update(self, history, damping):
+        M_hist = torch.stack(tuple(history), dim=1)
+        try:
+            U, S, _ = torch.svd_lowrank(M_hist, q=self.q, niter=self.niter)
+            if damping is not None and damping != 0: S.add_(damping)
+            return U, S
+        except torch.linalg.LinAlgError:
+            return None, None
+
+    def apply(self, g: torch.Tensor, U: torch.Tensor, S: torch.Tensor):
+        Utg = (U.T @ g).div_(S)
+        return U @ Utg
+
+class _RandomizedSVDSolver(_Solver):
+    def __init__(self, k: int = 3, driver: str | None = 'gesvda'):
+        self.driver = driver
+        self.k = k
+
+    def update(self, history, damping):
+        M_hist = torch.stack(tuple(history), dim=1)
+        device = None # driver is CUDA only
+        if self.driver is not None:
+            device = M_hist.device
+            M_hist = M_hist.cuda()
+
+        try:
+            U, S, _ = randomized_svd(M_hist, k=self.k, driver=self.driver)
+
+            if self.driver is not None:
+                U = U.to(device); S = S.to(device)
+
+            if damping is not None and damping != 0: S.add_(damping)
+            return U, S
+
+        except torch.linalg.LinAlgError:
+            return None, None
+
+    def apply(self, g: torch.Tensor, U: torch.Tensor, S: torch.Tensor):
+        Utg = (U.T @ g).div_(S)
+        return U @ Utg
+
+class _QRDiagonalSolver(_Solver):
+    def __init__(self, sqrt=True): self.sqrt = sqrt
+    def update(self, history, damping):
+        M_hist = torch.stack(tuple(history), dim=1)
+        try:
+            Q, R = torch.linalg.qr(M_hist, mode='reduced') # pylint:disable=not-callable
+            R_diag = R.diag().abs()
+            if damping is not None and damping != 0: R_diag.add_(damping)
+            if self.sqrt: R_diag.sqrt_()
+            return Q, R_diag
+        except torch.linalg.LinAlgError:
+            return None, None
+
+    def apply(self, g: torch.Tensor, Q: torch.Tensor, R_diag: torch.Tensor):
+        Qtg = (Q.T @ g).div_(R_diag)
+        return Q @ Qtg
+
+class _QRSolver(_Solver):
+    def __init__(self, sqrt=True): self.sqrt = sqrt
+    def update(self, history, damping):
+        M_hist = torch.stack(tuple(history), dim=1)
+        try:
+            # Q: d x k, R: k x k
+            Q, R = torch.linalg.qr(M_hist, mode='reduced') # pylint:disable=not-callable
+            A = R @ R.T
+            if damping is not None and damping != 0: A.diagonal(dim1=-2, dim2=-1).add_(damping)
+            if self.sqrt: A = matrix_power_eigh(A, 0.5)
+            return Q, A
+        except (torch.linalg.LinAlgError):
+            return None,None
+
+    def apply(self, g: torch.Tensor, Q: torch.Tensor, A: torch.Tensor) -> torch.Tensor:
+        g_proj = Q.T @ g
+        y, _ = torch.linalg.solve_ex(A, g_proj) # pylint:disable=not-callable
+        return Q @ y
+
+class _QRHouseholderSolver(_Solver):
+    def __init__(self, sqrt=True): self.sqrt = sqrt
+    def update(self, history, damping):
+        M_hist = torch.stack(tuple(history), dim=1)
+        try:
+            # Q: d x k, R: k x k
+            Q, R = qr_householder(M_hist, mode='reduced') # pylint:disable=not-callable
+            A = R @ R.T
+            if damping is not None and damping != 0: A.diagonal(dim1=-2, dim2=-1).add_(damping)
+            if self.sqrt: A = matrix_power_eigh(A, 0.5)
+            return Q, A
+        except (torch.linalg.LinAlgError):
+            return None,None
+
+    def apply(self, g: torch.Tensor, Q: torch.Tensor, A: torch.Tensor) -> torch.Tensor:
+        g_proj = Q.T @ g
+        y, _ = torch.linalg.solve_ex(A, g_proj) # pylint:disable=not-callable
+        return Q @ y
+
+
+class _EighSolver(_Solver):
+    def __init__(self, sqrt=True):
+        self.sqrt = sqrt
+
+    def update(self, history, damping):
+        M_hist = torch.stack(tuple(history), dim=1)
+        grams = M_hist @ M_hist.T # (d, d)
+        if damping is not None and damping != 0: grams.diagonal(dim1=-2, dim2=-1).add_(damping)
+        try:
+            L, Q = torch.linalg.eigh(grams) # L: (d,), Q: (d, d) # pylint:disable=not-callable
+            L = L.abs().clamp_(min=1e-12)
+            if self.sqrt: L = L.sqrt()
+            return Q, L
+        except torch.linalg.LinAlgError:
+            return None, None
+
+    def apply(self, g: torch.Tensor, Q: torch.Tensor, L: torch.Tensor) -> torch.Tensor:
+        Qtg = (Q.T @ g).div_(L)
+        return Q @ Qtg
+
+
+SOLVERS = {
+    "svd": _SVDSolver(), # fallbacks on "gesvd" which basically takes ages or just hangs completely
+    "svd_gesvdj": _SVDSolver("gesvdj"), # no fallback on slow "gesvd"
+    "svd_gesvda": _SVDSolver("gesvda"), # approximate method for wide matrices, sometimes better sometimes worse but faster
+    "svd_lowrank": _SVDLowRankSolver(), # maybe need to tune parameters for this, with current ones its slower and worse
+    "randomized_svd2": _RandomizedSVDSolver(2),
+    "randomized_svd3": _RandomizedSVDSolver(3),
+    "randomized_svd4": _RandomizedSVDSolver(4),
+    "randomized_svd5": _RandomizedSVDSolver(5),
+    "eigh": _EighSolver(), # this is O(n**2) storage, but is this more accurate?
+    "qr": _QRSolver(),
+    "qr_householder": _QRHouseholderSolver(), # this is slower... but maybe it won't freeze? I think svd_gesvda is better
+    "qrdiag": _QRDiagonalSolver(),
+}
+
+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
+        else: state_[key].lerp_(value, 1-beta)
+
+class SpectralPreconditioner(TensorwisePreconditioner):
+    """Whitening preconditioner via SVD on history of past gradients or gradient differences scaled by parameter differences.
+
+    Args:
+        history_size (int, optional): number of past gradients to store for preconditioning. Defaults to 10.
+        update_freq (int, optional): how often to re-compute the preconditioner. Defaults to 1.
+        damping (float, optional): damping term, makes it closer to GD. Defaults to 1e-7.
+        order (int, optional):
+            whitening order, 1 approximates FIM (maybe), 2 - hessian (maybe), 3+ - god knows what.
+        solver (str, optional): what to use for whitening. Defaults to 'svd'.
+        U_beta (float | None, optional): beta for U (probably a bad idea). Defaults to None.
+        S_beta (float | None, optional): beta for S (probably a bad idea). Defaults to None.
+        interval (int, optional): How often to update history. Defaults to 1 (every step).
+        concat_params (bool, optional):
+            whether to apply preconditioning to each tensor (False, default) or to all tensors concatenated into a vector (True). Latter will be slower but captures interactions between layers. Defaults to True.
+        scale_first (bool, optional): makes first step small, usually not needed. Defaults to False.
+        inner (Chainable | None, optional): Inner modules applied after updating preconditioner and before applying it. Defaults to None.
+    """
+    def __init__(
+        self,
+        history_size: int = 10,
+        update_freq: int = 1,
+        damping: float = 1e-12,
+        order: int = 1,
+        solver: Literal['svd', 'svd_gesvdj', 'svd_gesvda', 'svd_lowrank', 'eigh', 'qr', 'qrdiag', 'qr_householder'] | _Solver | str = 'svd_gesvda',
+        A_beta: float | None = None,
+        B_beta: float | None = None,
+        interval: int = 1,
+        concat_params: bool = False,
+        scale_first: bool = False,
+        inner: Chainable | None = None,
+    ):
+        if isinstance(solver, str): solver = SOLVERS[solver]
+        # 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, order=order, A_beta=A_beta, B_beta=B_beta, solver=solver)
+        super().__init__(defaults, uses_grad=False, concat_params=concat_params, scale_first=scale_first, inner=inner, update_freq=interval)
+
+    @torch.no_grad
+    def update_tensor(self, tensor, param, grad, state, settings):
+        order = settings['order']
+        history_size = settings['history_size']
+        update_freq = settings['update_freq']
+        damping = settings['damping']
+        A_beta = settings['A_beta']
+        B_beta = settings['B_beta']
+        solver: _Solver = settings['solver']
+
+        if 'history' not in state: state['history'] = deque(maxlen=history_size)
+        history = state['history']
+
+        if order == 1: history.append(tensor.clone().view(-1))
+        else:
+
+            # if order=2, history is of gradient differences, order 3 is differences between differences, etc
+            # normalized 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_k = cur_p - state[f'prev_p_{i}']
+                y_k = 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
+
+                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:
+            A, B = solver.update(history, damping=damping)
+            maybe_lerp_(state, A_beta, 'A', A)
+            maybe_lerp_(state, B_beta, 'B', B)
+
+        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, state, settings):
+        history_size = settings['history_size']
+        solver: _Solver = settings['solver']
+
+        A = state.get('A', None)
+        if A is None:
+            # make a conservative step to avoid issues due to different GD scaling
+            return tensor.div_(max(1, tensor.abs().sum())) # pyright:ignore[reportArgumentType]
+
+        B = state['B']
+        update = solver.apply(tensor.view(-1), A, B).view_as(tensor)
+
+        n = len(state['history'])
+        if n != history_size: update.mul_(n/history_size)
+        return update
+