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

torchzero/modules/second_order/newton_cg.py

@@ -13,21 +13,12 @@ from ..trust_region.trust_region import default_radius
 class NewtonCG(Module):
     """Newton's method with a matrix-free conjugate gradient or minimial-residual solver.
 
-    This optimizer implements Newton's method using a matrix-free conjugate
-    gradient (CG) or a minimal-residual (MINRES) solver to approximate the search direction. Instead of
-    forming the full Hessian matrix, it only requires Hessian-vector products
-    (HVPs). These can be calculated efficiently using automatic
-    differentiation or approximated using finite differences.
+    Notes:
+        * In most cases NewtonCGSteihaug 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::
-        In most cases NewtonCG 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.
+        * 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::
-        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).
-
-    .. warning::
+    Warning:
         CG may fail if hessian is not positive-definite.
 
     Args:
@@ -66,26 +57,24 @@ class NewtonCG(Module):
             NewtonCG will attempt to apply preconditioning to the output of this module.
 
     Examples:
-        Newton-CG with a backtracking line search:
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.NewtonCG(),
-                tz.m.Backtracking()
-            )
-
-        Truncated Newton method (useful for large-scale problems):
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.NewtonCG(maxiter=10, warm_start=True),
-                tz.m.Backtracking()
-            )
-
+    Newton-CG with a backtracking line search:
+
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.NewtonCG(),
+        tz.m.Backtracking()
+    )
+    ```
+
+    Truncated Newton method (useful for large-scale problems):
+    ```
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.NewtonCG(maxiter=10),
+        tz.m.Backtracking()
+    )
+    ```
 
     """
     def __init__(
@@ -95,7 +84,7 @@ class NewtonCG(Module):
         reg: float = 1e-8,
         hvp_method: Literal["forward", "central", "autograd"] = "autograd",
         solver: Literal['cg', 'minres', 'minres_npc'] = 'cg',
-        h: float = 1e-3,
+        h: float = 1e-3, # tuned 1e-4 or 1e-3
         miniter:int = 1,
         warm_start=False,
         inner: Chainable | None = None,
@@ -187,96 +176,95 @@ class NewtonCG(Module):
 
 
 class NewtonCGSteihaug(Module):
-    """Trust region Newton's method with a matrix-free Steihaug-Toint conjugate gradient or MINRES solver.
+    """Newton's method with trust region and a matrix-free Steihaug-Toint conjugate gradient solver.
 
-    This optimizer implements Newton's method using a matrix-free conjugate
-    gradient (CG) solver to approximate the search direction. Instead of
-    forming the full Hessian matrix, it only requires Hessian-vector products
-    (HVPs). These can be calculated efficiently using automatic
-    differentiation or approximated using finite differences.
+    Notes:
+        * In most cases NewtonCGSteihaug 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::
-        In most cases NewtonCGSteihaug 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.
-
-    .. 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).
-
-    .. warning::
-        CG may fail if hessian is not positive-definite.
+        * 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).
 
     Args:
-        maxiter (int | None, optional):
-            Maximum number of iterations for the conjugate gradient solver.
-            By default, this is set to the number of dimensions in the
-            objective function, which is the theoretical upper bound for CG
-            convergence. Setting this to a smaller value (truncated Newton)
-            can still generate good search directions. Defaults to None.
         eta (float, optional):
