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

torchzero/modules/quasi_newton/quasi_newton.py

@@ -5,10 +5,10 @@ from typing import Any, Literal
 
 import torch
 
-from ...core import Chainable, Module, TensorwiseTransform, Transform
+from ...core import Chainable, Module, TensorTransform, Transform
 from ...utils import TensorList, set_storage_, unpack_states, safe_dict_update_
-from ...utils.linalg import linear_operator
-from ..functional import initial_step_size, safe_clip
+from ...linalg import linear_operator
+from ..opt_utils import initial_step_size, safe_clip
 
 
 
@@ -17,7 +17,7 @@ def _maybe_lerp_(state, key, value: torch.Tensor, beta: float | None):
     elif state[key].shape != value.shape: state[key] = value
     else: state[key].lerp_(value, 1-beta)
 
-class HessianUpdateStrategy(TensorwiseTransform, ABC):
+class HessianUpdateStrategy(TensorTransform, ABC):
     """Base class for quasi-newton methods that store and update hessian approximation H or inverse B.
 
     This is an abstract class, to use it, subclass it and override ``update_H`` and/or ``update_B``,
@@ -106,11 +106,12 @@ class HessianUpdateStrategy(TensorwiseTransform, ABC):
         scale_first: bool = False,
         concat_params: bool = True,
         inverse: bool = True,
+        uses_loss: bool = False,
         inner: Chainable | None = None,
     ):
         if defaults is None: defaults = {}
         safe_dict_update_(defaults, dict(init_scale=init_scale, tol=tol, ptol=ptol, ptol_restart=ptol_restart, gtol=gtol, inverse=inverse, beta=beta, restart_interval=restart_interval, scale_first=scale_first))
-        super().__init__(defaults, uses_grad=False, concat_params=concat_params, update_freq=update_freq, inner=inner)
+        super().__init__(defaults, uses_loss=uses_loss, concat_params=concat_params, update_freq=update_freq, inner=inner)
 
     def reset_for_online(self):
         super().reset_for_online()
@@ -141,23 +142,27 @@ class HessianUpdateStrategy(TensorwiseTransform, ABC):
         return H
 
     # ------------------------------ common methods ------------------------------ #
-    def auto_initial_scale(self, s:torch.Tensor,y:torch.Tensor) -> torch.Tensor | float:
+    def auto_initial_scale(self, s:torch.Tensor,y:torch.Tensor) -> torch.Tensor | float | None:
         """returns multiplier to B on 2nd step if ``init_scale='auto'``. H should be divided by this!"""
         ys = y.dot(s)
         yy = y.dot(y)
-        if ys != 0 and yy != 0: return yy/ys
-        return 1
+        tiny = torch.finfo(ys.dtype).tiny * 2
+        if ys > tiny and yy > tiny: return yy/ys
+        return None
 
-    def reset_P(self, P: torch.Tensor, s:torch.Tensor,y:torch.Tensor, inverse:bool, init_scale: Any, state:dict[str,Any]) -> None:
+    def reset_P(self, P: torch.Tensor, s:torch.Tensor, y:torch.Tensor, inverse:bool, init_scale: Any, state:dict[str,Any]) -> None:
         """resets ``P`` which is either B or H"""
         set_storage_(P, self.initialize_P(s.numel(), device=P.device, dtype=P.dtype, is_inverse=inverse))
-        if init_scale == 'auto': init_scale = self.auto_initial_scale(s,y)
-        if init_scale >= 1:
+        if init_scale == 'auto':
+            init_scale = self.auto_initial_scale(s,y)
+            state["scaled"] = init_scale is not None
+
+        if init_scale is not None and init_scale != 1:
             if inverse: P /= init_scale
             else: P *= init_scale
 
     @torch.no_grad
-    def update_tensor(self, tensor, param, grad, loss, state, setting):
+    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
         p = param.view(-1); g = tensor.view(-1)
         inverse = setting['inverse']
         M_key = 'H' if inverse else 'B'
@@ -182,6 +187,7 @@ class HessianUpdateStrategy(TensorwiseTransform, ABC):
             state['f_prev'] = loss
             state['p_prev'] = p.clone()
             state['g_prev'] = g.clone()
+            state["scaled"] = False
             return
 
         state['f'] = loss
@@ -205,9 +211,13 @@ class HessianUpdateStrategy(TensorwiseTransform, ABC):
         if gtol is not None and y.abs().max() <= gtol:
             return
 
-        if step == 2 and init_scale == 'auto':
-            if inverse: M /= self.auto_initial_scale(s,y)
-            else: M *= self.auto_initial_scale(s,y)
+        # apply automatic initial scale if it hasn't been applied
+        if (not state["scaled"]) and (init_scale == 'auto'):
+            scale = self.auto_initial_scale(s,y)
+            if scale is not None:
+                state["scaled"] = True
+                if inverse: M /= self.auto_initial_scale(s,y)
+                else: M *= self.auto_initial_scale(s,y)
 
         beta = setting['beta']
         if beta is not None and beta != 0: M = M.clone() # because all of them update it in-place
@@ -223,7 +233,7 @@ class HessianUpdateStrategy(TensorwiseTransform, ABC):
         state['f_prev'] = loss
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
         step = state['step']
 
         if setting['scale_first'] and step == 1:
@@ -250,8 +260,8 @@ class HessianUpdateStrategy(TensorwiseTransform, ABC):
         self.global_state.clear()
         return tensor.mul_(initial_step_size(tensor))
 
-    def get_H(self, var):
-        param = var.params[0]
+    def get_H(self, objective):
+        param = objective.params[0]
         state = self.state[param]
         settings = self.settings[param]
         if "B" in state:
@@ -367,22 +377,21 @@ def bfgs_B_(B:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
     B += term1.sub_(term2)
     return B
 
-def bfgs_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
+
+def bfgs_H_(H: torch.Tensor, s: torch.Tensor, y: torch.Tensor, tol: float):
     sy = s.dot(y)
     if sy <= tol: return H
 
-    sy_sq = safe_clip(sy**2)
-
-    Hy = H@y
-    scale1 = (sy + y.dot(Hy)) / sy_sq
-    term1 = s.outer(s).mul_(scale1)
+    rho = 1.0 / sy
+    Hy = H @ y
 
-    num2 = (Hy.outer(s)).add_(s.outer(y @ H))
-    term2 = num2.div_(sy)
+    term1 = (s.outer(s)).mul_(rho * (1 + rho * y.dot(Hy)))
+    term2 = (Hy.outer(s) + s.outer(Hy)).mul_(rho)
 
-    H += term1.sub_(term2)
+    H.add_(term1).sub_(term2)
     return H
 
+
 class BFGS(_InverseHessianUpdateStrategyDefaults):
     """Broyden–Fletcher–Goldfarb–Shanno Quasi-Newton method. This is usually the most stable quasi-newton method.
 
