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

torchzero/modules/higher_order/higher_order_newton.py

@@ -13,7 +13,7 @@ import torch
 from ...core import Chainable, Module, apply_transform
 from ...utils import TensorList, vec_to_tensors, vec_to_tensors_
 from ...utils.derivatives import (
-    hessian_list_to_mat,
+    flatten_jacobian,
     jacobian_wrt,
 )
 
@@ -70,57 +70,94 @@ def _proximal_poly_H(x: np.ndarray, c, prox, x0: np.ndarray, derivatives):
 def _poly_minimize(trust_region, prox, de_iters: Any, c, x: torch.Tensor, derivatives):
     derivatives = [T.detach().cpu().numpy().astype(np.float64) for T in derivatives]
     x0 = x.detach().cpu().numpy().astype(np.float64) # taylor series center
-    bounds = None
-    if trust_region is not None: bounds = list(zip(x0 - trust_region, x0 + trust_region))
 
-    # if len(derivatives) is 1, only gradient is available, I use that to test proximal penalty and bounds
-    if bounds is None:
-        if len(derivatives) == 1: method = 'bfgs'
-        else: method = 'trust-exact'
+    # notes
+    # 1. since we have exact hessian we use trust methods
+
+    # 2. if len(derivatives) is 1, only gradient is available,
+    # thus use slsqp depending on whether trust region is enabled
+    # this is just so that I can test that trust region works
+    if trust_region is None:
+        if len(derivatives) == 1: raise RuntimeError("trust region must be enabled because 1st order has no minima")
+        method = 'trust-exact'
+        de_bounds = list(zip(x0 - 10, x0 + 10))
+        constraints = None
+
     else:
-        if len(derivatives) == 1: method = 'l-bfgs-b'
+        if len(derivatives) == 1: method = 'slsqp'
         else: method = 'trust-constr'
+        de_bounds = list(zip(x0 - trust_region, x0 + trust_region))
+
+        def l2_bound_f(x):
+            if x.ndim == 2: return np.sum((x - x0[:,None])**2, axis=0)[None,:] # DE passes (ndim, batch_size) and expects (M, S)
+            return np.sum((x - x0)**2, axis=0)
+
+        def l2_bound_g(x):
+            return 2 * (x - x0)
+
+        def l2_bound_h(x, v):
+            return v[0] * 2 * np.eye(x0.shape[0])
+
+        constraint = scipy.optimize.NonlinearConstraint(
+            fun=l2_bound_f,
+            lb=0, # 0 <= ||x-x0||^2
+            ub=trust_region**2, # ||x-x0||^2 <= R^2
+            jac=l2_bound_g, # pyright:ignore[reportArgumentType]
+            hess=l2_bound_h,
+            keep_feasible=False
+        )
+        constraints = [constraint]
 
     x_init = x0.copy()
     v0 = _proximal_poly_v(x0, c, prox, x0, derivatives)
+
+    # ---------------------------------- run DE ---------------------------------- #
     if de_iters is not None and de_iters != 0:
         if de_iters == -1: de_iters = None # let scipy decide
+
+        # DE needs bounds so use linf ig
         res = scipy.optimize.differential_evolution(
             _proximal_poly_v,
-            bounds if bounds is not None else list(zip(x0 - 10, x0 + 10)),
+            de_bounds,
             args=(c, prox, x0.copy(), derivatives),
             maxiter=de_iters,
             vectorized=True,
+            constraints = constraints,
+            updating='deferred',
         )
-        if res.fun < v0: x_init = res.x
-
-    res = scipy.optimize.minimize(
-        _proximal_poly_v,
-        x_init,
-        method=method,
-        args=(c, prox, x0.copy(), derivatives),
-        jac=_proximal_poly_g,
-        hess=_proximal_poly_H,
-        bounds=bounds
-    )
+        if res.fun < v0 and np.all(np.isfinite(res.x)): x_init = res.x
 
+    # ------------------------------- run minimize ------------------------------- #
+    try:
+        res = scipy.optimize.minimize(
+            _proximal_poly_v,
+            x_init,
+            method=method,
+            args=(c, prox, x0.copy(), derivatives),
+            jac=_proximal_poly_g,
+            hess=_proximal_poly_H,
+            constraints = constraints,
+        )
+    except ValueError:
+        return x, -float('inf')
     return torch.from_numpy(res.x).to(x), res.fun
 
 
 
 class HigherOrderNewton(Module):
-    """
-    A basic arbitrary order newton's method with optional trust region and proximal penalty.
-    It is recommended to enable at least one of trust region or proximal penalty.
+    """A basic arbitrary order newton's method with optional trust region and proximal penalty.
 
     This constructs an nth order taylor approximation via autograd and minimizes it with
