torchzero 0.3.15__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-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