torchzero 0.3.5__tar.gz → 0.3.8__tar.gz

{torchzero-0.3.5 → torchzero-0.3.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.3.5
+Version: 0.3.8
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 License: MIT License
@@ -156,7 +156,7 @@ for epoch in range(100):
   * `Newton`: Classic Newton's method.
   * `NewtonCG`: Matrix-free newton's method with conjugate gradient solver.
   * `NystromSketchAndSolve`: Nyström sketch-and-solve method.
-  * `NystromPCG`: NewtonCG with Nyström preconditioning (my current recommendation).
+  * `NystromPCG`: NewtonCG with Nyström preconditioning (usually beats NewtonCG).
 
 * **Quasi-Newton**: Approximate second-order optimization methods.
   * `LBFGS`: Limited-memory BFGS.

{torchzero-0.3.5 → torchzero-0.3.8}/README.md

@@ -117,7 +117,7 @@ for epoch in range(100):
   * `Newton`: Classic Newton's method.
   * `NewtonCG`: Matrix-free newton's method with conjugate gradient solver.
   * `NystromSketchAndSolve`: Nyström sketch-and-solve method.
-  * `NystromPCG`: NewtonCG with Nyström preconditioning (my current recommendation).
+  * `NystromPCG`: NewtonCG with Nyström preconditioning (usually beats NewtonCG).
 
 * **Quasi-Newton**: Approximate second-order optimization methods.
   * `LBFGS`: Limited-memory BFGS.

{torchzero-0.3.5 → torchzero-0.3.8}/pyproject.toml

@@ -2,7 +2,7 @@
 # STEP 1 - COMMIT NEW CHANGES BUT DON'T PUSH THEM YET
 # STEP 2 - BUMP VERSION AND COMMIT IT (DONT PUSH!!!!)
 # STEP 3 - CREATE TAG WITH THAT VERSION
-# STEP 4 - PUSH CHANGES
+# STEP 4 - PUSH (SYNC) CHANGES
 # STEP 5 - PUSH TAG
 
 [build-system]
@@ -13,7 +13,7 @@ build-backend = "setuptools.build_meta"
 name = "torchzero"
 description = "Modular optimization library for PyTorch."
 
-version = "0.3.5"
+version = "0.3.8"
 dependencies = [
   "torch",
   "numpy",

{torchzero-0.3.5 → torchzero-0.3.8}/tests/test_opts.py

@@ -745,7 +745,7 @@ SSVM = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.SSVM(1), tz.m.StrongWolfe()),
     sphere_opt=lambda p: tz.Modular(p, tz.m.SSVM(1), tz.m.StrongWolfe()),
     needs_closure=True,
-    func='rosen', steps=50, loss=1e-12, merge_invariant=True,
+    func='rosen', steps=50, loss=0.02, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 

{torchzero-0.3.5 → torchzero-0.3.8}/tests/test_tensorlist.py

@@ -1301,7 +1301,7 @@ def test_reduction_ops(simple_tl: TensorList, reduction_method, dim, keepdim):
         expected_tl = TensorList(expected_list)
         assert isinstance(result, TensorList)
         assert len(result) == len(expected_tl)
-        assert_tl_allclose(result, expected_tl, atol=1e-6) # Use allclose due to potential float variations
+        assert_tl_allclose(result, expected_tl, atol=1e-3) # Use allclose due to potential float variations
 
 # --- Grafting, Rescaling, Normalizing, Clipping ---
 
@@ -1381,8 +1381,8 @@ def test_rescale(simple_tl: TensorList, dim):
         assert torch.allclose(rescaled_scalar.global_min(), torch.tensor(min_val))
         assert torch.allclose(rescaled_scalar.global_max(), torch.tensor(max_val))
     else:
-        assert_tl_allclose(rescaled_scalar_min, TensorList([torch.full_like(t, min_val) for t in rescaled_scalar_min]),atol=1e-4)
-        assert_tl_allclose(rescaled_scalar_max, TensorList([torch.full_like(t, max_val) for t in rescaled_scalar_max]),atol=1e-4)
+        assert_tl_allclose(rescaled_scalar_min, TensorList([torch.full_like(t, min_val) for t in rescaled_scalar_min]),atol=1e-3)
+        assert_tl_allclose(rescaled_scalar_max, TensorList([torch.full_like(t, max_val) for t in rescaled_scalar_max]),atol=1e-3)
 
 
     # Rescale list
@@ -1402,8 +1402,8 @@ def test_rescale(simple_tl: TensorList, dim):
          assert global_max_rescaled < avg_max + 1.0 # Loose check
 
     else:
-        assert_tl_allclose(rescaled_list_min, TensorList([torch.full_like(t, mn) for t, mn in zip(rescaled_list_min, min_list)]),atol=1e-4)
-        assert_tl_allclose(rescaled_list_max, TensorList([torch.full_like(t, mx) for t, mx in zip(rescaled_list_max, max_list)]),atol=1e-4)
+        assert_tl_allclose(rescaled_list_min, TensorList([torch.full_like(t, mn) for t, mn in zip(rescaled_list_min, min_list)]),atol=1e-3)
+        assert_tl_allclose(rescaled_list_max, TensorList([torch.full_like(t, mx) for t, mx in zip(rescaled_list_max, max_list)]),atol=1e-3)
 
     # Rescale to 01 helper
     rescaled_01 = simple_tl.rescale_to_01(dim=dim, eps=eps)
@@ -1413,8 +1413,8 @@ def test_rescale(simple_tl: TensorList, dim):
         assert torch.allclose(rescaled_01.global_min(), torch.tensor(0.0))
         assert torch.allclose(rescaled_01.global_max(), torch.tensor(1.0))
     else:
-        assert_tl_allclose(rescaled_01_min, TensorList([torch.zeros_like(t) for t in rescaled_01_min]), atol=1e-4)
-        assert_tl_allclose(rescaled_01_max, TensorList([torch.ones_like(t) for t in rescaled_01_max]), atol=1e-4)
+        assert_tl_allclose(rescaled_01_min, TensorList([torch.zeros_like(t) for t in rescaled_01_min]), atol=1e-3)
+        assert_tl_allclose(rescaled_01_max, TensorList([torch.ones_like(t) for t in rescaled_01_max]), atol=1e-3)
 
 
     # Test inplace
@@ -1454,11 +1454,11 @@ def test_normalize(big_tl: TensorList, dim):
     normalized_scalar_var = normalized_scalar.var(dim=dim if dim != 'global' else None)
 
     if dim == 'global':
-        assert torch.allclose(normalized_scalar.global_mean(), torch.tensor(mean_val), atol=1e-4)
-        assert torch.allclose(normalized_scalar.global_var(), torch.tensor(var_val), atol=1e-4)
+        assert torch.allclose(normalized_scalar.global_mean(), torch.tensor(mean_val), atol=1e-3)
+        assert torch.allclose(normalized_scalar.global_var(), torch.tensor(var_val), atol=1e-3)
     else:
-        assert_tl_allclose(normalized_scalar_mean, TensorList([torch.full_like(t, mean_val) for t in normalized_scalar_mean]), atol=1e-4)
-        assert_tl_allclose(normalized_scalar_var, TensorList([torch.full_like(t, var_val) for t in normalized_scalar_var]), atol=1e-4)
+        assert_tl_allclose(normalized_scalar_mean, TensorList([torch.full_like(t, mean_val) for t in normalized_scalar_mean]), atol=1e-3)
+        assert_tl_allclose(normalized_scalar_var, TensorList([torch.full_like(t, var_val) for t in normalized_scalar_var]), atol=1e-3)
 
     # Normalize list mean/var
     normalized_list = simple_tl.normalize(mean_list, var_list, dim=dim)
@@ -1476,19 +1476,19 @@ def test_normalize(big_tl: TensorList, dim):
         #  assert torch.allclose(global_mean_rescaled, torch.tensor(avg_mean), rtol=1e-1, atol=1e-1) # Loose check
         #  assert torch.allclose(global_var_rescaled, torch.tensor(avg_var), rtol=1e-1, atol=1e-1) # Loose check
     else:
-        assert_tl_allclose(normalized_list_mean, TensorList([torch.full_like(t, m) for t, m in zip(normalized_list_mean, mean_list)]), atol=1e-4)
-        assert_tl_allclose(normalized_list_var, TensorList([torch.full_like(t, v) for t, v in zip(normalized_list_var, var_list)]), atol=1e-4)
+        assert_tl_allclose(normalized_list_mean, TensorList([torch.full_like(t, m) for t, m in zip(normalized_list_mean, mean_list)]), atol=1e-3)
+        assert_tl_allclose(normalized_list_var, TensorList([torch.full_like(t, v) for t, v in zip(normalized_list_var, var_list)]), atol=1e-3)
 
     # Z-normalize helper
     znorm = simple_tl.znormalize(dim=dim, eps=1e-10)
     znorm_mean = znorm.mean(dim=dim if dim != 'global' else None)
     znorm_var = znorm.var(dim=dim if dim != 'global' else None)
     if dim == 'global':
-        assert torch.allclose(znorm.global_mean(), torch.tensor(0.0), atol=1e-4)
-        assert torch.allclose(znorm.global_var(), torch.tensor(1.0), atol=1e-4)
+        assert torch.allclose(znorm.global_mean(), torch.tensor(0.0), atol=1e-3)
+        assert torch.allclose(znorm.global_var(), torch.tensor(1.0), atol=1e-3)
     else:
-        assert_tl_allclose(znorm_mean, TensorList([torch.zeros_like(t) for t in znorm_mean]), atol=1e-4)
-        assert_tl_allclose(znorm_var, TensorList([torch.ones_like(t) for t in znorm_var]), atol=1e-4)
+        assert_tl_allclose(znorm_mean, TensorList([torch.zeros_like(t) for t in znorm_mean]), atol=1e-3)
+        assert_tl_allclose(znorm_var, TensorList([torch.ones_like(t) for t in znorm_var]), atol=1e-3)
 
 
     # Test inplace

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/core/preconditioner.py

@@ -38,7 +38,7 @@ class Preconditioner(Transform):
 
 
     def _tensor_wise_transform(self, tensors:list[torch.Tensor], params:list[torch.Tensor], grads:list[torch.Tensor] | None, vars:Vars) -> list[torch.Tensor]:
-        step = self.global_state.get('step', 0)
+        step = self.global_state.get('__step', 0)
         states = [self.state[p] for p in params]
         settings = [self.settings[p] for p in params]
         global_settings = settings[0]
@@ -47,8 +47,10 @@ class Preconditioner(Transform):
         scale_first = global_settings['__scale_first']
         scale_factor = 0
         if scale_first and step == 0:
-            # initial step size guess from pytorch LBFGS
-            scale_factor = TensorList(tensors).abs().sum()
+            # initial step size guess from pytorch LBFGS was too unstable
+            # I switched to norm
+            tensors = TensorList(tensors)
+            scale_factor = tensors.abs().global_mean().clip(min=1)
 
         # update preconditioner
         if step % update_freq == 0:
@@ -65,11 +67,11 @@ class Preconditioner(Transform):
         if scale_first and step == 0:
             torch._foreach_div_(tensors, scale_factor)
 
-        self.global_state['step'] = step + 1
+        self.global_state['__step'] = step + 1
         return tensors
 
     def _concat_transform(self, tensors:list[torch.Tensor], params:list[torch.Tensor], grads:list[torch.Tensor] | None, vars:Vars) -> list[torch.Tensor]:
-        step = self.global_state.get('step', 0)
+        step = self.global_state.get('__step', 0)
         tensors_vec = torch.cat([t.ravel() for t in tensors])
         params_vec = torch.cat([p.ravel() for p in params])
         grads_vec = [torch.cat([g.ravel() for g in grads])] if grads is not None else None
@@ -82,8 +84,8 @@ class Preconditioner(Transform):
         scale_first = global_settings['__scale_first']
         scale_factor = 0
         if scale_first and step == 0:
-            # initial step size guess from pytorch LBFGS
-            scale_factor = tensors_vec.abs().sum()
+            # initial step size guess from pytorch LBFGS was too unstable
+            scale_factor = tensors_vec.abs().mean().clip(min=1)
 
         # update preconditioner
         if step % update_freq == 0:
@@ -99,11 +101,10 @@ class Preconditioner(Transform):
 
         # scale initial step, when preconditioner might not have been applied
         if scale_first and step == 0:
-            if scale_factor >= torch.finfo(tensors_vec.dtype).eps:
-                tensors_vec /= scale_factor
+            tensors_vec /= scale_factor
 
         tensors = vec_to_tensors(vec=tensors_vec, reference=tensors)
-        self.global_state['step'] = step + 1
+        self.global_state['__step'] = step + 1
         return tensors
 
     @torch.no_grad

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/experimental/__init__.py

@@ -3,7 +3,7 @@ from .adadam import Adadam
 from .adamY import AdamY
 from .adasoap import AdaSOAP
 from .curveball import CurveBall
-from .dsoap import DSOAP
+from .soapy import SOAPY
 from .gradmin import GradMin
 from .reduce_outward_lr import ReduceOutwardLR
 from .spectral import SpectralPreconditioner
@@ -11,4 +11,5 @@ from .subspace_preconditioners import (
     HistorySubspacePreconditioning,
     RandomSubspacePreconditioning,
 )
-from .tropical_newton import TropicalNewton
+from .tropical_newton import TropicalNewton
+from .newton_solver import NewtonSolver

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/experimental/absoap.py

@@ -140,11 +140,17 @@ Source=Literal['p','g','s','y', 'gy', 'sy', 'sn', 'yn', 'gys', 'sys','sn', 'yn']
 class ABSOAP(Transform):
     """SOAP but with two extra letters included in its name in order to improve converence
 
+    so what you can do is choose what goes into what ,and that is supposed to be good.
+
     new args
 
     scale by s whether to scale gradient differences by parameter differences
 
     y_to_ema2 whether to use gradient differences for exponential moving average too
+
+    okay I changed these args into another ones
+
+    BASICALLY THIS IS FOR MY EXPERIMENTS
     """
     def __init__(
         self,
@@ -213,7 +219,7 @@ class ABSOAP(Transform):
             if 'g_prev' not in state:
                 state['p_prev'] = p.clone()
                 state['g_prev'] = t.clone()
-                updates.append(tensors[i].sign())
+                updates.append(tensors[i].clip(-0.1,0.1))
                 continue
 
             p_prev = state['p_prev']
@@ -285,7 +291,7 @@ class ABSOAP(Transform):
                     state['Q'] = get_orthogonal_matrix(state['GG'])
 
                 state['step'] = 0
-                updates.append(tensors[i].sign())
+                updates.append(tensors[i].clip(-0.1,0.1))
                 continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
                 # I use sign instead as to not mess up with next modules. 1st Adam step is always sign anyway.
 

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/experimental/adadam.py

@@ -50,7 +50,7 @@ def adadam_(
     return None
 
 class Adadam(Module):
-    """Adam with a diagonally preconditioned preconditioner and a graceful name."""
+    """Adam with a diagonally preconditioned preconditioner."""
     def __init__(
         self,
         beta1: float = 0.9,

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/experimental/adamY.py

@@ -37,7 +37,7 @@ def adamy_(
         p_prev.copy_(p)
         g_prev.copy_(g)
 
-        update = g.sign().lazy_mul_(alpha*0.1)
+        update = g.clip(-0.1,0.1).lazy_mul_(alpha)
         if params_ is None: return update
         params_.sub_(update)
         return None

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/experimental/adasoap.py

@@ -218,7 +218,7 @@ class AdaSOAP(Transform):
                     state['Q'] = get_orthogonal_matrix(GG_precond)
 
                 state['step'] = 0
-                updates.append(tensors[i].sign())
+                updates.append(tensors[i].clip(-0.1,0.1))
                 continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
                 # I use sign instead as to not mess up with next modules. 1st Adam step is always sign anyway.
 

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/experimental/algebraic_newton.py

@@ -71,7 +71,7 @@ def tikhonov(H: torch.Tensor, reg: float, algebra: ta.Algebra = ta.TropicalSemir
 
 
 class AlgebraicNewton(Module):
-    """newton in other algebras, not practical because solving linear system is very hard."""
+    """newton in other algebras, not that it works."""
     def __init__(
         self,
         reg: float | None = None,

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/experimental/curveball.py

@@ -13,7 +13,7 @@ def curveball(
     momentum: float | NumberList,
     precond_lr: float | NumberList,
 ):
-    """returns z_, clone it!!!"""
+    """returns z_, clone it!!! (no just negate it)"""
     delta = Hz + tensors
     z_.mul_(momentum).sub_(delta.mul_(precond_lr)) # z ← ρz − βΔ
     return z_

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/experimental/gradmin.py

@@ -14,7 +14,7 @@ from ..smoothing.gaussian import Reformulation
 
 
 class GradMin(Reformulation):
-    """Reformulates the objective to minimize sum of gradient magnitudes via autograd.
+    """Reformulates the objective to minimize sum of gradient magnitudes via autograd. This is not expected to be practical.
 
     Args:
         loss_term (float, optional): adds loss value times this to sum of gradient magnitudes. Defaults to 1.

torchzero-0.3.8/torchzero/modules/experimental/newton_solver.py

@@ -0,0 +1,88 @@
+from collections.abc import Callable, Iterable
+from typing import Any, Literal, overload
+
+import torch
+
+from ...core import Chainable, Module, apply, Modular
+from ...utils import TensorList, as_tensorlist
+from ...utils.derivatives import hvp
+from ..quasi_newton import LBFGS
+
+class NewtonSolver(Module):
+    """Matrix free newton via with any custom solver (usually it is better to just use NewtonCG or NystromPCG is even better)"""
+    def __init__(
+        self,
+        solver: Callable[[list[torch.Tensor]], Any] = lambda p: Modular(p, LBFGS()),
+        maxiter=None,
+        tol=1e-3,
+        reg: float = 0,
+        warm_start=True,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(tol=tol, maxiter=maxiter, reg=reg, warm_start=warm_start, solver=solver)
+        super().__init__(defaults,)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, vars):
+        params = TensorList(vars.params)
+        closure = vars.closure
+        if closure is None: raise RuntimeError('NewtonCG requires closure')
+
+        settings = self.settings[params[0]]
+        solver_cls = settings['solver']
+        maxiter = settings['maxiter']
+        tol = settings['tol']
+        reg = settings['reg']
+        warm_start = settings['warm_start']
+
+        # ---------------------- Hessian vector product function --------------------- #
+        grad = vars.get_grad(create_graph=True)
+
+        def H_mm(x):
+            with torch.enable_grad():
+                Hvp = TensorList(hvp(params, grad, x, create_graph=True))
+                if reg != 0: Hvp = Hvp + (x*reg)
+                return Hvp
+
+        # -------------------------------- inner step -------------------------------- #
+        b = as_tensorlist(grad)
+        if 'inner' in self.children:
+            b = as_tensorlist(apply(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, vars=vars))
+
+        # ---------------------------------- run cg ---------------------------------- #
+        x0 = None
+        if warm_start: x0 = self.get_state('prev_x', params=params, cls=TensorList) # initialized to 0 which is default anyway
+        if x0 is None: x = b.zeros_like().requires_grad_(True)
+        else: x = x0.clone().requires_grad_(True)
+
+        solver = solver_cls(x)
+        def lstsq_closure(backward=True):
+            Hx = H_mm(x)
+            loss = (Hx-b).pow(2).global_mean()
+            if backward:
+                solver.zero_grad()
+                loss.backward(inputs=x)
+            return loss
+
+        if maxiter is None: maxiter = b.global_numel()
+        loss = None
+        initial_loss = lstsq_closure(False)
+        if initial_loss > tol:
+            for i in range(maxiter):
+                loss = solver.step(lstsq_closure)
+                assert loss is not None
+                if min(loss, loss/initial_loss) < tol: break
+
+        print(f'{loss = }')
+
+        if warm_start:
+            assert x0 is not None
+            x0.copy_(x)
+
+        vars.update = x.detach()
+        return vars
+
+

torchzero-0.3.5/torchzero/modules/experimental/dsoap.py → torchzero-0.3.8/torchzero/modules/experimental/soapy.py

@@ -3,7 +3,7 @@ from operator import itemgetter
 import torch
 
 from ...core import Chainable, Transform, apply
-from ...modules.optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
+from ..optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
 
 @torch.no_grad
 def update_soap_covariances_(
@@ -135,7 +135,7 @@ def get_orthogonal_matrix_QR(exp_avg_sq: torch.Tensor, GG: list[torch.Tensor | N
 
     return final, exp_avg_sq
 
-class DSOAP(Transform):
+class SOAPY(Transform):
     """SOAP but uses scaled gradient differences
 
     new args
@@ -195,7 +195,7 @@ class DSOAP(Transform):
             if 'g_prev' not in state:
                 state['p_prev'] = p.clone()
                 state['g_prev'] = t.clone()
-                updates.append(tensors[i].sign())
+                updates.append(tensors[i].clip(-0.1,0.1))
                 continue
 
             p_prev = state['p_prev']
@@ -228,7 +228,7 @@ class DSOAP(Transform):
                     state['Q'] = get_orthogonal_matrix(state['GG'])
 
                 state['step'] = 0
-                updates.append(tensors[i].sign())
+                updates.append(tensors[i].clip(-0.1,0.1))
                 continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
                 # I use sign instead as to not mess up with next modules. 1st Adam step is always sign anyway.
 

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/experimental/spectral.py

@@ -194,8 +194,10 @@ class SpectralPreconditioner(TensorwisePreconditioner):
         order (int, optional):
             whitening order, 1 approximates FIM (maybe), 2 - hessian (maybe), 3+ - god knows what.
         solver (str, optional): what to use for whitening. Defaults to 'svd'.
-        U_beta (float | None, optional): beta for U (probably a bad idea). Defaults to None.
-        S_beta (float | None, optional): beta for S (probably a bad idea). Defaults to None.
+        A_beta (float | None, optional):
+            beta for U (in SVD and other letters in other solvers) (probably a bad idea). Defaults to None.
+        B_beta (float | None, optional):
+            beta for S (in SVD and other letters in other solvers) (probably a bad idea). Defaults to None.
         interval (int, optional): How often to update history. Defaults to 1 (every step).
         concat_params (bool, optional):
             whether to apply preconditioning to each tensor (False, default) or to all tensors concatenated into a vector (True). Latter will be slower but captures interactions between layers. Defaults to True.
@@ -275,7 +277,7 @@ class SpectralPreconditioner(TensorwisePreconditioner):
         A = state.get('A', None)
         if A is None:
             # make a conservative step to avoid issues due to different GD scaling
-            return tensor.div_(max(1, tensor.abs().sum())) # pyright:ignore[reportArgumentType]
+            return tensor.clip_(-0.1, 0.1) # pyright:ignore[reportArgumentType]
 
         B = state['B']
         update = solver.apply(tensor.view(-1), A, B).view_as(tensor)

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/experimental/subspace_preconditioners.py

@@ -38,16 +38,19 @@ def apply_subspace_preconditioner(
     return basis @ update_projected # d
 
 class RandomSubspacePreconditioning(Transform):
-    """full matrix rmsprop in random subspace"""
-    def __init__(self, k: int, beta: float | None = 0.99):
-        defaults = dict(k=k, beta=beta)
+    """full matrix rmsprop in random slowly changing subspace"""
+    def __init__(self, k: int, beta: float | None = 0.99, basis_beta: float | None = 0.99, inner: Chainable | None = None):
+        defaults = dict(k=k, beta=beta, basis_beta=basis_beta)
         super().__init__(defaults, uses_grad=False)
 
+        if inner is not None: self.set_child('inner', inner)
+
     def transform(self, tensors, params, grads, vars):
         settings = self.settings[params[0]]
         g = torch.cat([t.view(-1) for t in tensors])
         k = settings['k']
         beta = settings['beta']
+        basis_beta = settings['basis_beta']
 
         if 'basis' not in self.global_state:
             self.global_state['basis'] = torch.randn(g.numel(), k, device=g.device, dtype=g.dtype)
@@ -56,13 +59,19 @@ class RandomSubspacePreconditioning(Transform):
         basis = self.global_state['basis']
         accumulator = self.global_state['accumulator']
 
+        if basis_beta is not None:
+            basis.lerp_(torch.randn_like(basis), 1-basis_beta)
+
         update_subspace_preconditioner_(g, basis, accumulator, beta)
+
+        if 'inner' in self.children:
+            tensors = apply(self.children['inner'], tensors, params, grads, vars)
+            g = torch.cat([t.view(-1) for t in tensors])
+
         try:
             preconditioned = apply_subspace_preconditioner(g, basis, accumulator)
         except torch.linalg.LinAlgError:
-            denom = g.abs().sum()
-            if denom <= 1e-10: denom = torch.ones_like(denom)
-            preconditioned = g / g.abs().sum()
+            preconditioned = g.clip(-0.1, 0.1)
         vec_to_tensors_(preconditioned, tensors)
 
         return tensors
@@ -119,9 +128,7 @@ class HistorySubspacePreconditioning(Transform):
         try:
             preconditioned = apply_subspace_preconditioner(g, basis, accumulator)
         except torch.linalg.LinAlgError:
-            denom = g.abs().sum()
-            if denom <= 1e-10: denom = torch.ones_like(denom)
-            preconditioned = g / g.abs().sum()
+            preconditioned = g.clip(-0.1,0.1)
         vec_to_tensors_(preconditioned, tensors)
 
         return tensors

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/optimizers/soap.py

@@ -222,8 +222,7 @@ class SOAP(Transform):
                     state['Q'] = get_orthogonal_matrix(state['GG'])
 
                 state['step'] = 0
-                updates.append(tensors[i].sign().div_(10))
-                # updates.append(tensors[i] / tensors[i].abs().sum())
+                updates.append(tensors[i].clip(-0.1, 0.1))
                 continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
                 # I use scaled update instead as to not mess up with next modules.
 

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/projections/projection.py

@@ -1,4 +1,5 @@
 import math
+from functools import partial
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from typing import Any, Literal
@@ -33,6 +34,25 @@ def _make_projected_closure(closure, vars: Vars, projection: "Projection",
 
     return projected_closure
 
+def _projected_get_grad_override(
+    retain_graph: bool | None = None,
+    create_graph: bool = False,
+    projection: Any = ...,
+    unprojected_vars: Any = ...,
+    self: Any = ...,
+):
+    assert isinstance(projection, Projection)
+    assert isinstance(unprojected_vars, Vars)
+    assert isinstance(self, Vars)
+
+    if self.grad is not None: return self.grad
+    grads = unprojected_vars.get_grad(retain_graph, create_graph)
+    projected_grads = list(projection.project(grads, self, current='grads'))
+    self.grad = projected_grads
+    for p, g in zip(self.params, projected_grads):
+        p.grad = g
+    return self.grad
+
 
 class Projection(Module, ABC):
     """
@@ -137,6 +157,12 @@ class Projection(Module, ABC):
 
         # step
         projected_vars.params = self._projected_params
+        projected_vars.get_grad = partial(
+            _projected_get_grad_override,
+            projection=self,
+            unprojected_vars=vars,
+            self=projected_vars,
+        )
         projected_vars = self.children['modules'].step(projected_vars)
 
         # empty fake params storage
@@ -149,7 +175,7 @@ class Projection(Module, ABC):
         unprojected_vars = projected_vars.clone(clone_update=False)
         unprojected_vars.closure = vars.closure
         unprojected_vars.params = vars.params
-        if unprojected_vars.grad is None: unprojected_vars.grad = vars.grad
+        unprojected_vars.grad = vars.grad
 
         if self._project_update:
             assert projected_vars.update is not None

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/quasi_newton/experimental/modular_lbfgs.py

@@ -37,10 +37,11 @@ def lbfgs(
     z_tfm: Any,
 ):
     if len(s_history) == 0 or y_k is None or ys_k is None:
-        # dir = params.grad.sign() # may work fine
 
-        # initial step size guess taken from pytorch L-BFGS
-        return tensors_.mul_(min(1.0, 1.0 / tensors_.abs().global_sum())) # pyright: ignore[reportArgumentType]
+        # initial step size guess modified from pytorch L-BFGS
+        scale = 1 / tensors_.abs().global_sum()
+        if scale < 1e-5: scale = 1 / tensors_.abs().mean()
+        return tensors_.mul_(min(1.0, scale)) # pyright: ignore[reportArgumentType]
 
     else:
         # 1st loop

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/quasi_newton/lbfgs.py

@@ -36,10 +36,11 @@ def lbfgs(
     step: int,
 ):
     if len(s_history) == 0 or y_k is None or ys_k is None:
-        # dir = params.grad.sign() # may work fine
 
-        # initial step size guess taken from pytorch L-BFGS
-        return tensors_.mul_(min(1.0, 1.0 / tensors_.abs().global_sum())) # pyright: ignore[reportArgumentType]
+        # initial step size guess modified from pytorch L-BFGS
+        scale = 1 / tensors_.abs().global_sum()
+        if scale < 1e-5: scale = 1 / tensors_.abs().mean()
+        return tensors_.mul_(min(1.0, scale)) # pyright: ignore[reportArgumentType]
 
     else:
         # 1st loop

{torchzero-0.3.5 → torchzero-0.3.8}/torchzero/modules/quasi_newton/lsr1.py

@@ -17,8 +17,9 @@ def lsr1_(
 ):
     if step == 0 or not s_history:
         # initial step size guess from pytorch
-        tensors_.div_(max(1.0, tensors_.abs().global_sum())) # pyright:ignore[reportArgumentType]
-        return tensors_
+        scale = 1 / tensors_.abs().global_sum()
+        if scale < 1e-5: scale = 1 / tensors_.abs().mean()
+        return tensors_.mul_(min(1.0, scale)) # pyright: ignore[reportArgumentType]
 
     m = len(s_history)
 
@@ -64,7 +65,9 @@ def lsr1_(
         Hx.add_(w_k, alpha=w_k.dot(tensors_) / wy) # pyright:ignore[reportArgumentType]
 
     if scale_second and step == 1:
-        Hx.div_(max(1.0, tensors_.abs().global_sum())) # pyright:ignore[reportArgumentType]
+        scale = 1 / tensors_.abs().global_sum()
+        if scale < 1e-5: scale = 1 / tensors_.abs().mean()
+        Hx.mul_(min(1.0, scale)) # pyright:ignore[reportArgumentType]
     return Hx