torchzero 0.3.14__py3-none-any.whl → 0.4.0__py3-none-any.whl

torchzero/modules/second_order/nystrom.py

@@ -1,150 +1,141 @@
-from collections.abc import Callable
-from typing import Literal, overload
-import warnings
-import torch
+from typing import Literal
 
-from ...utils import TensorList, as_tensorlist, generic_zeros_like, generic_vector_norm, generic_numel, vec_to_tensors
-from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
+import torch
 
-from ...core import Chainable, apply_transform, Module
-from ...utils.linalg.solve import nystrom_sketch_and_solve, nystrom_pcg
+from ...core import Chainable, Transform, HVPMethod
+from ...utils import TensorList, vec_to_tensors
+from ...linalg import nystrom_pcg, nystrom_sketch_and_solve, nystrom_approximation, cg
+from ...linalg.linear_operator import Eigendecomposition, ScaledIdentity
 
-class NystromSketchAndSolve(Module):
+class NystromSketchAndSolve(Transform):
     """Newton's method with a Nyström sketch-and-solve solver.
 
-    .. note::
-        This module requires the 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).
-
-    .. note::
-        In most cases NystromSketchAndSolve should be the first module in the chain because it relies on autograd. Use the :code:`inner` argument if you wish to apply Newton preconditioning to another module's output.
+    Notes:
+        - This module requires the 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).
 
-    .. note::
-        If this is unstable, increase the :code:`reg` parameter and tune the rank.
+        - In most cases NystromSketchAndSolve should be the first module in the chain because it relies on autograd. Use the ``inner`` argument if you wish to apply Newton preconditioning to another module's output.
 
-    .. note:
-        :code:`tz.m.NystromPCG` usually outperforms this.
+        - If this is unstable, increase the ``reg`` parameter and tune the rank.
 
     Args:
         rank (int): size of the sketch, this many hessian-vector products will be evaluated per step.
         reg (float, optional): regularization parameter. Defaults to 1e-3.
         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".
-        h (float, optional): finite difference step size if :code:`hvp_method` is "forward" or "central". Defaults to 1e-3.
+            Determines how Hessian-vector products are computed.
+
+            - ``"batched_autograd"`` - uses autograd with batched hessian-vector products to compute the preconditioner. Faster than ``"autograd"`` but uses more memory.
+            - ``"autograd"`` - uses autograd hessian-vector products, uses a for loop to compute the preconditioner. Slower than ``"batched_autograd"`` but uses less memory.
+            - ``"fd_forward"`` - uses gradient finite difference approximation with a less accurate forward formula which requires one extra gradient evaluation per hessian-vector product.
+            - ``"fd_central"`` - uses gradient finite difference approximation with a more accurate central formula which requires two gradient evaluations per hessian-vector product.
+
+            Defaults to ``"autograd"``.
+        h (float, optional):
+            The step size for finite difference if ``hvp_method`` is
+            ``"fd_forward"`` or ``"fd_central"``. Defaults to 1e-3.
         inner (Chainable | None, optional): modules to apply hessian preconditioner to. Defaults to None.
         seed (int | None, optional): seed for random generator. Defaults to None.
 
+
     Examples:
-        NystromSketchAndSolve with backtracking line search
+    NystromSketchAndSolve with backtracking line search
 
-        .. code-block:: python
+    ```py
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.NystromSketchAndSolve(100),
+        tz.m.Backtracking()
+    )
+    ```
 
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.NystromSketchAndSolve(10),
-                tz.m.Backtracking()
-            )
+    Trust region NystromSketchAndSolve
+
+    ```py
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.LevenbergMarquadt(tz.m.NystromSketchAndSolve(100)),
+    )
+    ```
+
+    References:
+    - [Frangella, Z., Rathore, P., Zhao, S., & Udell, M. (2024). SketchySGD: Reliable Stochastic Optimization via Randomized Curvature Estimates. SIAM Journal on Mathematics of Data Science, 6(4), 1173-1204.](https://arxiv.org/pdf/2211.08597)
+    - [Frangella, Z., Tropp, J. A., & Udell, M. (2023). Randomized nyström preconditioning. SIAM Journal on Matrix Analysis and Applications, 44(2), 718-752](https://arxiv.org/abs/2110.02820)
 
-    Reference:
-        Frangella, Z., Tropp, J. A., & Udell, M. (2023). Randomized nyström preconditioning. SIAM Journal on Matrix Analysis and Applications, 44(2), 718-752. https://arxiv.org/abs/2110.02820
     """
     def __init__(
         self,
         rank: int,
         reg: float = 1e-3,
-        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
+        hvp_method: HVPMethod = "batched_autograd",
         h: float = 1e-3,
+        update_freq: int = 1,
         inner: Chainable | None = None,
         seed: int | None = None,
     ):
