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

torchzero/modules/higher_order/higher_order_newton.py

@@ -70,57 +70,99 @@ 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.
 
-    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.
+    .. note::
+        In most cases HigherOrderNewton should be the first module in the chain because it relies on extra 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 higher order derivatives.
+        The closure must accept a ``backward`` argument (refer to documentation).
+
+    .. warning::
+        this uses roughly O(N^order) memory and solving the subproblem can be very expensive.
+
+    .. warning::
+        "none" and "proximal" trust methods may generate subproblems that have no minima, causing divergence.
 
     Args:
 
@@ -149,39 +191,38 @@ 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 = 2,
+        nminus: float = 0.25,
+        init: float | None = None,
+        eta: float = 1e-6,
+        max_attempts = 10,
         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)
         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)
-
-
         # ------------------------ calculate grad and hessian ------------------------ #
         with torch.enable_grad():
             loss = var.loss = var.loss_approx = closure(False)
@@ -205,52 +246,74 @@ class HigherOrderNewton(Module):
 
         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)
+            if trust_value < 1e-8 or trust_value > 1e16: trust_value = self.global_state['trust_region'] = settings['init']
+
+            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 < 0.25:
+                    self.global_state['trust_region'] = trust_value * nminus
+
+                # very good step
+                elif rho > 0.75:
+                    diff = trust_value - (x0 - x_star).abs_()
+                    if (diff.amin() / trust_value) > 1e-4: # hits 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/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 AdaptiveLineSearch
+from .backtracking import AdaptiveBacktracking, Backtracking
+from .line_search import LineSearchBase
 from .scipy import ScipyMinimizeScalar
-from .trust_region import TrustRegion
+from .strong_wolfe import StrongWolfe

torchzero/modules/line_search/adaptive.py

@@ -0,0 +1,99 @@
+import math
+from collections.abc import Callable
+from operator import itemgetter
+
+import torch
+
+from .line_search import LineSearchBase
+
+
+
+def adaptive_tracking(
+    f,
+    x_0,
+    maxiter: int,
+    nplus: float = 2,
+    nminus: float = 0.5,
+):
+    f_0 = f(0)
+
+    t = x_0
+    f_t = f(t)
+
+    # backtrack
+    if f_t > f_0:
+        while f_t > f_0:
+            maxiter -= 1
+            if maxiter < 0: return 0, f_0
+            t = t*nminus
+            f_t = f(t)
+        return t, f_t
+
+    # forwardtrack
+    f_prev = f_t
+    t *= nplus
+    f_t = f(t)
+    if f_prev < f_t: return t / nplus, f_prev
+    while f_prev >= f_t:
+        maxiter -= 1
+        if maxiter < 0: return t, f_t
+        f_prev = f_t
+        t *= nplus
+        f_t = f(t)
+    return t / nplus, f_prev
+
+class AdaptiveLineSearch(LineSearchBase):
+    """Adaptive line search, similar to backtracking but also has forward tracking mode.
+    Currently doesn't check for weak curvature condition.
+
+    Args:
+        init (float, optional): initial step size. Defaults to 1.0.
+        beta (float, optional): multiplies each consecutive step size by this value. Defaults to 0.5.
+        maxiter (int, optional): Maximum line search function evaluations. Defaults to 10.
+        adaptive (bool, optional):
+            when enabled, if line search failed, beta size is reduced.
+            Otherwise it is reset to initial value. Defaults to True.
+    """
+    def __init__(
+        self,
+        init: float = 1.0,
+        nplus: float = 2,
+        nminus: float = 0.5,
+        maxiter: int = 10,
+        adaptive=True,
+    ):
+        defaults=dict(init=init,nplus=nplus,nminus=nminus,maxiter=maxiter,adaptive=adaptive,)
+        super().__init__(defaults=defaults)
+        self.global_state['beta_scale'] = 1.0
+
+    def reset(self):
+        super().reset()
+        self.global_state['beta_scale'] = 1.0
+
+    @torch.no_grad
+    def search(self, update, var):
+        init, nplus, nminus, maxiter, adaptive = itemgetter(
+            'init', 'nplus', 'nminus', 'maxiter', 'adaptive')(self.settings[var.params[0]])
+
+        objective = self.make_objective(var=var)
+
+        # # directional derivative
+        # d = -sum(t.sum() for t in torch._foreach_mul(var.get_grad(), var.get_update()))
+
+        # scale beta (beta is multiplicative and i think may be better than scaling initial step size)
+        beta_scale = self.global_state.get('beta_scale', 1)
+        x_prev = self.global_state.get('prev_x', 1)
+
+        if adaptive: nminus = nminus * beta_scale
+
+
+        step_size, f = adaptive_tracking(objective, x_prev, maxiter, nplus=nplus, nminus=nminus)
+
+        # found an alpha that reduces loss
+        if step_size != 0:
+            self.global_state['beta_scale'] = min(1.0, self.global_state['beta_scale'] * math.sqrt(1.5))
+            return step_size
+
+        # on fail reduce beta scale value
+        self.global_state['beta_scale'] /= 1.5
+        return 0

