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

torchzero/modules/second_order/ifn.py

@@ -0,0 +1,89 @@
+import warnings
+from collections.abc import Callable
+from functools import partial
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, apply_transform, Var
+from ...utils import TensorList, vec_to_tensors
+from ...utils.linalg.linear_operator import DenseWithInverse, Dense
+from .newton import _get_H, _get_loss_grad_and_hessian, _newton_step
+
+
+class InverseFreeNewton(Module):
+    """Inverse-free newton's method
+
+    .. note::
+        In most cases Newton 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 the hessian.
+        The closure must accept a ``backward`` argument (refer to documentation).
+
+    .. warning::
+        this uses roughly O(N^2) memory.
+
+    Reference
+        [Massalski, Marcin, and Magdalena Nockowska-Rosiak. "INVERSE-FREE NEWTON'S METHOD." Journal of Applied Analysis & Computation 15.4 (2025): 2238-2257.](https://www.jaac-online.com/article/doi/10.11948/20240428)
+    """
+    def __init__(
+        self,
+        update_freq: int = 1,
+        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
+        vectorize: bool = True,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(hessian_method=hessian_method, vectorize=vectorize, update_freq=update_freq)
+        super().__init__(defaults)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def update(self, var):
+        update_freq = self.defaults['update_freq']
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        if step % update_freq == 0:
+            loss, g_list, H = _get_loss_grad_and_hessian(
+                var, self.defaults['hessian_method'], self.defaults['vectorize']
+            )
+            self.global_state["H"] = H
+
+            # inverse free part
+            if 'Y' not in self.global_state:
+                num = H.T
+                denom = (torch.linalg.norm(H, 1) * torch.linalg.norm(H, float('inf'))) # pylint:disable=not-callable
+
+                finfo = torch.finfo(H.dtype)
+                self.global_state['Y'] = num.div_(denom.clip(min=finfo.tiny * 2, max=finfo.max / 2))
+
+            else:
+                Y = self.global_state['Y']
+                I2 = torch.eye(Y.size(0), device=Y.device, dtype=Y.dtype).mul_(2)
+                I2 -= H @ Y
+                self.global_state['Y'] = Y @ I2
+
+
+    def apply(self, var):
+        Y = self.global_state["Y"]
+        params = var.params
+
+        # -------------------------------- inner step -------------------------------- #
+        update = var.get_update()
+        if 'inner' in self.children:
+            update = apply_transform(self.children['inner'], update, params=params, grads=var.grad, var=var)
+
+        g = torch.cat([t.ravel() for t in update])
+
+        # ----------------------------------- solve ---------------------------------- #
+        var.update = vec_to_tensors(Y@g, params)
+
+        return var
+
+    def get_H(self,var):
+        return DenseWithInverse(A = self.global_state["H"], A_inv=self.global_state["Y"])

torchzero/modules/second_order/inm.py

@@ -0,0 +1,105 @@
+from collections.abc import Callable
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module
+from ...utils import TensorList, vec_to_tensors
+from ..functional import safe_clip
+from .newton import _get_H, _get_loss_grad_and_hessian, _newton_step
+
+@torch.no_grad
+def inm(f:torch.Tensor, J:torch.Tensor, s:torch.Tensor, y:torch.Tensor):
+
+    yy = safe_clip(y.dot(y))
+    ss = safe_clip(s.dot(s))
+
+    term1 = y.dot(y - J@s) / yy
+    FbT = f.outer(s).mul_(term1 / ss)
+
+    P = FbT.add_(J)
+    return P
+
+def _eigval_fn(J: torch.Tensor, fn) -> torch.Tensor:
+    if fn is None: return J
+    L, Q = torch.linalg.eigh(J) # pylint:disable=not-callable
+    return (Q * L.unsqueeze(-2)) @ Q.mH
+
+class INM(Module):
+    """Improved Newton's Method (INM).
+
+    Reference:
+        [Saheya, B., et al. "A new Newton-like method for solving nonlinear equations." SpringerPlus 5.1 (2016): 1269.](https://d-nb.info/1112813721/34)
+    """
+
+    def __init__(
+        self,
+        damping: float = 0,
+        use_lstsq: bool = False,
+        update_freq: int = 1,
+        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
+        vectorize: bool = True,
+        inner: Chainable | None = None,
+        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,
+    ):
+        defaults = dict(damping=damping, hessian_method=hessian_method, use_lstsq=use_lstsq, vectorize=vectorize, 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):
+        update_freq = self.defaults['update_freq']
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        if step % update_freq == 0:
+            _, f_list, J = _get_loss_grad_and_hessian(
+                var, self.defaults['hessian_method'], self.defaults['vectorize']
+            )
+
+            f = torch.cat([t.ravel() for t in f_list])
+            J = _eigval_fn(J, self.defaults["eigval_fn"])
+
+            x_list = TensorList(var.params)
+            f_list = TensorList(var.get_grad())
+            x_prev, f_prev = self.get_state(var.params, "x_prev", "f_prev", cls=TensorList)
+
+            # initialize on 1st step, do Newton step
+            if step == 0:
+                x_prev.copy_(x_list)
+                f_prev.copy_(f_list)
+                self.global_state["P"] = J
+                return
+
+            # INM update
+            s_list = x_list - x_prev
+            y_list = f_list - f_prev
+            x_prev.copy_(x_list)
+            f_prev.copy_(f_list)
+
+            self.global_state["P"] = inm(f, J, s=s_list.to_vec(), y=y_list.to_vec())
+
+
+    @torch.no_grad
+    def apply(self, var):
+        params = var.params
+        update = _newton_step(
+            var=var,
+            H = self.global_state["P"],
+            damping=self.defaults["damping"],
+            inner=self.children.get("inner", None),
+            H_tfm=self.defaults["H_tfm"],
+            eigval_fn=None, # it is applied in `update`
+            use_lstsq=self.defaults["use_lstsq"],
+        )
+
+        var.update = vec_to_tensors(update, params)
+
+        return var
+
+    def get_H(self,var=...):
+        return _get_H(self.global_state["P"], eigval_fn=None)

torchzero/modules/second_order/newton.py

@@ -5,7 +5,7 @@ from typing import Literal
 
 import torch
 
-from ...core import Chainable, Module, apply_transform
+from ...core import Chainable, Module, apply_transform, Var
 from ...utils import TensorList, vec_to_tensors
 from ...utils.derivatives import (
     flatten_jacobian,
@@ -50,7 +50,88 @@ def _eigh_solve(H: torch.Tensor, g: torch.Tensor, tfm: Callable | None, search_n
         return None
 
 
+def _get_loss_grad_and_hessian(var: Var, hessian_method:str, vectorize:bool):
+    """returns (loss, g_list, H). Also sets var.loss and var.grad.
+    If hessian_method isn't 'autograd', loss is not set and returned as None"""
+    closure = var.closure
+    if closure is None:
+        raise RuntimeError("Second order methods requires a closure to be provided to the `step` method.")
 
+    params = var.params
+
+    # ------------------------ calculate grad and hessian ------------------------ #
+    loss = None
+    if hessian_method == 'autograd':
+        with torch.enable_grad():
+            loss = var.loss = var.loss_approx = closure(False)
+            g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
+            g_list = [t[0] for t in g_list] # remove leading dim from loss
+            var.grad = g_list
+            H = flatten_jacobian(H_list)
+
+    elif hessian_method in ('func', 'autograd.functional'):
+        strat = 'forward-mode' if vectorize else 'reverse-mode'
+        with torch.enable_grad():
+            g_list = var.get_grad(retain_graph=True)
+            H = hessian_mat(partial(closure, backward=False), params,
+                            method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
+
+    else:
+        raise ValueError(hessian_method)
+
+    return loss, g_list, H
+
+def _newton_step(var: Var, H: torch.Tensor, damping:float, inner: Module | None, H_tfm, eigval_fn, use_lstsq:bool, g_proj: Callable | None = None) -> torch.Tensor:
+    """returns the update tensor, then do vec_to_tensor(update, params)"""
+    params = var.params
+
+    if damping != 0:
+        H = H + torch.eye(H.size(-1), dtype=H.dtype, device=H.device).mul_(damping)
+
+    # -------------------------------- inner step -------------------------------- #
+    update = var.get_update()
+    if inner is not None:
+        update = apply_transform(inner, update, params=params, grads=var.grad, loss=var.loss, var=var)
+
+    g = torch.cat([t.ravel() for t in update])
+    if g_proj is not None: g = g_proj(g)
+
+    # ----------------------------------- solve ---------------------------------- #
+    update = None
+
+    if H_tfm is not None:
+        ret = H_tfm(H, g)
+
+        if isinstance(ret, torch.Tensor):
+            update = ret
+
+        else: # returns (H, is_inv)
+            H, is_inv = ret
+            if is_inv: update = H @ g
+
+    if eigval_fn is not None:
+        update = _eigh_solve(H, g, eigval_fn, search_negative=False)
+
+    if update is None and use_lstsq: update = _least_squares_solve(H, g)
+    if update is None: update = _cholesky_solve(H, g)
+    if update is None: update = _lu_solve(H, g)
+    if update is None: update = _least_squares_solve(H, g)
+
+    return update
+
+def _get_H(H: torch.Tensor, eigval_fn):
+    if eigval_fn is not None:
+        try:
+            L, Q = torch.linalg.eigh(H) # pylint:disable=not-callable
+            L: torch.Tensor = eigval_fn(L)
+            H = Q @ L.diag_embed() @ Q.mH
+            H_inv = Q @ L.reciprocal().diag_embed() @ Q.mH
+            return DenseWithInverse(H, H_inv)
+
+        except torch.linalg.LinAlgError:
+            pass
+
+    return Dense(H)
 
 class Newton(Module):
     """Exact newton's method via autograd.
@@ -81,7 +162,6 @@ class Newton(Module):
             how to calculate hessian. Defaults to "autograd".
         vectorize (bool, optional):
             whether to enable vectorized hessian. Defaults to True.
-        inner (Chainable | None, optional): modules to apply hessian preconditioner to. Defaults to None.
         H_tfm (Callable | None, optional):
             optional hessian transforms, takes in two arguments - `(hessian, gradient)`.
 
@@ -94,6 +174,7 @@ class Newton(Module):
         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.
+        inner (Chainable | None, optional): modules to apply hessian preconditioner to. Defaults to None.
 
     # See also
 
@@ -111,10 +192,9 @@ class Newton(Module):
     The linear system is solved via cholesky decomposition, if that fails, LU decomposition, and if that fails, least squares.
     Least squares can be forced by setting ``use_lstsq=True``, which may generate better search directions when linear system is overdetermined.
 
-    Additionally, if ``eigval_fn`` is specified or ``search_negative`` is ``True``,
-    eigendecomposition of the hessian is computed, ``eigval_fn`` is applied to the eigenvalues,
-    and ``(H + yI)⁻¹`` is computed using the computed eigenvectors and transformed eigenvalues.
-    This is more generally more computationally expensive.
+    Additionally, if ``eigval_fn`` is specified, eigendecomposition of the hessian is computed,
+    ``eigval_fn`` is applied to the eigenvalues, and ``(H + yI)⁻¹`` is computed using the computed eigenvectors and transformed eigenvalues. This is more generally more computationally expensive,
+    but not by much
 
     ## Handling non-convexity
 
@@ -167,16 +247,15 @@ class Newton(Module):
     def __init__(
         self,
         damping: float = 0,
-        search_negative: bool = False,
         use_lstsq: bool = False,
         update_freq: int = 1,
         hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
         vectorize: bool = True,
-        inner: Chainable | None = None,
         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,
+        inner: Chainable | None = None,
     ):
-        defaults = dict(damping=damping, hessian_method=hessian_method, use_lstsq=use_lstsq, vectorize=vectorize, H_tfm=H_tfm, eigval_fn=eigval_fn, search_negative=search_negative, update_freq=update_freq)
+        defaults = dict(damping=damping, hessian_method=hessian_method, use_lstsq=use_lstsq, vectorize=vectorize, H_tfm=H_tfm, eigval_fn=eigval_fn, update_freq=update_freq)
         super().__init__(defaults)
 
         if inner is not None:
@@ -184,200 +263,31 @@ class Newton(Module):
 
     @torch.no_grad
     def update(self, var):
-        params = TensorList(var.params)
-        closure = var.closure
-        if closure is None: raise RuntimeError('NewtonCG requires closure')
-
-        settings = self.settings[params[0]]
-        damping = settings['damping']
-        hessian_method = settings['hessian_method']
-        vectorize = settings['vectorize']
-        update_freq = settings['update_freq']
-
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
-        g_list = var.grad
-        H = None
-        if step % update_freq == 0:
-            # ------------------------ calculate grad and hessian ------------------------ #
-            if hessian_method == 'autograd':
-                with torch.enable_grad():
-                    loss = var.loss = var.loss_approx = closure(False)
-                    g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
-                    g_list = [t[0] for t in g_list] # remove leading dim from loss
-                    var.grad = g_list
-                    H = flatten_jacobian(H_list)
-
-            elif hessian_method in ('func', 'autograd.functional'):
-                strat = 'forward-mode' if vectorize else 'reverse-mode'
-                with torch.enable_grad():
-                    g_list = var.get_grad(retain_graph=True)
-                    H = hessian_mat(partial(closure, backward=False), params,
-                                    method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
-
-            else:
-                raise ValueError(hessian_method)
-
-            if damping != 0: H.add_(torch.eye(H.size(-1), dtype=H.dtype, device=H.device).mul_(damping))
-            self.global_state['H'] = H
+        if step % self.defaults['update_freq'] == 0:
+            loss, g_list, self.global_state['H'] = _get_loss_grad_and_hessian(
+                var, self.defaults['hessian_method'], self.defaults['vectorize']
+            )
 
     @torch.no_grad
     def apply(self, var):
-        H = self.global_state["H"]
-
         params = var.params
-        settings = self.settings[params[0]]
-        search_negative = settings['search_negative']
-        H_tfm = settings['H_tfm']
-        eigval_fn = settings['eigval_fn']
-        use_lstsq = settings['use_lstsq']
-
-        # -------------------------------- inner step -------------------------------- #
-        update = var.get_update()
-        if 'inner' in self.children:
-            update = apply_transform(self.children['inner'], update, params=params, grads=var.grad, var=var)
-
-        g = torch.cat([t.ravel() for t in update])
-
-        # ----------------------------------- solve ---------------------------------- #
-        update = None
-        if H_tfm is not None:
-            ret = H_tfm(H, g)
-
-            if isinstance(ret, torch.Tensor):
-                update = ret
-
-            else: # returns (H, is_inv)
-                H, is_inv = ret
-                if is_inv: update = H @ g
-
-        if search_negative or (eigval_fn is not None):
-            update = _eigh_solve(H, g, eigval_fn, search_negative=search_negative)
-
-        if update is None and use_lstsq: update = _least_squares_solve(H, g)
-        if update is None: update = _cholesky_solve(H, g)
-        if update is None: update = _lu_solve(H, g)
-        if update is None: update = _least_squares_solve(H, g)
+        update = _newton_step(
+            var=var,
+            H = self.global_state["H"],
+            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"],
+        )
 
         var.update = vec_to_tensors(update, params)
 
         return var
 
-    def get_H(self,var):
-        H = self.global_state["H"]
-        settings = self.defaults
-        if settings['eigval_fn'] is not None:
-            try:
-                L, Q = torch.linalg.eigh(H) # pylint:disable=not-callable
-                L = settings['eigval_fn'](L)
-                H = Q @ L.diag_embed() @ Q.mH
-                H_inv = Q @ L.reciprocal().diag_embed() @ Q.mH
-                return DenseWithInverse(H, H_inv)
-
-            except torch.linalg.LinAlgError:
-                pass
-
-        return Dense(H)
-
-
-class InverseFreeNewton(Module):
-    """Inverse-free newton's method
-
-    .. note::
-        In most cases Newton 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 the hessian.
-        The closure must accept a ``backward`` argument (refer to documentation).
-
-    .. warning::
-        this uses roughly O(N^2) memory.
-
-    Reference
-        Massalski, Marcin, and Magdalena Nockowska-Rosiak. "INVERSE-FREE NEWTON'S METHOD." Journal of Applied Analysis & Computation 15.4 (2025): 2238-2257.
-    """
-    def __init__(
-        self,
-        update_freq: int = 1,
-        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
-        vectorize: bool = True,
-        inner: Chainable | None = None,
-    ):
-        defaults = dict(hessian_method=hessian_method, vectorize=vectorize, update_freq=update_freq)
-        super().__init__(defaults)
-
-        if inner is not None:
-            self.set_child('inner', inner)
-
-    @torch.no_grad
-    def update(self, var):
-        params = TensorList(var.params)
-        closure = var.closure
-        if closure is None: raise RuntimeError('NewtonCG requires closure')
-
-        settings = self.settings[params[0]]
-        hessian_method = settings['hessian_method']
-        vectorize = settings['vectorize']
-        update_freq = settings['update_freq']
-
-        step = self.global_state.get('step', 0)
-        self.global_state['step'] = step + 1
-
-        g_list = var.grad
-        Y = None
-        if step % update_freq == 0:
-            # ------------------------ calculate grad and hessian ------------------------ #
-            if hessian_method == 'autograd':
-                with torch.enable_grad():
-                    loss = var.loss = var.loss_approx = closure(False)
-                    g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
-                    g_list = [t[0] for t in g_list] # remove leading dim from loss
-                    var.grad = g_list
-                    H = flatten_jacobian(H_list)
-
-            elif hessian_method in ('func', 'autograd.functional'):
-                strat = 'forward-mode' if vectorize else 'reverse-mode'
-                with torch.enable_grad():
-                    g_list = var.get_grad(retain_graph=True)
-                    H = hessian_mat(partial(closure, backward=False), params,
-                                    method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
-
-            else:
-                raise ValueError(hessian_method)
-
-            self.global_state["H"] = H
-
-            # inverse free part
-            if 'Y' not in self.global_state:
-                num = H.T
-                denom = (torch.linalg.norm(H, 1) * torch.linalg.norm(H, float('inf'))) # pylint:disable=not-callable
-                finfo = torch.finfo(H.dtype)
-                Y = self.global_state['Y'] = num.div_(denom.clip(min=finfo.tiny * 2, max=finfo.max / 2))
-
-            else:
-                Y = self.global_state['Y']
-                I = torch.eye(Y.size(0), device=Y.device, dtype=Y.dtype).mul_(2)
-                I -= H @ Y
-                Y = self.global_state['Y'] = Y @ I
-
-
-    def apply(self, var):
-        Y = self.global_state["Y"]
-        params = var.params
-
-        # -------------------------------- inner step -------------------------------- #
-        update = var.get_update()
-        if 'inner' in self.children:
-            update = apply_transform(self.children['inner'], update, params=params, grads=var.grad, var=var)
-
-        g = torch.cat([t.ravel() for t in update])
-
-        # ----------------------------------- solve ---------------------------------- #
-        var.update = vec_to_tensors(Y@g, params)
-
-        return var
+    def get_H(self,var=...):
+        return _get_H(self.global_state["H"], self.defaults["eigval_fn"])
 
-    def get_H(self,var):
-        return DenseWithInverse(A = self.global_state["H"], A_inv=self.global_state["Y"])