@@ -428,7 +437,7 @@ class BFGS(_InverseHessianUpdateStrategyDefaults):
     BFGS with backtracking line search:
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.BFGS(),
         tz.m.Backtracking()
@@ -437,7 +446,7 @@ class BFGS(_InverseHessianUpdateStrategyDefaults):
 
     BFGS with trust region
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.LevenbergMarquardt(tz.m.BFGS(inverse=False)),
     )
@@ -505,7 +514,7 @@ class SR1(_InverseHessianUpdateStrategyDefaults):
 
     SR1 with trust region
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.LevenbergMarquardt(tz.m.SR1(inverse=False)),
     )
@@ -1005,7 +1014,7 @@ def gradient_correction(g: TensorList, s: TensorList, y: TensorList):
     return g - (y * (s.dot(g) / sy))
 
 
-class GradientCorrection(Transform):
+class GradientCorrection(TensorTransform):
     """
     Estimates gradient at minima along search direction assuming function is quadratic.
 
@@ -1015,7 +1024,7 @@ class GradientCorrection(Transform):
     L-BFGS with gradient correction
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.LBFGS(inner=tz.m.GradientCorrection()),
         tz.m.Backtracking()
@@ -1027,9 +1036,9 @@ class GradientCorrection(Transform):
 
     """
     def __init__(self):