torchzero/modules/line_search/backtracking.py

@@ -4,7 +4,7 @@ from operator import itemgetter
 
 import torch
 
-from .line_search import LineSearch
+from .line_search import LineSearchBase
 
 
 def backtracking_line_search(
@@ -19,12 +19,12 @@ def backtracking_line_search(
     """
 
     Args:
-        objective_fn: evaluates step size along some descent direction.
-        dir_derivative: directional derivative along the descent direction.
-        alpha_init: initial step size.
+        f: evaluates step size along some descent direction.
+        g_0: directional derivative along the descent direction.
+        init: initial step size.
         beta: The factor by which to decrease alpha in each iteration
         c: The constant for the Armijo sufficient decrease condition
-        max_iter: Maximum number of backtracking iterations (default: 10).
+        maxiter: Maximum number of backtracking iterations (default: 10).
 
     Returns:
         step size
@@ -32,11 +32,15 @@ def backtracking_line_search(
 
     a = init
     f_x = f(0)
+    f_prev = None
 
     for iteration in range(maxiter):
         f_a = f(a)
 
-        if f_a <= f_x + c * a * min(g_0, 0): # pyright: ignore[reportArgumentType]
+        if (f_prev is not None) and (f_a > f_prev) and (f_prev < f_x): return a / beta
+        f_prev = f_a
+
+        if f_a < f_x + c * a * min(g_0, 0): # pyright: ignore[reportArgumentType]
             # found an acceptable alpha
             return a
 
@@ -59,7 +63,7 @@ def backtracking_line_search(
 
     return None
 
-class Backtracking(LineSearch):
+class Backtracking(LineSearchBase):
     """Backtracking line search satisfying the Armijo condition.
 
     Args:
@@ -68,9 +72,30 @@ class Backtracking(LineSearch):
         c (float, optional): acceptance value for Armijo condition. Defaults to 1e-4.
         maxiter (int, optional): Maximum line search function evaluations. Defaults to 10.
         adaptive (bool, optional):
-            when enabled, if line search failed, initial step size is reduced.
+            when enabled, if line search failed, beta is reduced.
             Otherwise it is reset to initial value. Defaults to True.
         try_negative (bool, optional): Whether to perform line search in opposite direction on fail. Defaults to False.
+
+    Examples:
+        Gradient descent with backtracking line search:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Backtracking()
+            )
+
+        LBFGS with backtracking line search:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.LBFGS(),
+                tz.m.Backtracking()
+            )
+
     """
     def __init__(
         self,
@@ -117,7 +142,7 @@ class Backtracking(LineSearch):
 def _lerp(start,end,weight):
     return start + weight * (end - start)
 
-class AdaptiveBacktracking(LineSearch):
+class AdaptiveBacktracking(LineSearchBase):
     """Adaptive backtracking line search. After each line search procedure, a new initial step size is set
     such that optimal step size in the procedure would be found on the second line search iteration.
 

