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

torchzero/modules/experimental/cubic_adam.py

@@ -0,0 +1,160 @@
+from typing import Any, Literal
+
+import torch
+
+from ...core import TensorTransform
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+from ..adaptive.lre_optimizers import LREOptimizerBase, _squared_reproject
+
+
+def signed_cbrt(x: TensorList | Any) -> Any:
+    return x.sign() * x.abs().pow(1/3)
+
+def _clip_min_magnitude(x: torch.Tensor, eps: float):
+    return x.sign() * x.abs().clamp(min=eps)
+
+_cubic_adam_mode = Literal["signed_cbrt", "unsigned_cbrt", "halve"]
+
+def _cubic_minimize(A: torch.Tensor | Any, B: torch.Tensor | Any, C: torch.Tensor | Any, eps):
+    """minimizes (A/3)x^3 + (A/2)x^2 + Cx"""
+    discriminant = B**2 - 4 * A * C
+
+    denom = _clip_min_magnitude(2 * A, eps)
+    root = discriminant.clamp(min=0).sqrt_()
+
+    x0 = (-B + root) / denom
+    x1 = (-B - root) / denom
+
+    f0 = (A/3)*x0**3 + (B/2)*x0**2 + C*x0
+    f1 = (A/3)*x1**3 + (B/2)*x1**2 + C*x1
+
+    x_star = x0.where(f0 < f1, x1)
+
+    adam = -C / (B + eps)
+    return adam.where(discriminant < 0, x_star)
+
+def cubic_adam_(
+    tensors: TensorList,
+    exp_avg_: TensorList,
+    exp_avg_sq_: TensorList,
+    exp_avg_cu_: TensorList,
+    alpha: float | NumberList,
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    beta3: float | NumberList,
+    eps: float | NumberList,
+    debiased: bool,
+    step: int,
+
+    mode: _cubic_adam_mode = 'signed_cbrt'
+):
+    exp_avg_.lerp_(tensors, 1-beta1)
+    exp_avg_sq_.lerp_(tensors**2, 1-beta2)
+    exp_avg_cu_.lerp_(tensors**3, 1-beta3)
+
+    if debiased:
+        m1 = exp_avg_ / (1 - beta1 ** step)
+        m2 = exp_avg_sq_ / (1 - beta2 ** step)
+        m3 = exp_avg_cu_ / (1 - beta3 ** step)
+    else:
+        m1, m2, m3 = exp_avg_, exp_avg_sq_, exp_avg_cu_
+
+    # adam minimizes ax^2 + bx
+    # we are going to minimize ax^3 + bx^2 + cx
+
+    if mode == "signed_cbrt": A = signed_cbrt(m3)
+    elif mode == "unsigned_cbrt": A = m3.abs().pow(1/3)
+    elif mode == 'halve': A = 0.5 * m3
+    else: raise ValueError(mode)
+
+    B = m2.sqrt()
+    C = m1
+    x_star = _cubic_minimize(A, B, C, eps)
+    return x_star.mul_(-alpha)
+
+class CubicAdam(TensorTransform):
+    """Adam which has 3rd momentum and minimizes a cubic polynomial."""
+    def __init__(
+        self,
+        beta1: float = 0.9,
+        beta2: float = 0.99,
+        beta3: float = 0.99,
+        eps: float = 1e-8,
+        debiased:bool=True,
+        alpha: float = 1.,
+
+        mode: _cubic_adam_mode = 'signed_cbrt'
+    ):
+        defaults=dict(beta1=beta1,beta2=beta2,beta3=beta3,eps=eps,debiased=debiased,alpha=alpha,mode=mode)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        beta1,beta2,beta3,eps,alpha=unpack_dicts(settings, 'beta1','beta2','beta3','eps','alpha', cls=NumberList)
+        exp_avg, exp_avg_sq, exp_avg_cu = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', 'exp_avg_cu', cls=TensorList)
+
+        return cubic_adam_(
+            tensors=TensorList(tensors),
+            exp_avg_=exp_avg,
+            exp_avg_sq_=exp_avg_sq,
+            exp_avg_cu_=exp_avg_cu,
+            alpha=alpha,
+            beta1=beta1,
+            beta2=beta2,
+            beta3=beta3,
+            eps=eps,
+            debiased=settings[0]['debiased'],
+            step=step,
+
+            mode=settings[0]["mode"]
+        )
+
+class SubspaceCubicAdam(LREOptimizerBase):
+    """Runs cubic Adam in low rank eigenbasis."""
+    def __init__(self, beta1=0.9, beta2=0.95, beta3=0.95, eps=1e-8, mode: _cubic_adam_mode = 'signed_cbrt', cautious:bool=False, exact_reproject:bool=True):
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.beta3 = beta3
+        self.eps = eps
+        self.cautious = cautious
+        self.mode: _cubic_adam_mode = mode
+        self.exact_reproject = exact_reproject
+
+    def step(self, g, L, Q, state):
+        g = Q.T @ g
+
+        if "exp_avg" not in state:
+            state["exp_avg"] = torch.zeros_like(g)
+            state["exp_avg_sq"] = torch.zeros_like(g)
+            state["exp_avg_cu"] = torch.zeros_like(g)
+            state["current_step"] = 1
+
+        dir = cubic_adam_(
+            tensors = TensorList([g]),
+            exp_avg_ = TensorList([state["exp_avg"]]),
+            exp_avg_sq_ = TensorList([state["exp_avg_sq"]]),
+            exp_avg_cu_ = TensorList([state["exp_avg_cu"]]),
+            alpha = 1,
+            beta1 = self.beta1,
+            beta2 = self.beta2,
+            beta3 = self.beta3,
+            eps = self.eps,
+            debiased = True,
+            step = state["current_step"],
+
+            mode=self.mode,
+        )[0]
+
+        state["current_step"] += 1
+        return Q @ dir
+
+    def reproject(self, L_old, Q_old, L_new, Q_new, state):
+        if  "exp_avg" not in state: return
+
+        C = Q_new.T @ Q_old
+
+        state["exp_avg"] = C @ state["exp_avg"]
+        state["exp_avg_sq"] = _squared_reproject(C, state["exp_avg_sq"], exact=self.exact_reproject)
+        state["exp_avg_cu"] = C.pow(3) @ state["exp_avg_cu"] # exact reproject with 1_000_000 is feasible