-        super().__init__(None, uses_grad=False)
+        super().__init__()
 
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         if 'p_prev' not in states[0]:
             p_prev = unpack_states(states, tensors, 'p_prev', init=params)
             g_prev = unpack_states(states, tensors, 'g_prev', init=tensors)
@@ -1154,6 +1163,7 @@ class NewSSM(HessianUpdateStrategy):
             scale_first=scale_first,
             concat_params=concat_params,
             inverse=True,
+            uses_loss=True,
             inner=inner,
         )
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, setting):
@@ -1171,13 +1181,18 @@ class NewSSM(HessianUpdateStrategy):
 
 # this is supposed to be equivalent (and it is)
 def shor_r_(H:torch.Tensor, y:torch.Tensor, alpha:float):
-    p = H@y
-    #(1-y)^2 (ppT)/(pTq)
-    #term = p.outer(p).div_(p.dot(y).clip(min=1e-32))
-    term = p.outer(p).div_(safe_clip(p.dot(y)))
-    H.sub_(term, alpha=1-alpha**2)
+    Hy = H @ y
+    yHy = safe_clip(y.dot(Hy))
+    term = Hy.outer(Hy).div_(yHy)
+    H.sub_(term, alpha=(1-alpha**2))
     return H
 
+# def projected_gradient_(H:torch.Tensor, y:torch.Tensor):
+#     Hy = H @ y
+#     yHy = safe_clip(y.dot(Hy))
+#     H -= (Hy.outer(y) @ H).div_(yHy)
+#     return H
+
 class ShorR(HessianUpdateStrategy):
     """Shor’s r-algorithm.
 

torchzero/modules/quasi_newton/sg2.py

@@ -1,29 +1,39 @@
 import torch
 
-from ...core import Module, Chainable, apply_transform
-from ...utils import TensorList, vec_to_tensors
-from ..second_order.newton import _newton_step, _get_H
+from ...core import Chainable, Transform
+from ...utils import TensorList, unpack_dicts, unpack_states, vec_to_tensors_
+from ...linalg.linear_operator import Dense
+
 
 def sg2_(
     delta_g: torch.Tensor,
     cd: torch.Tensor,
 ) -> torch.Tensor:
-    """cd is c * perturbation, and must be multiplied by two if hessian estimate is two-sided
-    (or divide delta_g by two)."""
+    """cd is c * perturbation."""
 
-    M = torch.outer(1.0 / cd, delta_g)
+    M = torch.outer(0.5 / cd, delta_g)
     H_hat = 0.5 * (M + M.T)
 
     return H_hat
 
 
 
-class SG2(Module):
+class SG2(Transform):
     """second-order stochastic gradient
 
+    2SPSA (second-order SPSA)
+    ```python
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.SPSA(),
+        tz.m.SG2(),
+        tz.m.LR(1e-2),
+    )
+    ```
+
     SG2 with line search
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.SG2(),
         tz.m.Backtracking()