-        defaults = dict(rank=rank, reg=reg, hvp_method=hvp_method, h=h, seed=seed)
-        super().__init__(defaults,)
-
-        if inner is not None:
-            self.set_child('inner', inner)
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner'], defaults["update_freq"]
+        super().__init__(defaults, update_freq=update_freq, inner=inner)
 
     @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-
-        closure = var.closure
-        if closure is None: raise RuntimeError('NewtonCG requires closure')
-
-        settings = self.settings[params[0]]
-        rank = settings['rank']
-        reg = settings['reg']
-        hvp_method = settings['hvp_method']
-        h = settings['h']
-
-        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']
+    def update_states(self, objective, states, settings):
+        params = TensorList(objective.params)
+        fs = settings[0]
 
         # ---------------------- Hessian vector product function --------------------- #
-        if hvp_method == 'autograd':
-            grad = var.get_grad(create_graph=True)
-
-            def H_mm(x):
-                with torch.enable_grad():
-                    Hvp = hvp(params, grad, params.from_vec(x), retain_graph=True)
-                    return torch.cat([t.ravel() for t in Hvp])
+        hvp_method = fs['hvp_method']
+        h = fs['h']
+        _, H_mv, H_mm = objective.tensor_Hvp_function(hvp_method=hvp_method, h=h, at_x0=True)
 
-        else:
+        # ---------------------------------- sketch ---------------------------------- #
+        ndim = sum(t.numel() for t in objective.params)
+        device = params[0].device
+        dtype = params[0].dtype
 
-            with torch.enable_grad():
-                grad = var.get_grad()
+        generator = self.get_generator(params[0].device, seed=fs['seed'])
+        try:
+            L, Q = nystrom_approximation(A_mv=H_mv, A_mm=H_mm, ndim=ndim, rank=fs['rank'],
+                                        dtype=dtype, device=device, generator=generator)
 
-            if hvp_method == 'forward':
-                def H_mm(x):
-                    Hvp = hvp_fd_forward(closure, params, params.from_vec(x), h=h, g_0=grad, normalize=True)[1]
-                    return torch.cat([t.ravel() for t in Hvp])
+            self.global_state["L"] = L
+            self.global_state["Q"] = Q
+        except torch.linalg.LinAlgError:
+            pass
 
-            elif hvp_method == 'central':
-                def H_mm(x):
-                    Hvp = hvp_fd_central(closure, params, params.from_vec(x), h=h, normalize=True)[1]
-                    return torch.cat([t.ravel() for t in Hvp])
+    def apply_states(self, objective, states, settings):
+        fs = settings[0]
+        b = objective.get_updates()
 
-            else:
-                raise ValueError(hvp_method)
+        # ----------------------------------- solve ---------------------------------- #
+        if "L" not in self.global_state:
+            return objective
 
+        L = self.global_state["L"]
+        Q = self.global_state["Q"]
+        x = nystrom_sketch_and_solve(L=L, Q=Q, b=torch.cat([t.ravel() for t in b]), reg=fs["reg"])
 
-        # -------------------------------- inner step -------------------------------- #
-        b = var.get_update()
-        if 'inner' in self.children:
-            b = apply_transform(self.children['inner'], b, params=params, grads=grad, var=var)
+        # -------------------------------- set update -------------------------------- #
+        objective.updates = vec_to_tensors(x, reference=objective.params)
+        return objective
 
-        # ------------------------------ sketch&n&solve ------------------------------ #
-        x = nystrom_sketch_and_solve(A_mm=H_mm, b=torch.cat([t.ravel() for t in b]), rank=rank, reg=reg, generator=generator)
-        var.update = vec_to_tensors(x, reference=params)
-        return var
+    def get_H(self, objective=...):
+        if "L" not in self.global_state:
+            return ScaledIdentity()
 
