torchzero 0.4.1__py3-none-any.whl → 0.4.2__py3-none-any.whl

torchzero/modules/adaptive/ggt.py

@@ -7,7 +7,7 @@ from ...core import Chainable, TensorTransform
 from ...linalg import torch_linalg, regularize_eigh
 from .lre_optimizers import LREOptimizerBase
 
-def ggt_update(history: deque[torch.Tensor] | torch.Tensor, damping, rdamping, truncate, eig_tol):
+def ggt_update(history: deque[torch.Tensor] | torch.Tensor, damping, rdamping, truncate, eig_tol, matrix_power=-1/2):
     """returns U ``(ndim, rank)``, L ``(rank, )``"""
     if isinstance(history, torch.Tensor):
         M = history
@@ -27,7 +27,7 @@ def ggt_update(history: deque[torch.Tensor] | torch.Tensor, damping, rdamping, t
         if L is None or Q is None: # this means there are no finite eigenvalues
             return None, None
 
-        U = (M @ Q) * L.rsqrt()
+        U = (M @ Q) * L.pow(matrix_power)
 
         # this damping is added after computing U, this is why I didn't use one in linalg.regularize_eig
         # that's because we damp singular values this way
@@ -44,14 +44,13 @@ class GGT(TensorTransform):
     """
     GGT method from https://arxiv.org/pdf/1806.02958
 
-    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.
+    The update rule is to stack recent gradients into M and
+    compute eigendecomposition of M M^T via eigendecomposition of M^T M.
 
     This is equivalent to full-matrix Adagrad on recent gradients.
 
     Args:
         history_size (int, optional): number of past gradients to store. Defaults to 10.
-        beta (float, optional): beta for momentum maintained in whitened space. Defaults to 0.0.
         update_freq (int, optional): frequency of updating the preconditioner (U and S). Defaults to 1.
         eig_tol (float, optional): removes eigenvalues this much smaller than largest eigenvalue. Defaults to 1e-7.
         truncate (int, optional): number of larges eigenvalues to keep. None to disable. Defaults to None.
@@ -105,7 +104,8 @@ class GGT(TensorTransform):
         truncate: int | None = None,
         damping: float = 1e-4,
         rdamping: float = 0,
-        eigenbasis_optimizer: LREOptimizerBase | None = None,
+        matrix_power: float = -1/2,
+        basis_optimizer: LREOptimizerBase | None = None,
         concat_params: bool = True,
 
         inner: Chainable | None = None,
@@ -114,6 +114,7 @@ class GGT(TensorTransform):
         del defaults['self'], defaults['inner'], defaults['concat_params']
 
         super().__init__(defaults, concat_params=concat_params, inner=inner)
+        self.add_projected_keys("grad", "history")
 
     @torch.no_grad
     def single_tensor_update(self, tensor, param, grad, loss, state, setting):
@@ -141,14 +142,15 @@ class GGT(TensorTransform):
                 rdamping=setting["rdamping"],
                 truncate=setting["truncate"],
                 eig_tol=setting["eig_tol"],
+                matrix_power=setting["matrix_power"],
             )
 
-            # reproject eigenbasis optimizer
-            eigenbasis_optimizer: LREOptimizerBase | None = setting["eigenbasis_optimizer"]
-            if eigenbasis_optimizer is not None:
+            # reproject basis optimizer
+            basis_optimizer: LREOptimizerBase | None = setting["basis_optimizer"]
+            if basis_optimizer is not None:
                 if (L is not None) and (U is not None) and (L_new is not None) and (U_new is not None):
-                    eigenbasis_state = state["eigenbasis_state"]
-                    eigenbasis_optimizer.reproject(L_old=L, Q_old=U, L_new=L_new, Q_new=U_new, state=eigenbasis_state)
+                    basis_state = state["basis_state"]
+                    basis_optimizer.reproject(L_old=L, Q_old=U, L_new=L_new, Q_new=U_new, state=basis_state)
 
 
             # store new factors
@@ -169,18 +171,18 @@ class GGT(TensorTransform):
 
         L = state['L']
 
-        # step with eigenbasis optimizer
-        eigenbasis_optimizer: LREOptimizerBase | None = setting["eigenbasis_optimizer"]
-        if eigenbasis_optimizer is not None:
+        # step with basis optimizer
+        basis_optimizer: LREOptimizerBase | None = setting["basis_optimizer"]
+        if basis_optimizer is not None:
 
-            if "eigenbasis_state" not in state: state["eigenbasis_state"] = {}
-            eigenbasis_state = state["eigenbasis_state"]
+            if "basis_state" not in state: state["basis_state"] = {}
+            basis_state = state["basis_state"]
 
-            update = eigenbasis_optimizer.step(g, L=L, Q=U, state=eigenbasis_state)
+            update = basis_optimizer.step(g, L=L, Q=U, state=basis_state)
             return update.view_as(tensor)
 
         # or just whiten
         z = U.T @ g
-        update = (U * L.rsqrt()) @ z
+        update = (U * L.pow(setting["matrix_power"])) @ z
         return update.view_as(tensor)
 

torchzero/modules/adaptive/lion.py

@@ -23,9 +23,11 @@ class Lion(TensorTransform):
         defaults = dict(beta1=beta1, beta2=beta2)
         super().__init__(defaults)
 
+        self.add_projected_keys("grad", "exp_avg")
+
     @torch.no_grad
     def multi_tensor_apply(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)
+        exp_avg = unpack_states(states, tensors, 'exp_avg', cls=TensorList)
         return lion_(TensorList(tensors), exp_avg, beta1, beta2)
 

torchzero/modules/adaptive/mars.py

@@ -6,13 +6,13 @@ from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 
 def mars_correction_(
     tensors_: TensorList,
-    prev_: TensorList,
+    g_prev_: TensorList,
     beta: float | NumberList,
     scaling: float | NumberList,
     max_norm: float | NumberList |  None,
 ):
-    dg = (tensors_ - prev_).mul_(scaling * beta / (1-beta))
-    prev_.copy_(tensors_)
+    dg = (tensors_ - g_prev_).mul_(scaling * beta / (1-beta))
+    g_prev_.copy_(tensors_)
 
     c = tensors_.add_(dg)
     if max_norm is not None:
@@ -63,16 +63,17 @@ class MARSCorrection(TensorTransform):
     ):
         defaults = dict(beta=beta, scaling=scaling, max_norm=max_norm)
         super().__init__(defaults)
+        self.add_projected_keys("grad", "g_prev")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
-        prev = unpack_states(states, tensors, 'prev', init=tensors, cls=TensorList)
+        g_prev = unpack_states(states, tensors, 'g_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,
+            g_prev_=g_prev,
             beta=beta,
             scaling=scaling,
             max_norm=max_norm,

torchzero/modules/adaptive/msam.py

@@ -121,6 +121,8 @@ class MSAMMomentum(TensorTransform):
         defaults = dict(lr = lr, momentum=momentum, rho=rho, nesterov=nesterov, lerp=lerp, weight_decay=weight_decay)
         super().__init__(defaults, uses_grad=False)
 
+        self.add_projected_keys("grad", "velocity")
+
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         velocity = unpack_states(states, tensors, 'velocity', cls=TensorList)
@@ -180,6 +182,7 @@ class MSAM(Transform):
         super().__init__(defaults)
 
         self.set_child('modules', modules)
+        self.add_projected_keys("grad", "velocity")
 
 
     @torch.no_grad

torchzero/modules/adaptive/rmsprop.py

@@ -38,6 +38,8 @@ class RMSprop(TensorTransform):
         super().__init__(defaults, inner=inner)
 
         self.set_child('exp_avg_sq', exp_avg_sq_tfm)
+        self.add_projected_keys("grad", "exp_avg")
+        self.add_projected_keys("grad_sq", "exp_avg_sq", "exp_avg_sq_max")
 
     @torch.no_grad
     def single_tensor_initialize(self, tensor, param, grad, loss, state, setting):

torchzero/modules/adaptive/rprop.py

@@ -128,15 +128,15 @@ def rprop_(
 
 class Rprop(TensorTransform):
     """
-    Resilient propagation. The update magnitude gets multiplied by `nplus` if gradient didn't change the sign,
-    or `nminus` if it did. Then the update is applied with the sign of the current gradient.
+    Resilient propagation. The update magnitude gets multiplied by ``nplus`` if gradient didn't change the sign,
+    or ``nminus`` if it did. Then the update is applied with the sign of the current gradient.
 
     Additionally, if gradient changes sign, the update for that weight is reverted.
     Next step, magnitude for that weight won't change.
 
     Compared to pytorch this also implements backtracking update when sign changes.
 
-    This implementation is identical to :code:`torch.optim.Rprop` if :code:`backtrack` is set to False.
+    This implementation is identical to ``torch.optim.Rprop`` if ``backtrack`` is set to False.
 
     Args:
         nplus (float): multiplicative increase factor for when ascent didn't change sign (default: 1.2).
@@ -164,6 +164,8 @@ class Rprop(TensorTransform):
         defaults = dict(nplus = nplus, nminus = nminus, alpha = alpha, lb = lb, ub = ub, backtrack=backtrack)
         super().__init__(defaults, uses_grad=False)
 
+        self.add_projected_keys("grad", "prev")
+
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         step = self.global_state.get('step', 0)
@@ -196,14 +198,14 @@ class Rprop(TensorTransform):
 
 class ScaleLRBySignChange(TensorTransform):
     """
-    learning rate gets multiplied by `nplus` if ascent/gradient didn't change the sign,
-    or `nminus` if it did.
+    learning rate gets multiplied by ``nplus`` if ascent/gradient didn't change the sign,
+    or ``nminus`` if it did.
 
     This is part of RProp update rule.
 
     Args:
-        nplus (float): learning rate gets multiplied by `nplus` if ascent/gradient didn't change the sign
-        nminus (float): learning rate gets multiplied by `nminus` if ascent/gradient changed the sign
+        nplus (float): learning rate gets multiplied by ``nplus`` if ascent/gradient didn't change the sign
+        nminus (float): learning rate gets multiplied by ``nminus`` if ascent/gradient changed the sign
         lb (float): lower bound for lr.
         ub (float): upper bound for lr.
         alpha (float): initial learning rate.

torchzero/modules/adaptive/shampoo.py

@@ -207,6 +207,9 @@ class Shampoo(TensorTransform):
         if setting["merge_small"]:
             tensor, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(tensor, setting["max_dim"])
 
+            if "inner" not in self.children:
+                state["merged"] = tensor
+
         if 'diagonal_accumulator' in state:
             update_diagonal_(tensor, state['diagonal_accumulator'], beta=setting["beta"])
         else:
@@ -227,10 +230,15 @@ class Shampoo(TensorTransform):
 
         state["step"] += 1
 
+
     @torch.no_grad
     def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
+
         if setting["merge_small"]:
-            tensor, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(tensor, setting["max_dim"])
+            if "inner" not in self.children:
+                tensor = state.pop("merged")
+            else:
+                tensor, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(tensor, setting["max_dim"])
 
         if 'diagonal_accumulator' in state:
             dir = apply_diagonal_(tensor, state['diagonal_accumulator'], eps=setting["adagrad_eps"])

torchzero/modules/adaptive/soap.py

@@ -1,12 +1,13 @@
-from operator import itemgetter
 import warnings
+from operator import itemgetter
 
 import torch
 
-from ...core import TensorTransform, Chainable
-from ...utils import unpack_dicts, unpack_states, TensorList, NumberList
-from ...modules.adaptive.shampoo import _merge_small_dims, _unmerge_small_dims
+from ...core import Chainable, TensorTransform
 from ...linalg import torch_linalg
+from ...modules.adaptive.shampoo import _merge_small_dims, _unmerge_small_dims
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+
 
 @torch.no_grad
 def update_soap_covariances_(
@@ -221,25 +222,38 @@ class SOAP(TensorTransform):
             return TensorList(tensors).clamp(-0.1, 0.1)
             # return TensorList(tensors).zero_()
 
-
         fs = settings[0]
-        merged = []
+        merged_updates = [] # for when exp_avg is maintained unprojected
+        merged_grads = [] # this doesn't go into preconditioner
         projected = []
-        # ---------------------------------- project --------------------------------- #
 
-        for tensor, state, setting in zip(tensors, states, settings):
+        # -------------------------------- inner step -------------------------------- #
+        updates = tensors
+        has_inner = "inner" in self.children
+        if has_inner:
+            updates = self.inner_step_tensors("inner", updates, clone=True,
+                                              params=params, grads=grads, loss=loss)
+
+        # ---------------------------------- project --------------------------------- #
+        for grad, update, state, setting in zip(tensors, updates, states, settings):
             if setting["merge_small"]:
-                tensor, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(tensor, setting["max_dim"])
+                update, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(update, setting["max_dim"])
+                if has_inner: # grad is a different tensor, merge it too
+                    grad, _, _ = _merge_small_dims(grad, setting["max_dim"])
+                else: # in this case update is still just grad
+                    grad = update
 
-            merged.append(tensor)
+            merged_updates.append(update)
+            merged_grads.append(grad)
 
             if state['GG'] is not None:
-                tensor = project(tensor, state['Q'])
+                update = project(update, state['Q'])
+
+            projected.append(update)
 
-            projected.append(tensor)
 
         # ------------------------ run adam in projected space ----------------------- #
-        exp_avg_proj, exp_avg_sq_proj = unpack_states(states, tensors, "exp_avg_proj", "exp_avg_sq_proj", must_exist=True, cls=TensorList)
+        exp_avg_proj, exp_avg_sq_proj = unpack_states(states, projected, "exp_avg_proj", "exp_avg_sq_proj", must_exist=True, cls=TensorList)
         alpha, beta1, beta2, eps = unpack_dicts(settings, "alpha", "beta1", "beta2", "eps", cls=NumberList)
 
         # lerp exp_avg in projected space
@@ -249,15 +263,17 @@ class SOAP(TensorTransform):
         # or lerp in original space and project
         else:
             exp_avg = exp_avg_proj
-            exp_avg.lerp_(merged, weight=1-beta1)
+            exp_avg.lerp_(merged_updates, weight=1-beta1)
             exp_avg_proj = []
             for t, state, setting in zip(exp_avg, states, settings):
                 if state['GG'] is not None:
                     t = project(t, state["Q"])
                 exp_avg_proj.append(t)
 
+        # lerp exp_avg_sq
         exp_avg_sq_proj.mul_(beta2).addcmul_(projected, projected, value=1-beta2)
 
+        # adam direction
         denom = exp_avg_sq_proj.sqrt().add_(eps)
         dirs_proj = exp_avg_proj / denom
 
@@ -272,27 +288,14 @@ class SOAP(TensorTransform):
 
             dirs.append(dir)
 
-
-        # -------------------------------- inner step -------------------------------- #
-        if "inner" in self.children:
-            tensors = self.inner_step_tensors("inner", tensors, clone=False,
-                                              params=params, grads=grads,loss=loss)
-
-            # we now have to re-merge small dims on updated tensors
-            merged = []
-            for tensor, state, setting in zip(tensors, states, settings):
-                if setting["merge_small"]:
-                    tensor, _, _ = _merge_small_dims(tensor, setting["max_dim"])
-                    merged.append(tensor)
-
         # -------------------------- update preconditioners -------------------------- #
         # Update is done after the gradient step to avoid using current gradients in the projection.
 
-        for tensor, state, setting in zip(merged, states, settings):
+        for grad, state, setting in zip(merged_grads, states, settings):
             if state['GG'] is not None:
 
                 # lerp covariances
-                update_soap_covariances_(tensor, state['GG'], beta=setting["shampoo_beta"])
+                update_soap_covariances_(grad, state['GG'], beta=setting["shampoo_beta"])
 
                 # (state['step'] - 1) since we start updating on 2nd step
                 if (state['step'] - 1) % setting['precond_freq'] == 0:

torchzero/modules/basis/__init__.py

@@ -0,0 +1,2 @@
+from .soap_basis import SOAPBasis
+from .ggt_basis import GGTBasis

torchzero/modules/basis/ggt_basis.py

@@ -0,0 +1,199 @@
+from collections import deque
+
+import torch
+
+from ...core import Chainable, TensorTransform
+from ...utils import set_storage_
+from ..adaptive.ggt import ggt_update
+
+
+def _cubic_reproject(C: torch.Tensor, cu: torch.Tensor, approx:bool):
+    if approx: return C.pow(3) @ cu
+
+    n = cu.numel()
+    T = torch.zeros([n,n,n], device=cu.device, dtype=cu.dtype)
+    T[range(n),range(n),range(n)] = cu
+    T = torch.einsum('ai,bj,ck,ijk->abc', C, C, C, T)
+    n2 = T.size(0)
+    return T[range(n2), range(n2), range(n2)]
+
+class GGTBasis(TensorTransform):
+    """
+    Run another optimizer in GGT eigenbasis. The eigenbasis is ``rank``-sized, so it is possible to run expensive
+    methods such as Full-matrix Adagrad/Adam.
+
+    The update rule is to stack recent gradients into M and
+    compute eigendecomposition of M M^T via eigendecomposition of M^T M.
+
+    This is equivalent to full-matrix Adagrad on recent gradients.
+
+    Note:
+        the buffers of the ``basis_opt`` are re-projected whenever basis changes. The reprojection logic is not implemented on all modules. Some supported modules are:
+
+        ``Adagrad``, ``FullMatrixAdagrad``, ``Adam``, ``Adan``, ``Lion``, ``MARSCorrection``, ``MSAMMomentum``, ``RMSprop``, ``GGT``, ``EMA``, ``HeavyBall``, ``NAG``, ``ClipNormByEMA``, ``ClipValueByEMA``, ``NormalizeByEMA``, ``ClipValueGrowth``, ``CoordinateMomentum``, ``CubicAdam``.
+
+        Additionally most modules with no internal buffers are supported, e.g. ``Cautious``, ``Sign``, ``ClipNorm``, ``Orthogonalize``, etc. However modules that use weight values, such as ``WeighDecay`` can't be supported, as weights can't be projected.
+
+        Also, if you say use ``EMA`` on output of ``Pow(2)``, the exponential average will be reprojected as gradient and not as squared gradients. Use modules like ``EMASquared``, ``SqrtEMASquared`` to get correct reprojections.
+
+
+    Args:
+        basis_opt (Chainable): module or modules to run in GGT eigenbasis.
+        history_size (int, optional): number of past gradients to store, and rank of preconditioner. Defaults to 10.
+        update_freq (int, optional): frequency of updating the preconditioner (U and S). Defaults to 1.
+        eig_tol (float, optional): removes eigenvalues this much smaller than largest eigenvalue. Defaults to 1e-7.
+        truncate (int, optional): number of larges eigenvalues to keep. None to disable. Defaults to None.
+        damping (float, optional): damping value. Defaults to 1e-4.
+        rdamping (float, optional): value of damping relative to largest eigenvalue. Defaults to 0.
+        concat_params (bool, optional): if True, treats all parameters as a single vector. Defaults to True.
+        inner (Chainable | None, optional):
+            output of this module is projected and ``basis_opt`` will run on it, but preconditioners are updated
+            from original gradients.
+
+    ## Examples:
+
+    Examples:
+    Adam in GGT eigenbasis:
+    ```python
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.GGTBasis(tz.m.Adam(beta2=0.99)),
+        tz.m.LR(1e-3)
+    )
+    ```
+
+    Full-matrix Adam in GGT eigenbasis. We can define full-matrix Adam through ``FullMatrixAdagrad``.
+    ```python
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.GGTBasis(
+            [tz.m.FullMatrixAdagrad(beta=0.99, inner=tz.m.EMA(0.9, debias=True))]
+        ),
+        tz.m.LR(1e-3)
+    )
+    ```
+
+    LaProp in GGT eigenbasis:
+    ```python
+
+    # we define LaProp through other modules, moved it out for brevity
+    laprop = (
+        tz.m.RMSprop(0.95),
+        tz.m.Debias(beta1=None, beta2=0.95),
+        tz.m.EMA(0.95),
+        tz.m.Debias(beta1=0.95, beta2=None),
+    )
+
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.GGTBasis(laprop),
+        tz.m.LR(1e-3)
+    )
+    ```
+
+    Reference:
+        Agarwal N. et al. Efficient full-matrix adaptive regularization //International Conference on Machine Learning. – PMLR, 2019. – С. 102-110.
+    """
+
+    def __init__(
+        self,
+        basis_opt: Chainable,
+        history_size: int = 100,
+        update_freq: int = 1,
+        eig_tol: float = 1e-7,
+        truncate: int | None = None,
+        damping: float = 1e-4,
+        rdamping: float = 0,
+        matrix_power: float = -1/2,
+        approx_sq_reproject:bool = False,
+        approx_cu_reproject:bool = False,
+
+        inner: Chainable | None = None,
+    ):
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner']
+
+        super().__init__(defaults, concat_params=True, inner=inner)
+        self.set_child("basis_opt", basis_opt)
+
+    @torch.no_grad
+    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
+        history_size = setting['history_size']
+        update_freq = setting['update_freq']
+
+        if 'history' not in state: state['history'] = deque(maxlen=history_size)
+        history = state['history']
+
+        t = tensor.clone().view(-1)
+        history.append(t)
+
+        step = state.get('step', 0)
+        state['step'] = step + 1
+
+        if step % update_freq == 0 :
+
+            # compute new factors
+            L = state.get("L", None)
+            U = state.get("U", None)
+
+            L_new, U_new = ggt_update(
+                history,
+                damping=setting["damping"],
+                rdamping=setting["rdamping"],
+                truncate=setting["truncate"],
+                eig_tol=setting["eig_tol"],
+                matrix_power=setting["matrix_power"],
+            )
+
+            if (L is not None) and (U is not None) and (L_new is not None) and (U_new is not None):
+                # reproject basis optimizer
+                # this happens after first step, so basis opt is initialized by then
+                # note that because we concatenate parameters, each buffer will a single rank-length vector
+                C = U_new.T @ U # change of basis matrix
+
+                # reproject gradient-like buffers
+                for (buff,) in self.get_child_projected_buffers("basis_opt", "grad"):
+                    set_storage_(buff, C @ buff)
+
+                # reproject covariance diagonal-like buffers
+                for (buff,) in self.get_child_projected_buffers("basis_opt", "grad_sq"):
+                    if setting["approx_sq_reproject"]: set_storage_(buff, C.pow(2) @ buff)
+                    else: set_storage_(buff, (C @ buff.diag_embed() @ C.T).diagonal())
+
+                # reproject third order diagonal-like buffers
+                for (buff,) in self.get_child_projected_buffers("basis_opt", "grad_cu"):
+                    buff_r = _cubic_reproject(C, buff, setting["approx_cu_reproject"])
+                    set_storage_(buff, buff_r)
+
+                # reproject covariance-like buffers
+                for (buff,) in self.get_child_projected_buffers("basis_opt", "covariance"):
+                    set_storage_(buff, C @ buff @ C.T)
+
+            # store new factors
+            if L_new is not None: state["L"] = L_new
+            if U_new is not None: state["U"] = U_new
+
+
+    @torch.no_grad
+    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
+        g = tensor.view(-1)
+        U = state.get('U', None)
+
+        if U is None:
+            # fallback to element-wise preconditioning
+            history = torch.stack(tuple(state["history"]), 0)
+            g /= history.square().mean(0).sqrt().add(1e-8)
+            return g.view_as(tensor)
+
+        # project
+        g_proj = U.T @ g
+
+        # step
+        dir_proj = self.inner_step_tensors("basis_opt", tensors=[g_proj], clone=False, grads=[g_proj])[0]
+
+        # unproject
+        update = U @ dir_proj
+
+        # update = (U * L.pow(setting["matrix_power"])) @ z
+        return update.view_as(tensor)
+