torchzero/modules/experimental/curveball.py

@@ -1,25 +1,25 @@
 from typing import Literal
-from collections.abc import Callable
+
 import torch
 
-from ...core import Module, Target, Transform, Chainable, apply_transform
-from ...utils import NumberList, TensorList, as_tensorlist
-from ...utils.derivatives import hvp, hvp_fd_forward, hvp_fd_central
+from ...core import Chainable, Transform, step, HVPMethod
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+
 
 def curveball(
     tensors: TensorList,
     z_: TensorList,
-    Hz: TensorList,
+    Hzz: TensorList,
     momentum: float | NumberList,
     precond_lr: float | NumberList,
 ):
     """returns z_, clone it!!! (no just negate it)"""
-    delta = Hz + tensors
+    delta = Hzz + tensors
     z_.mul_(momentum).sub_(delta.mul_(precond_lr)) # z ← ρz − βΔ
     return z_
 
 
-class CurveBall(Module):
+class CurveBall(Transform):
     """CurveBall method from https://arxiv.org/pdf/1805.08095#page=4.09.
 
     For now this implementation does not include automatic ρ, α and β hyper-parameters in closed form, therefore it is expected to underperform compared to official implementation (https://github.com/jotaf98/pytorch-curveball/tree/master) so I moved this to experimental.
@@ -36,7 +36,7 @@ class CurveBall(Module):
         self,
         precond_lr: float=1e-3,
         momentum: float=0.9,
-        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
+        hvp_method: HVPMethod = "autograd",
         h: float = 1e-3,
         reg: float = 1,
         inner: Chainable | None = None,
@@ -44,46 +44,30 @@ class CurveBall(Module):
         defaults = dict(precond_lr=precond_lr, momentum=momentum, hvp_method=hvp_method, h=h, reg=reg)
         super().__init__(defaults)
 
-        if inner is not None: self.set_child('inner', inner)
+        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']
-        h = settings['h']
+    def apply_states(self, objective, states, settings):
+        params = objective.params
+        fs = settings[0]
+        hvp_method = fs['hvp_method']
+        h = fs['h']
 
-        precond_lr, momentum, reg = self.get_settings(params, 'precond_lr', 'momentum', 'reg', cls=NumberList)
+        precond_lr, momentum, reg = unpack_dicts(settings, 'precond_lr', 'momentum', 'reg', cls=NumberList)
 
-
-        closure = var.closure
+        closure = objective.closure
         assert closure is not None
 
-        z, Hz = self.get_state(params, 'z', 'Hz', cls=TensorList)
-
-        if hvp_method == 'autograd':
-            grad = var.get_grad(create_graph=True)
-            Hvp = hvp(params, grad, z)
-
-        elif hvp_method == 'forward':
-            loss, Hvp = hvp_fd_forward(closure, params, z, h=h, g_0=var.get_grad(), normalize=True)
-
-        elif hvp_method == 'central':
-            loss, Hvp = hvp_fd_central(closure, params, z, h=h, normalize=True)
-
-        else:
-            raise ValueError(hvp_method)
-
-
-        Hz.set_(Hvp + z*reg)
+        z, Hz = unpack_states(states, params, 'z', 'Hz', cls=TensorList)
+        Hz, _ = objective.hessian_vector_product(z, rgrad=None, at_x0=True, hvp_method=hvp_method, h=h)
 
+        Hz = TensorList(Hz)
+        Hzz = Hz.add_(z * reg)
 
-        update = var.get_update()
-        if 'inner' in self.children:
-            update = apply_transform(self.children['inner'], update, params, grads=var.grad, var=var)
+        objective = self.inner_step("inner", objective, must_exist=False)
+        updates = objective.get_updates()
 
-        z = curveball(TensorList(update), z, Hz, momentum=momentum, precond_lr=precond_lr)
-        var.update = z.neg()
+        z = curveball(TensorList(updates), z, Hzz, momentum=momentum, precond_lr=precond_lr)
+        objective.updates = z.neg()
 
-        return var
+        return objective

torchzero/modules/experimental/eigen_sr1.py

@@ -0,0 +1,182 @@
+import torch
+
+from ...core import Transform
+from ...linalg.orthogonalize import orthogonalize, OrthogonalizeMethod
+from ...linalg.eigh import eigh_plus_uuT, regularize_eigh
+from ...utils import TensorList, unpack_states, vec_to_tensors_
+from ..opt_utils import safe_clip
+from .eigengrad import _eigengrad_update_state_, eigengrad_apply
+
+
+def sr1_u(L: torch.Tensor, Q: torch.Tensor, s:torch.Tensor, y: torch.Tensor, tol:float):
+    """u from u u^T correction and its sign"""
+    r = y - torch.linalg.multi_dot([Q, L.diag_embed(), Q.T, s]) # pylint:disable=not-callable
+    rs = r.dot(s)
+
+    if rs.abs() < tol * torch.linalg.vector_norm(r) * torch.linalg.vector_norm(s): # pylint:disable=not-callable
+        return None, None
+
+    u = r / rs.abs().sqrt()
+    return u, torch.sign(rs)
+
+class EigenSR1(Transform):
+    def __init__(
+        self,
+        rank: int = 100,
+        tol: float = 1e-32,
+        eig_tol: float | None = None,
+        damping: float = 0,
+        rdamping: float = 0,
+        abs: bool = False,
+        mm_tol: float = 1e-7,
+        mm_truncate: int | None = None,
+        mm_damping: float = 1e-4,
+        mm_rdamping: float = 0,
+        mm_abs: bool = True,
+        id_reg: float | None = None,
+        column_space_tol=1e-9,
+        beta: float = 0.95,
+        balance_tol: float = 10,
+        balance_strength: float = 1e-1,
+
+        eigenbasis_optimizer = None,
+        update_freq: int = 1,
+        init_steps: int = 10,
+        orthogonalize_interval: int | None = 1,
+        orthogonalize_method: OrthogonalizeMethod = 'qr',
+
+        hvp_method = "autograd",
+        h = 1e-3,
+        inner = None,
+
+    ):
+        defaults = locals().copy()
+        for k in ["self", "inner"]:
+            del defaults[k]
+
+        super().__init__(defaults)
+
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
+        step = self.increment_counter("step", 0)
+
+        if step % fs["update_freq"] == 0:
+
+            params = TensorList(objective.params)
+
+            # compute y as hessian-vector product with s (random vecs during init steps)
+            if ("p_prev" not in self.global_state) or (step < fs["init_steps"]):
+                s_list = params.sample_like('rademacher')
+
+            else:
+                p_prev = self.global_state["p_prev"]
+                s_list = params - p_prev
+
+            if s_list.dot(s_list) < torch.finfo(s_list[0].dtype).tiny * 2:
+                s_list = params.sample_like('rademacher')
+
+            self.global_state["p_prev"] = params
+
+            # compute y as hessian-vector product with s
+            Hz_list, _ = objective.hessian_vector_product(s_list, rgrad=None, at_x0=True, hvp_method=fs["hvp_method"], h=fs["h"])
+
+            s = torch.cat([t.ravel() for t in s_list])
+            y = torch.cat([t.ravel() for t in Hz_list])
+
+            # keep track of exponential moving average of hessian diagonal and balance eigenvalues
+            if (fs["balance_strength"] != 0) and (step > fs["init_steps"]) and ("L" in self.global_state):
+
+                D = s * y # hutchinson estimator
+                exp_avg = self.global_state.get("exp_avg", None)
+
+                if exp_avg is None:
+                    exp_avg = self.global_state["exp_avg"] = D
+
+                else:
+                    exp_avg.lerp_(D, weight=1-fs["beta"])
+
+                L = self.global_state["L"]
+                L_abs = L.abs()
+                tau = L_abs.amax() / exp_avg.abs().amax()
+
+                if tau > fs["balance_tol"]:
+                    L_balanced = L_abs.pow((1 / tau) ** (1 / fs["balance_strength"])).copysign(L)
+                    self.global_state["L"] = torch.where(L_abs > 1, L_balanced, L)
+
+            # initialize L and Q on 1st step
+            if "L" not in self.global_state:
+
+                L = torch.zeros(1, dtype=s.dtype, device=s.device) # rank, rank
+                Q = torch.zeros([s.numel(), 1], dtype=s.dtype, device=s.device) # ndim, rank
+
+                u, sign = sr1_u(L=L, Q=Q, s=s, y=y, tol=0)
+                assert u is not None and sign is not None
+
+                # for uu^T u is eigenvector and u^T u is eigenvalue
+                norm = torch.linalg.vector_norm(u).clip(min=torch.finfo(u.dtype).tiny * 2) # pylint:disable=not-callable
+
+                self.global_state["L"] = self.global_state["L_reg"] = (u.dot(u).unsqueeze(0) / norm) * sign # (rank,)
+                self.global_state["Q"] = self.global_state["Q_reg"] = u.unsqueeze(-1) / norm # (m, rank)
+
+            # update hessian
+            else:
+                try:
+                    L = self.global_state["L"]
+                    Q = self.global_state["Q"]
+
+                    H_step = self.increment_counter("H_step", start=0)
+                    if H_step % fs["orthogonalize_interval"] == 0:
+                        Q = orthogonalize(Q, method=fs["orthogonalize_method"])
+
+                    u, sign = sr1_u(L=L, Q=Q, s=s, y=y, tol=fs["tol"])
+
+                    if (u is not None) and (sign is not None):
+
+                        # compute new factors
+                        L_new, Q_new = eigh_plus_uuT(L, Q, u, tol=fs["column_space_tol"], alpha=sign.item(), retry_float64=True)
+
+                        # truncate/regularize new factors (those go into the accumulator)
+                        L_new, Q_new = regularize_eigh(L=L_new, Q=Q_new, truncate=min(fs["rank"], s.numel()),
+                                                      tol=fs["eig_tol"], damping=fs["damping"], rdamping=fs["rdamping"])
+
+                        _eigengrad_update_state_(state=self.global_state, setting=fs, L_new=L_new, Q_new=Q_new)
+
+                except torch.linalg.LinAlgError:
+                    pass
+
+
+
+    def apply_states(self, objective, states, settings):
+        fs = settings[0]
+        updates = objective.get_updates()
+
+        if "eigenbasis_state" not in self.global_state:
+            self.global_state["eigenbasis_state"] = {}
+
+        step = self.global_state["step"] # starts at 0
+        if step < fs["init_steps"]:
+
+            # skip update first init_steps to let hessian kick-start
+            objective.stop = True
+            objective.skip_update = True
+            return objective
+
+        if "L_reg" not in self.global_state:
+            TensorList(updates).clip_(-0.1, 0.1)
+            return objective
+
+        dir = eigengrad_apply(
+            tensor = torch.cat([t.ravel() for t in updates]),
+            L_reg = self.global_state["L_reg"],
+            Q_reg = self.global_state["Q_reg"],
+            beta = None,
+            step = None,
+            debias = False,
+            id_reg = fs["id_reg"],
+            eigenbasis_optimizer = fs["eigenbasis_optimizer"],
+            eigenbasis_state = self.global_state["eigenbasis_state"],
+            whiten_fn = lambda x: x
+        )
+
+        vec_to_tensors_(dir, updates)
+        return objective

torchzero/modules/experimental/eigengrad.py

@@ -0,0 +1,207 @@
+# pylint: disable = non-ascii-name
+from collections.abc import Mapping
+
+import torch
+
+from ...core import Chainable, TensorTransform
+from ...linalg.eigh import eigh_plus_uuT, regularize_eigh
+from ...linalg.orthogonalize import OrthogonalizeMethod, orthogonalize
+from ...linalg.linear_operator import Eigendecomposition
+from ..adaptive.lre_optimizers import LREOptimizerBase
+
+
+def _eigengrad_update_state_(state:dict, setting: Mapping, L_new: torch.Tensor | None, Q_new:torch.Tensor | None):
+    """stores L, Q, L_reg, Q_reg and reprojects eigenbasis opt (this is also used on other eigen based modules)"""
+    if (L_new is not None) and (Q_new is not None):
+
+        # re-orthogonalize
+        orthogonalize_interval = setting["orthogonalize_interval"]
+        if orthogonalize_interval is not None:
+            Q_step = state.get("Q_step", 0)
+            state["Q_step"] = Q_step + 1
+            if Q_step % orthogonalize_interval == 0:
+                Q_new = orthogonalize(Q_new, method=setting["orthogonalize_method"])
+
+        # take absolute value (for hessian)
+        if setting.get("abs", False):
+            L_new = L_new.abs()
+
+        # store
+        state["L"] = L_new
+        state["Q"] = Q_new
+
+        # absolute value for matmul
+        if setting.get("mm_abs", False):
+            L_new = L_new.abs()
+
+        # regularize for matmul
+        # this second round of regularization is only used for preconditioning
+        # and doesn't affect the accumulator
+        L_reg_new, Q_reg_new = regularize_eigh(L=L_new, Q=Q_new,
+            truncate=setting["mm_truncate"],
+            tol=setting["mm_tol"],
+            damping=setting["mm_damping"],
+            rdamping=setting["mm_rdamping"],
+        )
+
+        # print(f'{state["L_reg"] = }, {L_reg_new = }')
+
+        # reproject eigenbasis optimizer
+        if (L_reg_new is not None) and (Q_reg_new is not None):
+            eigenbasis_optimizer: LREOptimizerBase | None = setting["eigenbasis_optimizer"]
+            if eigenbasis_optimizer is not None:
+                eigenbasis_optimizer.reproject(L_old=state["L_reg"], Q_old=state["Q_reg"], L_new=L_reg_new,
+                                                Q_new=Q_reg_new, state=state["eigenbasis_state"])
+
+            state["L_reg"] = L_reg_new
+            state["Q_reg"] = Q_reg_new
+
+
+def eigengrad_apply(
+    tensor: torch.Tensor,
+    L_reg: torch.Tensor,
+    Q_reg: torch.Tensor,
+    beta: float | None,
+    step: int | None,
+    debias: bool,
+    id_reg: float | None,
+    eigenbasis_optimizer: LREOptimizerBase | None,
+    eigenbasis_state: dict,
+
+    whiten_fn = torch.sqrt
+):
+    # debias
+    if debias:
+        assert beta is not None and step is not None
+        L_reg = L_reg / (1 - beta **step)
+
+    # step with eigenbasis optimizer
+    if eigenbasis_optimizer is not None:
+        if (id_reg is not None) and (id_reg != 0):
+            raise RuntimeError("id_reg is not compatible with eigenbasis_optimizer")
+
+        update = eigenbasis_optimizer.step(tensor.ravel(), L=L_reg, Q=Q_reg, state=eigenbasis_state)
+        return update.view_as(tensor)
+
+    # or just whiten
+    # L_reg = L_reg.clip(min=torch.finfo(L_reg.dtype).tiny * 2)
+
+    if id_reg is None or id_reg == 0:
+        G = Eigendecomposition(whiten_fn(L_reg), Q_reg, use_nystrom=False)
+        dir = G.solve(tensor.ravel())
+
+    else:
+        G = Eigendecomposition(whiten_fn(L_reg), Q_reg, use_nystrom=True)
+        dir = G.solve_plus_diag(tensor.ravel(), diag=id_reg)
+
+    return dir.view_as(tensor)
+
+class Eigengrad(TensorTransform):
+    """we can easily compute rank 1 symmetric update to a low rank eigendecomposition.
+    So this stores covariance matrix as it.
+
+
+    Args:
+        rank (int): maximum allowed rank
+        beta (float, optional): beta for covariance matrix exponential moving average. Defaults to 0.95.
+        eig_tol (float, optional):
+            removes eigenvalues this much smaller than largest eigenvalue when updating the preconditioner. Defaults to 1e-7.
+        damping (float, optional):
+            added to eigenvalues when updating the preconditioner. Defaults to 1e-8.
+        rdamping (float, optional):
+            added to eigenvalues when updating the preconditioner, relative to largest eigenvalue. Defaults to 0.
+        mm_tol (float, optional):
+            removes eigenvalues this much smaller than largest eigenvalue when computing the update. Defaults to 1e-7.
+        mm_truncate (int | None, optional):
+            uses top k eigenvalues to compute the update. Defaults to None.
+        mm_damping (float, optional):
+            added to eigenvalues when computing the update. Defaults to 1e-4.
+        mm_rdamping (float, optional):
+            added to eigenvalues when computing the update, relative to largest eigenvalue. Defaults to 0.
+        id_reg (float, optional):
+            multiplier to identity matrix added to preconditioner before computing update
+            If this value is given, solution from Nyström sketch-and-solve will be used to compute the update.
+            This value can't be too small (i.e. less than 1e-5) or the solver will be very unstable. Defaults to None.
+        column_space_tol (float, optional):
+            tolerance for deciding if new eigenvector is within column space of the covariance matrix. Defaults to 1e-9.
+        concat_params (bool, optional):
+            whether to precondition all parameters at once if True, or each separately if False. Defaults to True.
+        update_freq (int, optional): update frequency. Defaults to 1.
+        inner (Chainable | None, optional): inner modules. Defaults to None.
+
+    """
+
+    def __init__(
+        self,
+        rank: int = 100,
+        beta=0.95,
+        eig_tol: float | None = 1e-5,
+        damping: float = 0,
+        rdamping: float = 0,
+        mm_tol: float = 0,
+        mm_truncate: int | None = None,
+        mm_damping: float = 1e-4,
+        mm_rdamping: float = 0,
+        id_reg: float | None = None,
+        column_space_tol = 1e-9,
+
+        orthogonalize_interval: int | None = None,
+        orthogonalize_method: OrthogonalizeMethod = 'qr',
+
+        eigenbasis_optimizer: LREOptimizerBase | None = None,
+        concat_params: bool = True,
+        update_freq: int = 1,
+        inner: Chainable | None = None,
+    ):
+        defaults = locals().copy()
+        for k in ["self", "concat_params", "inner", "update_freq"]:
+            del defaults[k]
+
+        super().__init__(defaults, concat_params=concat_params, inner=inner, update_freq=update_freq)
+
+    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
+        state["step"] = state.get("step", 0) + 1
+        beta = setting["beta"]
+
+        if "L" not in state:
+            # for uu^T u is eigenvector and u^T u is eigenvalue
+            norm = torch.linalg.vector_norm(tensor).clip(min=torch.finfo(tensor.dtype).tiny * 2) # pylint:disable=not-callable
+
+            state["L"] = state["L_reg"] = (tensor.dot(tensor).unsqueeze(0) / norm) # (rank,)
+            state["Q"] = state["Q_reg"] = tensor.unsqueeze(-1) / norm # (m, rank)
+
+        else:
+            try:
+                L = state["L"]
+                Q = state["Q"]
+
+                # compute new factors
+                L_new, Q_new = eigh_plus_uuT(L*beta, Q, tensor, alpha=(1-beta), tol=setting["column_space_tol"], retry_float64=True)
+
+                # truncate/regularize new factors (those go into the accumulator)
+                L_new, Q_new = regularize_eigh(L=L_new, Q=Q_new, truncate=setting["rank"], tol=setting["eig_tol"],
+                                              damping=setting["damping"], rdamping=setting["rdamping"])
+
+                _eigengrad_update_state_(state=state, setting=setting, L_new=L_new, Q_new=Q_new)
+
+            except torch.linalg.LinAlgError:
+                pass
+
+    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
+        if "L_reg" not in state:
+            return tensor.clip(-0.1, 0.1)
+
+        if "eigenbasis_state" not in state:
+            state["eigenbasis_state"] = {}
+
+        return eigengrad_apply(
+            tensor = tensor,
+            L_reg = state["L_reg"],
+            Q_reg = state["Q_reg"],
+            beta = setting["beta"],
+            step = state["step"],
+            debias = True,
+            id_reg = setting["id_reg"],
+            eigenbasis_optimizer = setting["eigenbasis_optimizer"],
+            eigenbasis_state = state["eigenbasis_state"]
+        )

torchzero/modules/experimental/gradmin.py

@@ -5,7 +5,7 @@ from typing import Literal
 
 import torch
 
-from ...core import Module, Var, Chainable
+from ...core import Module, Objective, Chainable
 from ...utils import NumberList, TensorList
 from ...utils.derivatives import jacobian_wrt
 from ..grad_approximation import GradApproximator, GradTarget
@@ -43,7 +43,7 @@ class GradMin(Reformulation):
         super().__init__(defaults, modules=modules)
 
     @torch.no_grad
-    def closure(self, backward, closure, params, var):
+    def closure(self, backward, closure, params, objective):
         settings = self.settings[params[0]]
         loss_term = settings['loss_term']
         relative = settings['relative']