-    scipy.optimize.minimize trust region newton solvers with optional proximal penalty.
+    ``scipy.optimize.minimize`` trust region newton solvers with optional proximal penalty.
 
-    This uses n^order memory, where n is number of decision variables, and I am not aware
-    of any problems where this is more efficient than newton's method. It can minimize
-    rosenbrock in a single step, but that step probably takes more time than newton.
-    And there are way more efficient tensor methods out there but they tend to be
-    significantly more complex.
+    The hessian of taylor approximation is easier to evaluate, plus it can be evaluated in a batched mode,
+    so it can be more efficient in very specific instances.
+
+    Notes:
+        - In most cases HigherOrderNewton should be the first module in the chain because it relies on extra autograd. Use the ``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 higher order derivatives. The closure must accept a ``backward`` argument (refer to documentation).
+        - this uses roughly O(N^order) memory and solving the subproblem is very expensive.
+        - "none" and "proximal" trust methods may generate subproblems that have no minima, causing divergence.
 
     Args:
 
@@ -136,7 +173,7 @@ class HigherOrderNewton(Module):
         increase (float, optional): trust region multiplier on good steps. Defaults to 1.5.
         decrease (float, optional): trust region multiplier on bad steps. Defaults to 0.75.
         trust_init (float | None, optional):
-            initial trust region size. If none, defaults to 1 on :code:`trust_method="bounds"` and 0.1 on :code:`"proximal"`. Defaults to None.
+            initial trust region size. If none, defaults to 1 on :code:`trust_method="bounds"` and 0.1 on ``"proximal"``. Defaults to None.
         trust_tol (float, optional):
             Maximum ratio of expected loss reduction to actual reduction for trust region increase.
             Should 1 or higer. Defaults to 2.
@@ -149,38 +186,43 @@ class HigherOrderNewton(Module):
         self,
         order: int = 4,
         trust_method: Literal['bounds', 'proximal', 'none'] | None = 'bounds',
-        increase: float = 1.5,
-        decrease: float = 0.75,
-        trust_init: float | None = None,
-        trust_tol: float = 2,
+        nplus: float = 3.5,
+        nminus: float = 0.25,
+        rho_good: float = 0.99,
+        rho_bad: float = 1e-4,
+        init: float | None = None,
+        eta: float = 1e-6,
+        max_attempts = 10,
+        boundary_tol: float = 1e-2,
         de_iters: int | None = None,
         vectorize: bool = True,
     ):
-        if trust_init is None:
-            if trust_method == 'bounds': trust_init = 1
-            else: trust_init = 0.1
+        if init is None:
+            if trust_method == 'bounds': init = 1
+            else: init = 0.1
 
-        defaults = dict(order=order, trust_method=trust_method, increase=increase, decrease=decrease, trust_tol=trust_tol, trust_init=trust_init, vectorize=vectorize, de_iters=de_iters)
+        defaults = dict(order=order, trust_method=trust_method, nplus=nplus, nminus=nminus, eta=eta, init=init, vectorize=vectorize, de_iters=de_iters, max_attempts=max_attempts, boundary_tol=boundary_tol, rho_good=rho_good, rho_bad=rho_bad)
         super().__init__(defaults)
 
     @torch.no_grad
     def step(self, var):
         params = TensorList(var.params)
         closure = var.closure
-        if closure is None: raise RuntimeError('NewtonCG requires closure')
+        if closure is None: raise RuntimeError('HigherOrderNewton requires closure')
 
         settings = self.settings[params[0]]
         order = settings['order']
-        increase = settings['increase']
-        decrease = settings['decrease']
-        trust_tol = settings['trust_tol']
-        trust_init = settings['trust_init']
+        nplus = settings['nplus']
+        nminus = settings['nminus']
+        eta = settings['eta']
+        init = settings['init']
         trust_method = settings['trust_method']
         de_iters = settings['de_iters']
+        max_attempts = settings['max_attempts']
         vectorize = settings['vectorize']
-
-        trust_value = self.global_state.get('trust_value', trust_init)
-
+        boundary_tol = settings['boundary_tol']
+        rho_good = settings['rho_good']
+        rho_bad = settings['rho_bad']
 
         # ------------------------ calculate grad and hessian ------------------------ #
         with torch.enable_grad():
@@ -200,57 +242,86 @@ class HigherOrderNewton(Module):
                 T_list = jacobian_wrt([T], params, create_graph=not is_last, batched=vectorize)
                 with torch.no_grad() if is_last else nullcontext():
                     # the shape is (ndim, ) * order
-                    T = hessian_list_to_mat(T_list).view(n, n, *T.shape[1:])
+                    T = flatten_jacobian(T_list).view(n, n, *T.shape[1:])
                     derivatives.append(T)
 
         x0 = torch.cat([p.ravel() for p in params])
 