+        L = self.global_state["L"]
+        Q = self.global_state["Q"]
+        return Eigendecomposition(L, Q)
 
 
-class NystromPCG(Module):
+class NystromPCG(Transform):
     """Newton's method with a Nyström-preconditioned conjugate gradient solver.
     This tends to outperform NewtonCG but requires tuning sketch size.
     An adaptive version exists in https://arxiv.org/abs/2110.02820, I might implement it too at some point.
 
-    .. note::
-        This module requires the a closure passed to the optimizer step,
+    Notes:
+        - This module requires the 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).
 
-    .. note::
-        In most cases NystromPCG should be the first module in the chain because it relies on autograd. Use the :code:`inner` argument if you wish to apply Newton preconditioning to another module's output.
+        - In most cases NystromPCG should be the first module in the chain because it relies on autograd. Use the ``inner`` argument if you wish to apply Newton preconditioning to another module's output.
 
     Args:
         sketch_size (int):
@@ -159,31 +150,31 @@ class NystromPCG(Module):
         tol (float, optional): relative tolerance for conjugate gradient solver. Defaults to 1e-4.
         reg (float, optional): regularization parameter. Defaults to 1e-8.
         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".
-        h (float, optional): finite difference step size if :code:`hvp_method` is "forward" or "central". Defaults to 1e-3.
+            Determines how Hessian-vector products are computed.
+
+            - ``"batched_autograd"`` - uses autograd with batched hessian-vector products to compute the preconditioner. Faster than ``"autograd"`` but uses more memory.
+            - ``"autograd"`` - uses autograd hessian-vector products, uses a for loop to compute the preconditioner. Slower than ``"batched_autograd"`` but uses less memory.
+            - ``"fd_forward"`` - uses gradient finite difference approximation with a less accurate forward formula which requires one extra gradient evaluation per hessian-vector product.
+            - ``"fd_central"`` - uses gradient finite difference approximation with a more accurate central formula which requires two gradient evaluations per hessian-vector product.
+
+            Defaults to ``"autograd"``.
+        h (float, optional):
+            The step size for finite difference if ``hvp_method`` is
+            ``"fd_forward"`` or ``"fd_central"``. Defaults to 1e-3.
         inner (Chainable | None, optional): modules to apply hessian preconditioner to. Defaults to None.
         seed (int | None, optional): seed for random generator. Defaults to None.
 
     Examples:
 
-        NystromPCG with backtracking line search
+    NystromPCG with backtracking line search
 
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.NystromPCG(10),
-                tz.m.Backtracking()
-            )
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.NystromPCG(10),
+        tz.m.Backtracking()
+    )
+    ```
 
     Reference:
         Frangella, Z., Tropp, J. A., & Udell, M. (2023). Randomized nyström preconditioning. SIAM Journal on Matrix Analysis and Applications, 44(2), 718-752. https://arxiv.org/abs/2110.02820
@@ -191,81 +182,70 @@ class NystromPCG(Module):
     """
     def __init__(
         self,
-        sketch_size: int,
+        rank: int,
         maxiter=None,
-        tol=1e-3,
+        tol=1e-8,
         reg: float = 1e-6,
-        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
+        update_freq: int = 1, # here update_freq is within update_states
+        hvp_method: HVPMethod = "batched_autograd",
         h=1e-3,
         inner: Chainable | None = None,
         seed: int | None = None,
     ):
-        defaults = dict(sketch_size=sketch_size, reg=reg, maxiter=maxiter, tol=tol, hvp_method=hvp_method, h=h, seed=seed)
-        super().__init__(defaults,)
-
-        if inner is not None:
-            self.set_child('inner', inner)
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner']
+        super().__init__(defaults, inner=inner)
 
     @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-
-        closure = var.closure
-        if closure is None: raise RuntimeError('NewtonCG requires closure')
-
-        settings = self.settings[params[0]]
-        sketch_size = settings['sketch_size']
-        maxiter = settings['maxiter']
-        tol = settings['tol']
-        reg = settings['reg']
-        hvp_method = settings['hvp_method']
-        h = settings['h']
-
-
-        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']
-
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
 
         # ---------------------- Hessian vector product function --------------------- #
-        if hvp_method == 'autograd':
-            grad = var.get_grad(create_graph=True)
-
-            def H_mm(x):
-                with torch.enable_grad():
-                    Hvp = hvp(params, grad, params.from_vec(x), retain_graph=True)
-                    return torch.cat([t.ravel() for t in Hvp])
-
-        else:
+        # this should run on every update_states
+        hvp_method = fs['hvp_method']
+        h = fs['h']
+        _, H_mv, H_mm = objective.tensor_Hvp_function(hvp_method=hvp_method, h=h, at_x0=True)
+        objective.temp = H_mv
 
