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

torchzero/modules/second_order/nystrom.py

@@ -1,153 +1,187 @@
-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, regularize_eigh, OrthogonalizeMethod
+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.
+        reg (float | None, optional):
+            scale of identity matrix added to hessian. Note that if this is specified, nystrom sketch-and-solve
+            is used to compute ``(Q diag(L) Q.T + reg*I)x = b``. It is very unstable when ``reg`` is small,
+            i.e. smaller than 1e-4. If this is None,``(Q diag(L) Q.T)x = b`` is computed by simply taking
+            reciprocal of eigenvalues. Defaults to 1e-3.
+        eigv_tol (float, optional):
+            all eigenvalues smaller than largest eigenvalue times ``eigv_tol`` are removed. Defaults to None.
+        truncate (int | None, optional):
+            keeps top ``truncate`` eigenvalues. Defaults to None.
+        damping (float, optional): scalar added to eigenvalues. Defaults to 0.
+        rdamping (float, optional): scalar multiplied by largest eigenvalue and added to eigenvalues. Defaults to 0.
+        update_freq (int, optional): frequency of updating preconditioner. Defaults to 1.
         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.Optimizer(
+        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.Optimizer(
+        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",
+        reg: float | None = 1e-2,
+        eigv_tol: float = 0,
+        truncate: int | None = None,
+        damping: float = 0,
+        rdamping: float = 0,
+        update_freq: int = 1,
+        orthogonalize_method: OrthogonalizeMethod = 'qr',
+        hvp_method: HVPMethod = "batched_autograd",
         h: float = 1e-3,
         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)
+        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)
+
+        # ---------------------------------- sketch ---------------------------------- #
+        ndim = sum(t.numel() for t in objective.params)
+        device = params[0].device
+        dtype = params[0].dtype
+
+        generator = self.get_generator(params[0].device, seed=fs['seed'])
+        try:
+            # compute the approximation
+            L, Q = nystrom_approximation(
+                A_mv=H_mv,
+                A_mm=H_mm,
+                ndim=ndim,
+                rank=min(fs["rank"], ndim),
+                eigv_tol=fs["eigv_tol"],
+                orthogonalize_method=fs["orthogonalize_method"],
+                dtype=dtype,
+                device=device,
+                generator=generator,
+            )
 
-            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])
+            # regularize
+            L, Q = regularize_eigh(
+                L=L,
+                Q=Q,
+                truncate=fs["truncate"],
+                tol=fs["eigv_tol"],
+                damping=fs["damping"],
+                rdamping=fs["rdamping"],
+            )
 
-        else:
+            # store
+            if L is not None:
+                self.global_state["L"] = L
+                self.global_state["Q"] = Q
 
-            with torch.enable_grad():
-                grad = var.get_grad()
+        except torch.linalg.LinAlgError as e:
+            warnings.warn(f"Nystrom approximation failed with: {e}")
 
-            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])
+    def apply_states(self, objective, states, settings):
+        if "L" not in self.global_state:
+            return objective
 
-            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])
+        fs = settings[0]
+        updates = objective.get_updates()
+        b=torch.cat([t.ravel() for t in updates])
 
-            else:
-                raise ValueError(hvp_method)
+        # ----------------------------------- solve ---------------------------------- #
+        L = self.global_state["L"]
+        Q = self.global_state["Q"]
 
+        if fs["reg"] is None:
+            x = Q @ ((Q.mH @ b) / L)
+        else:
+            x = nystrom_sketch_and_solve(L=L, Q=Q, b=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):
+        rank (int):
             size of the sketch for preconditioning, this many hessian-vector products will be evaluated before
             running the conjugate gradient solver. Larger value improves the preconditioning and speeds up
             conjugate gradient.
@@ -159,31 +193,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
-
-        .. code-block:: python
+    NystromPCG with backtracking line search
 
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.NystromPCG(10),
-                tz.m.Backtracking()
-            )
+    ```python
+    opt = tz.Optimizer(
+        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 +225,78 @@ 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
+        eigv_tol: float = 0,
+        orthogonalize_method: OrthogonalizeMethod = 'qr',
+        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:
-
-            with torch.enable_grad():
-                grad = var.get_grad()
-
-            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])
-
-            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])
-
-            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
-
+        # this should run on every update_states
+        _, H_mv, H_mm = objective.tensor_Hvp_function(hvp_method=fs['hvp_method'], h=fs['h'], at_x0=True)
+        objective.temp = H_mv
+
+        # --------------------------- update preconditioner -------------------------- #
+        step = self.increment_counter("step", 0)
+        if step % fs["update_freq"] == 0:
+
+            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'])
+
+            try:
+                L, Q = nystrom_approximation(
+                    A_mv=None,
+                    A_mm=H_mm,
+                    ndim=ndim,
+                    rank=min(fs["rank"], ndim),
+                    eigv_tol=fs["eigv_tol"],
+                    orthogonalize_method=fs["orthogonalize_method"],
+                    dtype=dtype,
+                    device=device,
+                    generator=generator,
+                )
+
+                self.global_state["L"] = L
+                self.global_state["Q"] = Q
+
+            except torch.linalg.LinAlgError as e:
+                warnings.warn(f"Nystrom approximation failed with: {e}")
 
+    @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