-        if trust_method is None: trust_method = 'none'
-        else: trust_method = trust_method.lower()
-
-        if trust_method == 'none':
-            trust_region = None
-            prox = 0
-
-        elif trust_method == 'bounds':
-            trust_region = trust_value
-            prox = 0
-
-        elif trust_method == 'proximal':
-            trust_region = None
-            prox = 1 / trust_value
-
+        success = False
+        x_star = None
+        while not success:
+            max_attempts -= 1
+            if max_attempts < 0: break
+
+            # load trust region value
+            trust_value = self.global_state.get('trust_region', init)
+
+            # make sure its not too small or too large
+            finfo = torch.finfo(x0.dtype)
+            if trust_value < finfo.tiny*2 or trust_value > finfo.max / (2*nplus):
+                trust_value = self.global_state['trust_region'] = settings['init']
+
+            # determine tr and prox values
+            if trust_method is None: trust_method = 'none'
+            else: trust_method = trust_method.lower()
+
+            if trust_method == 'none':
+                trust_region = None
+                prox = 0
+
+            elif trust_method == 'bounds':
+                trust_region = trust_value
+                prox = 0
+
+            elif trust_method == 'proximal':
+                trust_region = None
+                prox = 1 / trust_value
+
+            else:
+                raise ValueError(trust_method)
+
+            # minimize the model
+            x_star, expected_loss = _poly_minimize(
+                trust_region=trust_region,
+                prox=prox,
+                de_iters=de_iters,
+                c=loss.item(),
+                x=x0,
+                derivatives=derivatives,
+            )
+
+            # update trust region
+            if trust_method == 'none':
+                success = True
+            else:
+                pred_reduction = loss - expected_loss
+
+                vec_to_tensors_(x_star, params)
+                loss_star = closure(False)
+                vec_to_tensors_(x0, params)
+                reduction = loss - loss_star
+
+                rho = reduction / (max(pred_reduction, 1e-8))
+                # failed step
+                if rho < rho_bad:
+                    self.global_state['trust_region'] = trust_value * nminus
+
+                # very good step
+                elif rho > rho_good:
+                    step = (x_star - x0)
+                    magn = torch.linalg.vector_norm(step) # pylint:disable=not-callable
+                    if trust_method == 'proximal' or (trust_value - magn) / trust_value <= boundary_tol:
+                        # close to boundary
+                        self.global_state['trust_region'] = trust_value * nplus
+
+                # if the ratio is high enough then accept the proposed step
+                success = rho > eta
+
+        assert x_star is not None
+        if success:
+            difference = vec_to_tensors(x0 - x_star, params)
+            var.update = list(difference)
         else:
-            raise ValueError(trust_method)
-
-        x_star, expected_loss = _poly_minimize(
-            trust_region=trust_region,
-            prox=prox,
-            de_iters=de_iters,
-            c=loss.item(),
-            x=x0,
-            derivatives=derivatives,
-        )
-
-        # trust region
-        if trust_method != 'none':
-            expected_reduction = loss - expected_loss
-
-            vec_to_tensors_(x_star, params)
-            loss_star = closure(False)
-            vec_to_tensors_(x0, params)
-            reduction = loss - loss_star
-
-            # failed step
-            if reduction <= 0:
-                x_star = x0
-                self.global_state['trust_value'] = trust_value * decrease
-
-            # very good step
-            elif expected_reduction / reduction <= trust_tol:
-                self.global_state['trust_value'] = trust_value * increase
-
-        difference = vec_to_tensors(x0 - x_star, params)
-        var.update = list(difference)
+            var.update = params.zeros_like()
         return var
 

torchzero/modules/least_squares/__init__.py

@@ -0,0 +1 @@
+from .gn import SumOfSquares, GaussNewton

torchzero/modules/least_squares/gn.py