-            with torch.enable_grad():
-                grad = var.get_grad()
+        # --------------------------- update preconditioner -------------------------- #
+        step = self.increment_counter("step", 0)
+        update_freq = self.defaults["update_freq"]
 
-            if hvp_method == 'forward':
-                def H_mm(x):
-                    Hvp = hvp_fd_forward(closure, params, params.from_vec(x), h=h, g_0=grad, normalize=True)[1]
-                    return torch.cat([t.ravel() for t in Hvp])
+        if step % update_freq == 0:
 
-            elif hvp_method == 'central':
-                def H_mm(x):
-                    Hvp = hvp_fd_central(closure, params, params.from_vec(x), h=h, normalize=True)[1]
-                    return torch.cat([t.ravel() for t in Hvp])
+            rank = fs['rank']
+            ndim = sum(t.numel() for t in objective.params)
+            device = objective.params[0].device
+            dtype = objective.params[0].dtype
+            generator = self.get_generator(device, seed=fs['seed'])
 
-            else:
-                raise ValueError(hvp_method)
-
-
-        # -------------------------------- inner step -------------------------------- #
-        b = var.get_update()
-        if 'inner' in self.children:
-            b = apply_transform(self.children['inner'], b, params=params, grads=grad, var=var)
-
-        # ------------------------------ sketch&n&solve ------------------------------ #
-        x = nystrom_pcg(A_mm=H_mm, b=torch.cat([t.ravel() for t in b]), sketch_size=sketch_size, reg=reg, tol=tol, maxiter=maxiter, x0_=None, generator=generator)
-        var.update = vec_to_tensors(x, reference=params)
-        return var
+            try:
+                L, Q = nystrom_approximation(A_mv=None, A_mm=H_mm, ndim=ndim, rank=rank,
+                                            dtype=dtype, device=device, generator=generator)
 
+                self.global_state["L"] = L
+                self.global_state["Q"] = Q
+            except torch.linalg.LinAlgError:
+                pass
 
+    @torch.no_grad
+    def apply_states(self, objective, states, settings):
+        b = objective.get_updates()
+        H_mv = objective.poptemp()
+        fs = self.settings[objective.params[0]]
+
+        # ----------------------------------- solve ---------------------------------- #
+        if "L" not in self.global_state:
+            # fallback on cg
+            sol = cg(A_mv=H_mv, b=TensorList(b), tol=fs["tol"], reg=fs["reg"], maxiter=fs["maxiter"])
+            objective.updates = sol.x
+            return objective
+
+        L = self.global_state["L"]
+        Q = self.global_state["Q"]
+        x = nystrom_pcg(L=L, Q=Q, A_mv=H_mv, b=torch.cat([t.ravel() for t in b]),
+                        reg=fs['reg'], tol=fs["tol"], maxiter=fs["maxiter"])
+
+        # -------------------------------- set update -------------------------------- #
+        objective.updates = vec_to_tensors(x, reference=objective.params)
+        return objective

torchzero/modules/second_order/rsn.py