-            whenever actual to predicted loss reduction ratio is larger than this, a step is accepted.
-        nplus (float, optional):
-            trust region multiplier on successful steps.
-        nminus (float, optional):
-            trust region multiplier on unsuccessful steps.
-        init (float, optional): initial trust region.
+            if ratio of actual to predicted rediction is larger than this, step is accepted. Defaults to 0.0.
+        nplus (float, optional): increase factor on successful steps. Defaults to 1.5.
+        nminus (float, optional): decrease factor on unsuccessful steps. Defaults to 0.75.
+        rho_good (float, optional):
+            if ratio of actual to predicted rediction is larger than this, trust region size is multiplied by `nplus`.
+        rho_bad (float, optional):
+            if ratio of actual to predicted rediction is less than this, trust region size is multiplied by `nminus`.
+        init (float, optional): Initial trust region value. Defaults to 1.
+        max_attempts (max_attempts, optional):
+            maximum number of trust radius reductions per step. A zero update vector is returned when
+            this limit is exceeded. Defaults to 10.
+        max_history (int, optional):
+            CG will store this many intermediate solutions, reusing them when trust radius is reduced
+            instead of re-running CG. Each solution storage requires 2N memory. Defaults to 100.
+        boundary_tol (float | None, optional):
+            The trust region only increases when suggested step's norm is at least `(1-boundary_tol)*trust_region`.
+            This prevents increasing trust region when solution is not on the boundary. Defaults to 1e-2.
+
+        maxiter (int | None, optional):
+            maximum number of CG iterations per step. Each iteration requies one backward pass if `hvp_method="forward"`, two otherwise. Defaults to None.
+        miniter (int, optional):
+            minimal number of CG iterations. This prevents making no progress
         tol (float, optional):
-            Relative tolerance for the conjugate gradient solver to determine
-            convergence. Defaults to 1e-4.
-        reg (float, optional):
-            Regularization parameter (damping) added to the Hessian diagonal.
-            This helps ensure the system is positive-definite. Defaults to 1e-8.
+            terminates CG when norm of the residual is less than this value. Defaults to 1e-8.
+            when initial guess is below tolerance. Defaults to 1.
+        reg (float, optional): hessian regularization. Defaults to 1e-8.
+        solver (str, optional): solver, "cg" or "minres". "cg" is recommended. Defaults to 'cg'.
+        adapt_tol (bool, optional):
+            if True, whenever trust radius collapses to smallest representable number,
+            the tolerance is multiplied by 0.1. Defaults to True.
+        npc_terminate (bool, optional):
+            whether to terminate CG/MINRES whenever negative curvature is detected. Defaults to False.
+
         hvp_method (str, optional):
-            Determines how Hessian-vector products are evaluated.
+            either "forward" to use forward formula which requires one backward pass per Hvp, or "central" to use a more accurate central formula which requires two backward passes. "forward" is usually accurate enough. Defaults to "forward".
+        h (float, optional): finite difference step size. Defaults to 1e-3.
 
-            - ``"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):
-            The step size for finite differences if :code:`hvp_method` is
-            ``"forward"`` or ``"central"``. Defaults to 1e-3.
         inner (Chainable | None, optional):
-            NewtonCG will attempt to apply preconditioning to the output of this module.
+            applies preconditioning to output of this module. Defaults to None.
 
