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

torchzero/modules/misc/split.py

@@ -81,7 +81,7 @@ class Split(Module):
     Muon with Adam fallback using same hyperparams as https://github.com/KellerJordan/Muon
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.NAG(0.95),
         tz.m.Split(

torchzero/modules/misc/switch.py

@@ -19,7 +19,7 @@ class Alternate(Module):
 
     ```python
 
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Alternate(
             tz.m.Adam(),
@@ -89,7 +89,7 @@ class Switch(Alternate):
     Start with Adam, switch to L-BFGS after 1000th step and Truncated Newton on 2000th step.
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Switch(
             [tz.m.Adam(), tz.m.LR(1e-3)],

torchzero/modules/momentum/cautious.py

@@ -57,7 +57,7 @@ class Cautious(TensorTransform):
     Cautious Adam
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         bench.parameters(),
         tz.m.Adam(),
         tz.m.Cautious(),
@@ -173,7 +173,7 @@ class ScaleByGradCosineSimilarity(TensorTransform):
 
     Scaled Adam
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         bench.parameters(),
         tz.m.Adam(),
         tz.m.ScaleByGradCosineSimilarity(),
@@ -211,7 +211,7 @@ class ScaleModulesByCosineSimilarity(Module):
 
     Adam scaled by similarity to RMSprop
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         bench.parameters(),
         tz.m.ScaleModulesByCosineSimilarity(
             main = tz.m.Adam(),

torchzero/modules/momentum/momentum.py

@@ -6,7 +6,7 @@ import torch
 
 from ...core import  TensorTransform
 from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
-from ..functional import debias, ema_
+from ..opt_utils import debias, ema_
 
 
 class EMA(TensorTransform):

torchzero/modules/ops/higher_level.py

@@ -6,7 +6,7 @@ import torch
 
 from ...core import  TensorTransform
 from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
-from ..functional import (
+from ..opt_utils import (
     centered_ema_sq_,
     debias,
     debias_second_momentum,

torchzero/modules/ops/multi.py

@@ -144,7 +144,7 @@ class Graft(MultiOperationBase):
 
     Shampoo grafted to Adam
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.GraftModules(
             direction = tz.m.Shampoo(),

torchzero/modules/projections/projection.py

@@ -149,8 +149,11 @@ class ProjectionBase(Module, ABC):
             Iterable[torch.Tensor]: unprojected tensors of the same shape as params
         """
 
+    def update(self, objective: Objective): raise RuntimeError("projections don't support update/apply")
+    def apply(self, objective: Objective): raise RuntimeError("projections don't support update/apply")
+
     @torch.no_grad
-    def apply(self, objective: Objective):
+    def step(self, objective: Objective):
         params = objective.params
         settings = [self.settings[p] for p in params]
 
@@ -266,7 +269,7 @@ class ProjectionBase(Module, ABC):
 
         # ----------------------------------- step ----------------------------------- #
         projected_obj.params = projected_params
-        projected_obj = self.children['modules'].apply(projected_obj)
+        projected_obj = self.children['modules'].step(projected_obj)
 
         # empty fake params storage
         # this doesn't affect update/grad because it is a different python object, set_ changes storage on an object

torchzero/modules/quasi_newton/__init__.py

@@ -30,4 +30,4 @@ from .quasi_newton import (
     ThomasOptimalMethod,
 )
 
-from .sg2 import SG2, SPSA2
+from .sg2 import SG2

torchzero/modules/quasi_newton/damping.py

@@ -5,7 +5,7 @@ import torch
 
 from ...utils import TensorList
 from ...linalg.linear_operator import DenseInverse, LinearOperator
-from ..functional import safe_clip
+from ..opt_utils import safe_clip
 
 
 class DampingStrategy(Protocol):

torchzero/modules/quasi_newton/diagonal_quasi_newton.py

@@ -9,7 +9,7 @@ from .quasi_newton import (
     _InverseHessianUpdateStrategyDefaults,
 )
 
-from ..functional import safe_clip
+from ..opt_utils import safe_clip
 
 
 def diagonal_bfgs_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):

torchzero/modules/quasi_newton/lbfgs.py

@@ -7,7 +7,7 @@ import torch
 from ...core import Chainable, TensorTransform
 from ...utils import TensorList, as_tensorlist, unpack_states
 from ...linalg.linear_operator import LinearOperator
-from ..functional import initial_step_size
+from ..opt_utils import initial_step_size
 from .damping import DampingStrategyType, apply_damping
 
 
@@ -188,7 +188,7 @@ class LBFGS(TensorTransform):
 
     L-BFGS with line search
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.LBFGS(100),
         tz.m.Backtracking()
@@ -197,7 +197,7 @@ class LBFGS(TensorTransform):
 
     L-BFGS with trust region
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.TrustCG(tz.m.LBFGS())
     )

torchzero/modules/quasi_newton/lsr1.py

@@ -7,7 +7,7 @@ import torch
 from ...core import Chainable, Module, TensorTransform, Objective, step
 from ...utils import NumberList, TensorList, as_tensorlist, generic_finfo_tiny, unpack_states, vec_to_tensors_
 from ...linalg.linear_operator import LinearOperator
-from ..functional import initial_step_size
+from ..opt_utils import initial_step_size
 from .damping import DampingStrategyType, apply_damping
 
 
@@ -110,7 +110,7 @@ class LSR1(TensorTransform):
 
     L-SR1 with line search
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.SR1(),
         tz.m.StrongWolfe(c2=0.1, fallback=True)
@@ -119,7 +119,7 @@ class LSR1(TensorTransform):
 
     L-SR1 with trust region
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.TrustCG(tz.m.LSR1())
     )

torchzero/modules/quasi_newton/quasi_newton.py

@@ -8,7 +8,7 @@ import torch
 from ...core import Chainable, Module, TensorTransform, Transform
 from ...utils import TensorList, set_storage_, unpack_states, safe_dict_update_
 from ...linalg import linear_operator
-from ..functional import initial_step_size, safe_clip
+from ..opt_utils import initial_step_size, safe_clip
 
 
 
@@ -106,11 +106,12 @@ class HessianUpdateStrategy(TensorTransform, 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,18 +142,22 @@ class HessianUpdateStrategy(TensorTransform, 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
 
@@ -182,6 +187,7 @@ class HessianUpdateStrategy(TensorTransform, 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(TensorTransform, 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
@@ -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)),
     )
@@ -1015,7 +1024,7 @@ class GradientCorrection(TensorTransform):
     L-BFGS with gradient correction
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.LBFGS(inner=tz.m.GradientCorrection()),
         tz.m.Backtracking()
@@ -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.