@@ -0,0 +1,237 @@
+import math
+from collections import deque
+from collections.abc import Callable
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Transform, HVPMethod
+from ...utils import vec_to_tensors
+from ...linalg.linear_operator import Sketched
+
+from .newton import _newton_step
+
+def _qr_orthonormalize(A:torch.Tensor):
+    m,n = A.shape
+    if m < n:
+        q, _ = torch.linalg.qr(A.T) # pylint:disable=not-callable
+        return q.T
+
+    q, _ = torch.linalg.qr(A) # pylint:disable=not-callable
+    return q
+
+def _orthonormal_sketch(m, n, dtype, device, generator):
+    return _qr_orthonormalize(torch.randn(m, n, dtype=dtype, device=device, generator=generator))
+
+def _gaussian_sketch(m, n, dtype, device, generator):
+    return torch.randn(m, n, dtype=dtype, device=device, generator=generator) / math.sqrt(m)
+
+def _rademacher_sketch(m, n, dtype, device, generator):
+    rademacher = torch.bernoulli(torch.full((m,n), 0.5), generator = generator).mul_(2).sub_(1)
+    return rademacher.mul_(1 / math.sqrt(m))
+
+class SubspaceNewton(Transform):
+    """Subspace Newton. Performs a Newton step in a subspace (random or spanned by past gradients).
+
+    Args:
+        sketch_size (int):
+            size of the random sketch. This many hessian-vector products will need to be evaluated each step.
+        sketch_type (str, optional):
+            - "orthonormal" - random orthonormal basis. Orthonormality is necessary to use linear operator based modules such as trust region, but it can be slower to compute.
+            - "rademacher" - approximately orthonormal scaled random rademacher basis.
+            - "gaussian" - random gaussian (not orthonormal) basis.
+            - "common_directions" - uses history steepest descent directions as the basis[2]. It is orthonormalized on-line using Gram-Schmidt.
+            - "mixed" - random orthonormal basis but with four directions set to gradient, slow and fast gradient EMAs, and previous update direction (default).
+        damping (float, optional): hessian damping (scale of identity matrix added to hessian). Defaults to 0.
+        hvp_method (str, optional):
+            How to compute hessian-matrix product:
+            - "batched_autograd" - uses batched autograd
+            - "autograd" - uses unbatched autograd
+            - "forward" - uses finite difference with forward formula, performing 1 backward pass per Hvp.
+            - "central" - uses finite difference with a more accurate central formula, performing 2 backward passes per Hvp.
+
+            . Defaults to "batched_autograd".
+        h (float, optional): finite difference step size. Defaults to 1e-2.
+        use_lstsq (bool, optional): whether to use least squares to solve ``Hx=g``. Defaults to False.
+        update_freq (int, optional): frequency of updating the hessian. Defaults to 1.
+        H_tfm (Callable | None, optional):
+            optional hessian transforms, takes in two arguments - `(hessian, gradient)`.
+
+            must return either a tuple: `(hessian, is_inverted)` with transformed hessian and a boolean value
+            which must be True if transform inverted the hessian and False otherwise.
+
+            Or it returns a single tensor which is used as the update.
+
+            Defaults to None.
+        eigval_fn (Callable | None, optional):
+            optional eigenvalues transform, for example ``torch.abs`` or ``lambda L: torch.clip(L, min=1e-8)``.
+            If this is specified, eigendecomposition will be used to invert the hessian.
+        seed (int | None, optional): seed for random generator. Defaults to None.
+        inner (Chainable | None, optional): preconditions output of this module. Defaults to None.
+
+    ### Examples
+
+    RSN with line search
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.RSN(),
+        tz.m.Backtracking()
+    )
+    ```
+
+    RSN with trust region
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.LevenbergMarquardt(tz.m.RSN()),
+    )
+    ```
+
+
+    References:
+        1. [Gower, Robert, et al. "RSN: randomized subspace Newton." Advances in Neural Information Processing Systems 32 (2019).](https://arxiv.org/abs/1905.10874)
+        2. Wang, Po-Wei, Ching-pei Lee, and Chih-Jen Lin. "The common-directions method for regularized empirical risk minimization." Journal of Machine Learning Research 20.58 (2019): 1-49.
+    """
+
+    def __init__(
+        self,
+        sketch_size: int,
+        sketch_type: Literal["orthonormal", "gaussian", "common_directions", "mixed"] = "mixed",
+        damping:float=0,
+        hvp_method: HVPMethod = "batched_autograd",
+        h: float = 1e-2,
+        use_lstsq: bool = True,
+        update_freq: int = 1,
+        H_tfm: Callable[[torch.Tensor, torch.Tensor], tuple[torch.Tensor, bool]] | Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
+        eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        seed: int | None = None,
+        inner: Chainable | None = None,
+    ):
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner'], defaults["update_freq"]
+        super().__init__(defaults, update_freq=update_freq, inner=inner)
+
+    @torch.no_grad
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
+        params = objective.params
+        generator = self.get_generator(params[0].device, fs["seed"])
+
+        ndim = sum(p.numel() for p in params)
+
+        device=params[0].device
+        dtype=params[0].dtype
+
+        # sample sketch matrix S: (ndim, sketch_size)
+        sketch_size = min(fs["sketch_size"], ndim)
+        sketch_type = fs["sketch_type"]
+        hvp_method = fs["hvp_method"]
+
+        if sketch_type in ('normal', 'gaussian'):
+            S = _gaussian_sketch(ndim, sketch_size, device=device, dtype=dtype, generator=generator)
+
+        elif sketch_type == "rademacher":
+            S = _rademacher_sketch(ndim, sketch_size, device=device, dtype=dtype, generator=generator)
+
+        elif sketch_type == 'orthonormal':
+            S = _orthonormal_sketch(ndim, sketch_size, device=device, dtype=dtype, generator=generator)
+
+        elif sketch_type == 'common_directions':
+            # Wang, Po-Wei, Ching-pei Lee, and Chih-Jen Lin. "The common-directions method for regularized empirical risk minimization." Journal of Machine Learning Research 20.58 (2019): 1-49.
+            g_list = objective.get_grads(create_graph=hvp_method in ("batched_autograd", "autograd"))
+            g = torch.cat([t.ravel() for t in g_list])
+
+            # initialize directions deque
+            if "directions" not in self.global_state:
+
+                g_norm = torch.linalg.vector_norm(g) # pylint:disable=not-callable
+                if g_norm < torch.finfo(g.dtype).tiny * 2:
+                    g = torch.randn_like(g)
+                    g_norm = torch.linalg.vector_norm(g) # pylint:disable=not-callable
+
+                self.global_state["directions"] = deque([g / g_norm], maxlen=sketch_size)
+                S = self.global_state["directions"][0].unsqueeze(1)
+
+            # add new steepest descent direction orthonormal to existing columns
+            else:
+                S = torch.stack(tuple(self.global_state["directions"]), dim=1)
+                p = g - S @ (S.T @ g)
+                p_norm = torch.linalg.vector_norm(p) # pylint:disable=not-callable
+                if p_norm > torch.finfo(p.dtype).tiny * 2:
+                    p = p / p_norm
+                    self.global_state["directions"].append(p)
+                    S = torch.cat([S, p.unsqueeze(1)], dim=1)
+
+        elif sketch_type == "mixed":
+            g_list = objective.get_grads(create_graph=hvp_method in ("batched_autograd", "autograd"))
+            g = torch.cat([t.ravel() for t in g_list])
+
+            # initialize state
+            if "slow_ema" not in self.global_state:
+                self.global_state["slow_ema"] = torch.randn_like(g) * 1e-2
+                self.global_state["fast_ema"] = torch.randn_like(g) * 1e-2
+                self.global_state["p_prev"] = torch.randn_like(g)
+
+            # previous update direction
+            p_cur = torch.cat([t.ravel() for t in params])
+            prev_dir = p_cur - self.global_state["p_prev"]
+            self.global_state["p_prev"] = p_cur
+
+            # EMAs
+            slow_ema = self.global_state["slow_ema"]
+            fast_ema = self.global_state["fast_ema"]
+            slow_ema.lerp_(g, 0.001)
+            fast_ema.lerp_(g, 0.1)
+
+            # form and orthogonalize sketching matrix
+            S = torch.stack([g, slow_ema, fast_ema, prev_dir], dim=1)
+            if sketch_size > 4:
+                S_random = _gaussian_sketch(ndim, sketch_size - 3, device=device, dtype=dtype, generator=generator)
+                S = torch.cat([S, S_random], dim=1)
+
+            S = _qr_orthonormalize(S)
+
+        else:
+            raise ValueError(f'Unknown sketch_type {sketch_type}')
+
+        # form sketched hessian
+        HS, _ = objective.hessian_matrix_product(S, rgrad=None, at_x0=True,
+                                                 hvp_method=fs["hvp_method"], h=fs["h"])
+        H_sketched = S.T @ HS
+
+        self.global_state["H_sketched"] = H_sketched
+        self.global_state["S"] = S
+
+    def apply_states(self, objective, states, settings):
+        S: torch.Tensor = self.global_state["S"]
+
+        d_proj = _newton_step(
+            objective=objective,
+            H=self.global_state["H_sketched"],
+            damping=self.defaults["damping"],
+            H_tfm=self.defaults["H_tfm"],
+            eigval_fn=self.defaults["eigval_fn"],
+            use_lstsq=self.defaults["use_lstsq"],
+            g_proj = lambda g: S.T @ g
+        )
+
+        d = S @ d_proj
+        objective.updates = vec_to_tensors(d, objective.params)
+        return objective
+
+    def get_H(self, objective=...):
+        eigval_fn = self.defaults["eigval_fn"]
+        H_sketched: torch.Tensor = self.global_state["H_sketched"]
+        S: torch.Tensor = self.global_state["S"]
+
+        if eigval_fn is not None:
+            try:
+                L, Q = torch.linalg.eigh(H_sketched) # pylint:disable=not-callable
+                L: torch.Tensor = eigval_fn(L)
+                H_sketched = Q @ L.diag_embed() @ Q.mH
+
+            except torch.linalg.LinAlgError:
+                pass
+
+        return Sketched(S, H_sketched)