@@ -0,0 +1,161 @@
+import torch
+from ...core import Module
+
+from ...utils.derivatives import jacobian_wrt, flatten_jacobian
+from ...utils import vec_to_tensors
+from ...utils.linalg import linear_operator
+class SumOfSquares(Module):
+    """Sets loss to be the sum of squares of values returned by the closure.
+
+    This is meant to be used to test least squares methods against ordinary minimization methods.
+
+    To use this, the closure should return a vector of values to minimize sum of squares of.
+    Please add the `backward` argument, it will always be False but it is required.
+    """
+    def __init__(self):
+        super().__init__()
+
+    @torch.no_grad
+    def step(self, var):
+        closure = var.closure
+
+        if closure is not None:
+            def sos_closure(backward=True):
+                if backward:
+                    var.zero_grad()
+                    with torch.enable_grad():
+                        loss = closure(False)
+                        loss = loss.pow(2).sum()
+                        loss.backward()
+                    return loss
+
+                loss = closure(False)
+                return loss.pow(2).sum()
+
+            var.closure = sos_closure
+
+        if var.loss is not None:
+            var.loss = var.loss.pow(2).sum()
+
+        if var.loss_approx is not None:
+            var.loss_approx = var.loss_approx.pow(2).sum()
+
+        return var
+
+
+class GaussNewton(Module):
+    """Gauss-newton method.
+
+    To use this, the closure should return a vector of values to minimize sum of squares of.
+    Please add the ``backward`` argument, it will always be False but it is required.
+    Gradients will be calculated via batched autograd within this module, you don't need to
+    implement the backward pass. Please see below for an example.
+
+    Note:
+        This method requires ``ndim^2`` memory, however, if it is used within ``tz.m.TrustCG`` trust region,
+        the memory requirement is ``ndim*m``, where ``m`` is number of values in the output.
+
+    Args:
+        reg (float, optional): regularization parameter. Defaults to 1e-8.
+        batched (bool, optional): whether to use vmapping. Defaults to True.
+
+    Examples:
+
+    minimizing the rosenbrock function:
+    ```python
+    def rosenbrock(X):
+        x1, x2 = X
+        return torch.stack([(1 - x1), 100 * (x2 - x1**2)])
+
+    X = torch.tensor([-1.1, 2.5], requires_grad=True)
+    opt = tz.Modular([X], tz.m.GaussNewton(), tz.m.Backtracking())
+
+    # define the closure for line search
+    def closure(backward=True):
+        return rosenbrock(X)
+
+    # minimize
+    for iter in range(10):
+        loss = opt.step(closure)
+        print(f'{loss = }')
+    ```
+
+    training a neural network with a matrix-free GN trust region:
+    ```python
+    X = torch.randn(64, 20)
+    y = torch.randn(64, 10)
+
+    model = nn.Sequential(nn.Linear(20, 64), nn.ELU(), nn.Linear(64, 10))
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.TrustCG(tz.m.GaussNewton()),
+    )
+
+    def closure(backward=True):
+        y_hat = model(X) # (64, 10)
+        return (y_hat - y).pow(2).mean(0) # (10, )
+
+    for i in range(100):
+        losses = opt.step(closure)
+        if i % 10 == 0:
+            print(f'{losses.mean() = }')
+    ```
+    """
+    def __init__(self, reg:float = 1e-8, batched:bool=True, ):
+        super().__init__(defaults=dict(batched=batched, reg=reg))
+
+    @torch.no_grad
+    def update(self, var):
+        params = var.params
+        batched = self.defaults['batched']
+
+        closure = var.closure
+        assert closure is not None
+
+        # gauss newton direction
+        with torch.enable_grad():
+            f = var.get_loss(backward=False) # n_out
+            assert isinstance(f, torch.Tensor)
+            G_list = jacobian_wrt([f.ravel()], params, batched=batched)
+
+        var.loss = f.pow(2).sum()
+
+        G = self.global_state["G"] = flatten_jacobian(G_list) # (n_out, ndim)
+        Gtf = G.T @ f.detach() # (ndim)
+        self.global_state["Gtf"] = Gtf
+        var.grad = vec_to_tensors(Gtf, var.params)
+
+        # set closure to calculate sum of squares for line searches etc
+        if var.closure is not None:
+            def sos_closure(backward=True):
+                if backward:
+                    var.zero_grad()
+                    with torch.enable_grad():
+                        loss = closure(False).pow(2).sum()
+                        loss.backward()
+                    return loss
+
+                loss = closure(False).pow(2).sum()
+                return loss
+
+            var.closure = sos_closure
+
+    @torch.no_grad
+    def apply(self, var):
+        reg = self.defaults['reg']
+
+        G = self.global_state['G']
+        Gtf = self.global_state['Gtf']
+
+        GtG = G.T @ G # (ndim, ndim)
+        if reg != 0:
+            GtG.add_(torch.eye(GtG.size(0), device=GtG.device, dtype=GtG.dtype).mul_(reg))
+
+        v = torch.linalg.lstsq(GtG, Gtf).solution # pylint:disable=not-callable
+
+        var.update = vec_to_tensors(v, var.params)
+        return var
+
+    def get_H(self, var):
+        G = self.global_state['G']
+        return linear_operator.AtA(G)

torchzero/modules/line_search/__init__.py

@@ -1,5 +1,5 @@
-from .line_search import LineSearch, GridLineSearch
-from .backtracking import backtracking_line_search, Backtracking, AdaptiveBacktracking
-from .strong_wolfe import StrongWolfe
+from .adaptive import AdaptiveTracking
+from .backtracking import AdaptiveBacktracking, Backtracking
+from .line_search import LineSearchBase
 from .scipy import ScipyMinimizeScalar
-from .trust_region import TrustRegion
+from .strong_wolfe import StrongWolfe