torchzero/modules/line_search/line_search.py

@@ -15,8 +15,9 @@ from ...utils import tofloat
 class MaxLineSearchItersReached(Exception): pass
 
 
-class LineSearch(Module, ABC):
+class LineSearchBase(Module, ABC):
     """Base class for line searches.
+
     This is an abstract class, to use it, subclass it and override `search`.
 
     Args:
@@ -26,6 +27,62 @@ class LineSearch(Module, ABC):
             the objective this many times, and step size with the lowest loss value will be used.
             This is useful when passing `make_objective` to an external library which
             doesn't have a maxiter option. Defaults to None.
+
+    Other useful methods:
+        * `evaluate_step_size` - returns loss with a given scalar step size
+        * `evaluate_step_size_loss_and_derivative` - returns loss and directional derivative with a given scalar step size
+        * `make_objective` - creates a function that accepts a scalar step size and returns loss. This can be passed to a scalar solver, such as scipy.optimize.minimize_scalar.
+        * `make_objective_with_derivative` - creates a function that accepts a scalar step size and returns a tuple with loss and directional derivative. This can be passed to a scalar solver.
+
+    Examples:
+        #### Basic line search
+
+        This evaluates all step sizes in a range by using the :code:`self.evaluate_step_size` method.
+
+        .. code-block:: python
+
+            class GridLineSearch(LineSearch):
+                def __init__(self, start, end, num):
+                    defaults = dict(start=start,end=end,num=num)
+                    super().__init__(defaults)
+
+                @torch.no_grad
+                def search(self, update, var):
+                    settings = self.settings[var.params[0]]
+                    start = settings["start"]
+                    end = settings["end"]
+                    num = settings["num"]
+
+                    lowest_loss = float("inf")
+                    best_step_size = best_step_size
+
+                    for step_size in torch.linspace(start,end,num):
+                        loss = self.evaluate_step_size(step_size.item(), var=var, backward=False)
+                        if loss < lowest_loss:
+                            lowest_loss = loss
+                            best_step_size = step_size
+
+                    return best_step_size
+
+        #### Using external solver via self.make_objective
+
+        Here we let :code:`scipy.optimize.minimize_scalar` solver find the best step size via :code:`self.make_objective`
+
+        .. code-block:: python
+
+            class ScipyMinimizeScalar(LineSearch):
+                def __init__(self, method: str | None = None):
+                    defaults = dict(method=method)
+                    super().__init__(defaults)
+
+                @torch.no_grad
+                def search(self, update, var):
+                    objective = self.make_objective(var=var)
+                    method = self.settings[var.params[0]]["method"]
+
+                    res = self.scopt.minimize_scalar(objective, method=method)
+                    return res.x
+
     """
     def __init__(self, defaults: dict[str, Any] | None, maxiter: int | None = None):
         super().__init__(defaults)
@@ -165,17 +222,18 @@ class LineSearch(Module, ABC):
         return var
 
 
-class GridLineSearch(LineSearch):
-    """Mostly for testing, this is not practical"""
-    def __init__(self, start, end, num):
-        defaults = dict(start=start,end=end,num=num)
-        super().__init__(defaults)
 
-    @torch.no_grad
-    def search(self, update, var):
-        start,end,num=itemgetter('start','end','num')(self.settings[var.params[0]])
+# class GridLineSearch(LineSearch):
+#     """Mostly for testing, this is not practical"""
+#     def __init__(self, start, end, num):
+#         defaults = dict(start=start,end=end,num=num)
+#         super().__init__(defaults)
+
+#     @torch.no_grad
+#     def search(self, update, var):
+#         start,end,num=itemgetter('start','end','num')(self.settings[var.params[0]])
 
-        for lr in torch.linspace(start,end,num):
-            self.evaluate_step_size(lr.item(), var=var, backward=False)
+#         for lr in torch.linspace(start,end,num):
+#             self.evaluate_step_size(lr.item(), var=var, backward=False)
 
-        return self._best_step_size
+#         return self._best_step_size