-    Examples:
-        Trust-region Newton-CG:
-
-        .. code-block:: python
+    ### Examples:
+    Trust-region Newton-CG:
 
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.NewtonCGSteihaug(),
-            )
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.NewtonCGSteihaug(),
+    )
+    ```
 
-    Reference:
+    ### Reference:
         Steihaug, Trond. "The conjugate gradient method and trust regions in large scale optimization." SIAM Journal on Numerical Analysis 20.3 (1983): 626-637.
     """
     def __init__(
         self,
-        maxiter: int | None = None,
+        # trust region settings
         eta: float= 0.0,
         nplus: float = 3.5,
         nminus: float = 0.25,
         rho_good: float = 0.99,
         rho_bad: float = 1e-4,
         init: float = 1,
-        tol: float = 1e-8,
-        reg: float = 1e-8,
-        hvp_method: Literal["forward", "central"] = "forward",
-        solver: Literal['cg', "minres"] = 'cg',
-        h: float = 1e-3,
         max_attempts: int = 100,
         max_history: int = 100,
-        boundary_tol: float = 1e-1,
+        boundary_tol: float = 1e-6, # tuned
+
+        # cg settings
+        maxiter: int | None = None,
         miniter: int = 1,
-        rms_beta: float | None = None,
+        tol: float = 1e-8,
+        reg: float = 1e-8,
+        solver: Literal['cg', "minres"] = 'cg',
         adapt_tol: bool = True,
         npc_terminate: bool = False,
+
+        # hvp settings
+        hvp_method: Literal["forward", "central"] = "central",
+        h: float = 1e-3, # tuned 1e-4 or 1e-3
+
+        # inner
         inner: Chainable | None = None,
     ):
         defaults = locals().copy()
@@ -336,19 +324,8 @@ class NewtonCGSteihaug(Module):
                 raise ValueError(hvp_method)
 
 
-        # ------------------------- update RMS preconditioner ------------------------ #
-        b = var.get_update()
-        P_mm = None
-        rms_beta = self.defaults["rms_beta"]
-        if rms_beta is not None:
-            exp_avg_sq = self.get_state(params, "exp_avg_sq", init=b, cls=TensorList)
-            exp_avg_sq.mul_(rms_beta).addcmul(b, b, value=1-rms_beta)
-            exp_avg_sq_sqrt = exp_avg_sq.sqrt().add_(1e-8)
-            def _P_mm(x):
-                return x / exp_avg_sq_sqrt
-            P_mm = _P_mm
-
         # -------------------------------- inner step -------------------------------- #
+        b = var.get_update()
         if 'inner' in self.children:
             b = apply_transform(self.children['inner'], b, params=params, grads=grad, var=var)
         b = as_tensorlist(b)
@@ -392,7 +369,6 @@ class NewtonCGSteihaug(Module):
                         miniter=miniter,
                         npc_terminate=npc_terminate,
                         history_size=max_history,
-                        P_mm=P_mm,
                     )
 
                 elif solver == 'minres':

torchzero/modules/second_order/nystrom.py

@@ -193,7 +193,7 @@ class NystromPCG(Module):
         self,
         sketch_size: int,
         maxiter=None,
-        tol=1e-3,
+        tol=1e-8,
         reg: float = 1e-6,
         hvp_method: Literal["forward", "central", "autograd"] = "autograd",
         h=1e-3,

torchzero/modules/second_order/rsn.py

@@ -0,0 +1,227 @@
+import math
+from collections import deque
+from collections.abc import Callable
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, apply_transform
+from ...utils import Distributions, TensorList, vec_to_tensors
+from ...utils.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
+    else:
+        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)
+
+class RSN(Module):
+    """Randomized Subspace Newton. Performs a Newton step in a random subspace.
+
+    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.
+            - "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 three directions set to gradient, slow EMA and fast EMA (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" - 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".
+        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: Literal["batched", "autograd", "forward", "central"] = "batched",
+        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 = dict(sketch_size=sketch_size, sketch_type=sketch_type,seed=seed,hvp_method=hvp_method, h=h, damping=damping, use_lstsq=use_lstsq, H_tfm=H_tfm, eigval_fn=eigval_fn, update_freq=update_freq)
+        super().__init__(defaults)
+
+        if inner is not None:
+            self.set_child("inner", inner)
+
+    @torch.no_grad
+    def update(self, var):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        if step % self.defaults['update_freq'] == 0:
+
+            closure = var.closure
+            if closure is None:
+                raise RuntimeError("RSN requires closure")
+            params = var.params
+            generator = self.get_generator(params[0].device, self.defaults["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(self.defaults["sketch_size"], ndim)
+            sketch_type = self.defaults["sketch_type"]
+            hvp_method = self.defaults["hvp_method"]
+
+            if sketch_type in ('normal', 'gaussian'):
+                S = _gaussian_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 = var.get_grad(create_graph=hvp_method in ("batched", "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 = var.get_grad(create_graph=hvp_method in ("batched", "autograd"))
+                g = torch.cat([t.ravel() for t in g_list])
+
+                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
+
+                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)
+
+                S = torch.stack([g, slow_ema, fast_ema], dim=1)
+                if sketch_size > 3:
+                    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, _ = var.hessian_matrix_product(S, at_x0=True, rgrad=None, hvp_method=self.defaults["hvp_method"], normalize=True, retain_graph=False, h=self.defaults["h"])
+            H_sketched = S.T @ HS
+
+            self.global_state["H_sketched"] = H_sketched
+            self.global_state["S"] = S
+
+    def apply(self, var):
+        S: torch.Tensor = self.global_state["S"]
+        d_proj = _newton_step(
+            var=var,
+            H=self.global_state["H_sketched"],
+            damping=self.defaults["damping"],
+            inner=self.children.get("inner", None),
+            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
+        var.update = vec_to_tensors(d, var.params)
+
+        return var
+
+    def get_H(self, var=...):
+        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)

torchzero/modules/trust_region/levenberg_marquardt.py

@@ -14,13 +14,13 @@ class LevenbergMarquardt(TrustRegionBase):
         hess_module (Module | None, optional):
             A module that maintains a hessian approximation (not hessian inverse!).
             This includes all full-matrix quasi-newton methods, ``tz.m.Newton`` and ``tz.m.GaussNewton``.
-            When using quasi-newton methods, set `inverse=False` when constructing them.
+            When using quasi-newton methods, set ``inverse=False`` when constructing them.
         y (float, optional):
             when ``y=0``, identity matrix is added to hessian, when ``y=1``, diagonal of the hessian approximation
             is added. Values between interpolate. This should only be used with Gauss-Newton. Defaults to 0.
         eta (float, optional):
             if ratio of actual to predicted rediction is larger than this, step is accepted.
-            When :code:`hess_module` is GaussNewton, this can be set to 0. Defaults to 0.15.
+            When ``hess_module`` is ``Newton`` or ``GaussNewton``, this can be set to 0. Defaults to 0.15.
         nplus (float, optional): increase factor on successful steps. Defaults to 1.5.
         nminus (float, optional): decrease factor on unsuccessful steps. Defaults to 0.75.
         rho_good (float, optional):

torchzero/modules/trust_region/trust_cg.py

@@ -60,17 +60,19 @@ class TrustCG(TrustRegionBase):
         nminus: float = 0.25,
         rho_good: float = 0.99,
         rho_bad: float = 1e-4,
-        boundary_tol: float | None = 1e-1,
+        boundary_tol: float | None = 1e-6, # tuned
         init: float = 1,
         max_attempts: int = 10,
         radius_strategy: _RadiusStrategy | _RADIUS_KEYS = 'default',
         reg: float = 0,
-        cg_tol: float = 1e-4,
+        maxiter: int | None = None,
+        miniter: int = 1,
+        cg_tol: float = 1e-8,
         prefer_exact: bool = True,
         update_freq: int = 1,
         inner: Chainable | None = None,
     ):
-        defaults = dict(reg=reg, prefer_exact=prefer_exact, cg_tol=cg_tol)
+        defaults = dict(reg=reg, prefer_exact=prefer_exact, cg_tol=cg_tol, maxiter=maxiter, miniter=miniter)
         super().__init__(
             defaults=defaults,
             hess_module=hess_module,
@@ -93,5 +95,5 @@ class TrustCG(TrustRegionBase):
         if settings['prefer_exact'] and isinstance(H, linear_operator.ScaledIdentity):
             return H.solve_bounded(g, radius)
 
-        x, _ = cg(H.matvec, g, trust_radius=radius, reg=settings['reg'], tol=settings["cg_tol"])
+        x, _ = cg(H.matvec, g, trust_radius=radius, reg=settings['reg'], maxiter=settings["maxiter"], miniter=settings["miniter"], tol=settings["cg_tol"])
         return x