@@ -32,9 +42,9 @@ class SG2(Module):
 
     SG2 with trust region
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
-        tz.m.LevenbergMarquardt(tz.m.SG2()),
+        tz.m.LevenbergMarquardt(tz.m.SG2(beta=0.75. n_samples=4)),
     )
     ```
 
@@ -43,61 +53,51 @@ class SG2(Module):
     def __init__(
         self,
         n_samples: int = 1,
-        h: float = 1e-2,
+        n_first_step_samples: int = 10,
+        start_step: int = 10,
         beta: float | None = None,
-        damping: float = 0,
-        eigval_fn=None,
-        one_sided: bool = False, # one-sided hessian
-        use_lstsq: bool = True,
+        damping: float = 1e-4,
+        h: float = 1e-2,
         seed=None,
+        update_freq: int = 1,
         inner: Chainable | None = None,
     ):
-        defaults = dict(n_samples=n_samples, h=h, beta=beta, damping=damping, eigval_fn=eigval_fn, one_sided=one_sided, seed=seed, use_lstsq=use_lstsq)
-        super().__init__(defaults)
-
-        if inner is not None: self.set_child('inner', inner)
+        defaults = dict(n_samples=n_samples, h=h, beta=beta, damping=damping, seed=seed, start_step=start_step, n_first_step_samples=n_first_step_samples)
+        super().__init__(defaults, update_freq=update_freq, inner=inner)
 
     @torch.no_grad
-    def update(self, var):
-        k = self.global_state.get('step', 0) + 1
-        self.global_state["step"] = k
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
+        k = self.increment_counter("step", 0)
 
-        params = TensorList(var.params)
-        closure = var.closure
+        params = TensorList(objective.params)
+        closure = objective.closure
         if closure is None:
             raise RuntimeError("closure is required for SG2")
         generator = self.get_generator(params[0].device, self.defaults["seed"])
 
-        h = self.get_settings(params, "h")
+        h = unpack_dicts(settings, "h")
         x_0 = params.clone()
-        n_samples = self.defaults["n_samples"]
+        n_samples = fs["n_samples"]
+        if k == 0: n_samples = fs["n_first_step_samples"]
         H_hat = None
 
+        # compute new approximation
         for i in range(n_samples):
             # generate perturbation
             cd = params.rademacher_like(generator=generator).mul_(h)
 
-            # one sided
-            if self.defaults["one_sided"]:
-                g_0 = TensorList(var.get_grad())
-                params.add_(cd)
-                closure()
+            # two sided hessian approximation
+            params.add_(cd)
+            closure()
+            g_p = params.grad.fill_none_(params)
 
-                g_p = params.grad.fill_none_(params)
-                delta_g = (g_p - g_0) * 2
+            params.copy_(x_0)
+            params.sub_(cd)
+            closure()
+            g_n = params.grad.fill_none_(params)
 
-            # two sided
-            else:
-                params.add_(cd)
-                closure()
-                g_p = params.grad.fill_none_(params)
-
-                params.copy_(x_0)
-                params.sub_(cd)
-                closure()
-                g_n = params.grad.fill_none_(params)
-
-                delta_g = g_p - g_n
+            delta_g = g_p - g_n
 
             # restore params
             params.set_(x_0)
@@ -114,179 +114,43 @@ class SG2(Module):
         assert H_hat is not None
         if n_samples > 1: H_hat /= n_samples
 
+        # add damping
+        if fs["damping"] != 0:
+            reg = torch.eye(H_hat.size(0), device=H_hat.device, dtype=H_hat.dtype).mul_(fs["damping"])
+            H_hat += reg
+
         # update H
         H = self.global_state.get("H", None)
         if H is None: H = H_hat
         else:
-            beta = self.defaults["beta"]
-            if beta is None: beta = k / (k+1)
+            beta = fs["beta"]
+            if beta is None: beta = (k+1) / (k+2)
             H.lerp_(H_hat, 1-beta)
 
         self.global_state["H"] = H
 
 
     @torch.no_grad
-    def apply(self, var):
-        dir = _newton_step(
-            var=var,
-            H = self.global_state["H"],
-            damping = self.defaults["damping"],
-            inner = self.children.get("inner", None),
-            H_tfm=None,
-            eigval_fn=self.defaults["eigval_fn"],
-            use_lstsq=self.defaults["use_lstsq"],
-            g_proj=None,
-        )
-
-        var.update = vec_to_tensors(dir, var.params)
-        return var
-
-    def get_H(self,var=...):
-        return _get_H(self.global_state["H"], self.defaults["eigval_fn"])
-
-
-
-
-# two sided
-# we have g via x + d, x - d
-# H via g(x + d), g(x - d)
-# 1 is x, x+2d
-# 2 is x, x-2d
-# 5 evals in total
-
-# one sided
-# g via x, x + d
-# 1 is x, x + d
-# 2 is x, x - d
-# 3 evals and can use two sided for g_0
-
-class SPSA2(Module):
-    """second-order SPSA
-
-    SPSA2 with line search
-    ```python
-    opt = tz.Modular(
-        model.parameters(),
-        tz.m.SPSA2(),
-        tz.m.Backtracking()
-    )
-    ```
-
-    SPSA2 with trust region
-    ```python
-    opt = tz.Modular(
-        model.parameters(),
-        tz.m.LevenbergMarquardt(tz.m.SPSA2()),
-    )
-    ```
-    """
-
-    def __init__(
-        self,
-        n_samples: int = 1,
-        h: float = 1e-2,
-        beta: float | None = None,
-        damping: float = 0,
-        eigval_fn=None,
-        use_lstsq: bool = True,
-        seed=None,
-        inner: Chainable | None = None,
-    ):
-        defaults = dict(n_samples=n_samples, h=h, beta=beta, damping=damping, eigval_fn=eigval_fn, seed=seed, use_lstsq=use_lstsq)
-        super().__init__(defaults)
-
-        if inner is not None: self.set_child('inner', inner)
-
-    @torch.no_grad
-    def update(self, var):
-        k = self.global_state.get('step', 0) + 1
-        self.global_state["step"] = k
+    def apply_states(self, objective, states, settings):
+        fs = settings[0]
+        updates = objective.get_updates()
 
-        params = TensorList(var.params)
-        closure = var.closure
-        if closure is None:
-            raise RuntimeError("closure is required for SPSA2")
+        H: torch.Tensor = self.global_state["H"]
+        k = self.global_state["step"]
+        if k < fs["start_step"]:
+            # don't precondition yet
+            # I guess we can try using trace to scale the update
+            # because it will have horrible scaling otherwise
+            torch._foreach_div_(updates, H.trace())
+            return objective
 
-        generator = self.get_generator(params[0].device, self.defaults["seed"])
+        b = torch.cat([t.ravel() for t in updates])
+        sol = torch.linalg.lstsq(H, b).solution # pylint:disable=not-callable
 
-        h = self.get_settings(params, "h")
-        x_0 = params.clone()
-        n_samples = self.defaults["n_samples"]
-        H_hat = None
-        g_0 = None
-
-        for i in range(n_samples):
-            # perturbations for g and H
-            cd_g = params.rademacher_like(generator=generator).mul_(h)
-            cd_H = params.rademacher_like(generator=generator).mul_(h)
-
-            # evaluate 4 points
-            x_p = x_0 + cd_g
-            x_n = x_0 - cd_g
+        vec_to_tensors_(sol, updates)
+        return objective
 
-            params.set_(x_p)
-            f_p = closure(False)
-            params.add_(cd_H)
-            f_pp = closure(False)
+    def get_H(self, objective=...):
+        return Dense(self.global_state["H"])
 
-            params.set_(x_n)
-            f_n = closure(False)
-            params.add_(cd_H)
-            f_np = closure(False)
-
-            g_p_vec = (f_pp - f_p) / cd_H
-            g_n_vec = (f_np - f_n) / cd_H
-            delta_g = g_p_vec - g_n_vec
-
-            # restore params
-            params.set_(x_0)
 
-            # compute grad
-            g_i = (f_p - f_n) / (2 * cd_g)
-            if g_0 is None: g_0 = g_i
-            else: g_0 += g_i
-
-            # compute H hat
-            H_i = sg2_(
-                delta_g = delta_g.to_vec().div_(2.0),
-                cd = cd_g.to_vec(), # The interval is measured by the original 'cd'
-            )
-            if H_hat is None: H_hat = H_i
-            else: H_hat += H_i
-
-        assert g_0 is not None and H_hat is not None
-        if n_samples > 1:
-            g_0 /= n_samples
-            H_hat /= n_samples
-
-        # set grad to approximated grad
-        var.grad = g_0
-
-        # update H
-        H = self.global_state.get("H", None)
-        if H is None: H = H_hat
-        else:
-            beta = self.defaults["beta"]
-            if beta is None: beta = k / (k+1)
-            H.lerp_(H_hat, 1-beta)
-
-        self.global_state["H"] = H
-
-    @torch.no_grad
-    def apply(self, var):
-        dir = _newton_step(
-            var=var,
-            H = self.global_state["H"],
-            damping = self.defaults["damping"],
-            inner = self.children.get("inner", None),
-            H_tfm=None,
-            eigval_fn=self.defaults["eigval_fn"],
-            use_lstsq=self.defaults["use_lstsq"],
-            g_proj=None,
-        )
-
-        var.update = vec_to_tensors(dir, var.params)
-        return var
-
-    def get_H(self,var=...):
-        return _get_H(self.global_state["H"], self.defaults["eigval_fn"])