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

torchzero/modules/adaptive/psgd/psgd_lra_whiten.py

@@ -0,0 +1,116 @@
+# pylint:disable=not-callable
+"""all functions are from https://github.com/lixilinx/psgd_torch/blob/master/psgd.py"""
+import math
+import warnings
+
+import torch
+
+from ....core import Chainable, TensorTransform
+from ._psgd_utils import _initialize_lra_state_
+from .psgd import lift2single, precond_grad_lra, update_precond_lra_whiten
+
+# matches
+class PSGDLRAWhiten(TensorTransform):
+    """Low rank whitening preconditioner from Preconditioned Stochastic Gradient Descent (see https://github.com/lixilinx/psgd_torch)
+
+    Args:
+        rank (int, optional):
+            Preconditioner has a diagonal part and a low rank part, whose rank is decided by this setting. Defaults to 10.
+        init_scale (float | None, optional):
+            initial scale of the preconditioner. If None, determined based on a heuristic. Defaults to None.
+        lr_preconditioner (float, optional): learning rate of the preconditioner. Defaults to 0.1.
+        betaL (float, optional): EMA factor for the L-smoothness constant wrt Q. Defaults to 0.9.
+        damping (float, optional):
+            adds small noise to hessian-vector product when updating the preconditioner. Defaults to 1e-9.
+        grad_clip_max_norm (float, optional): clips norm of the update. Defaults to float("inf").
+        update_probability (float, optional): probability of updating preconditioner on each step. Defaults to 1.0.
+        concat_params (bool, optional):
+            if True, treats all parameters as concatenated to a single vector.
+            If False, each parameter is preconditioned separately. Defaults to True.
+        inner (Chainable | None, optional): preconditioning will be applied to output of this module. Defaults to None.
+
+    ###Examples:
+
+    Pure PSGD LRA:
+    ```py
+    optimizer = tz.Optimizer(
+        model.parameters(),
+        tz.m.LRAWhiten(),
+        tz.m.LR(1e-3),
+    )
+    ```
+
+    Momentum into preconditioner (whitens momentum):
+    ```py
+    optimizer = tz.Optimizer(
+        model.parameters(),
+        tz.m.EMA(0.9),
+        tz.m.LRAWhiten(),
+        tz.m.LR(1e-3),
+    )
+    ```
+
+    Updating the preconditioner from gradients and applying it to momentum:
+    ```py
+    optimizer = tz.Optimizer(
+        model.parameters(),
+        tz.m.LRAWhiten(inner=tz.m.EMA(0.9)),
+        tz.m.LR(1e-3),
+    )
+    ```
+
+    """
+    def __init__(
+        self,
+        rank: int = 10,
+        init_scale: float | None = None,
+        lr_preconditioner=0.1,
+        betaL=0.9,
+        damping=1e-9,
+        grad_clip_max_amp=float("inf"),
+        update_probability=1.0,
+
+        concat_params: bool = True,
+        inner: Chainable | None = None,
+    ):
+        defaults = locals().copy()
+        del defaults["inner"], defaults["self"]
+        super().__init__(defaults, concat_params=concat_params, inner=inner)
+
+    @torch.no_grad
+    def single_tensor_initialize(self, tensor, param, grad, loss, state, setting):
+        _initialize_lra_state_(tensor, state, setting)
+
+    @torch.no_grad
+    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
+
+        g = tensor.ravel().unsqueeze(1) # column vector
+
+        UVd = state["UVd"]
+        if UVd[2] is None: # initialize d on the fly
+            UVd[2] = (torch.mean(g**4) + setting["damping"]**4)**(-1/8) * torch.ones_like(g)
+
+        if torch.rand([]) < setting["update_probability"]:  # update preconditioner
+            update_precond_lra_whiten(
+                UVd=UVd,
+                Luvd=state["Luvd"],
+                g=g,
+                lr=setting["lr_preconditioner"],
+                betaL=setting["betaL"],
+                damping=setting["damping"],
+            )
+
+    @torch.no_grad
+    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
+
+        g = tensor.ravel().unsqueeze(1)
+        pre_grad = precond_grad_lra(UVd=state["UVd"], g=g)
+
+        # norm clipping
+        grad_clip_max_amp = setting["grad_clip_max_amp"]
+        if grad_clip_max_amp < float("inf"): # clip preconditioned gradient
+            amp = torch.sqrt(torch.mean(pre_grad * pre_grad))
+            if amp > grad_clip_max_amp:
+                pre_grad *= grad_clip_max_amp/amp
+
+        return pre_grad.view_as(tensor)

torchzero/modules/adaptive/rprop.py

@@ -304,7 +304,7 @@ class SignConsistencyMask(TensorTransform):
     GD that skips update for weights where gradient sign changed compared to previous gradient.
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Mul(tz.m.SignConsistencyMask()),
         tz.m.LR(1e-2)
@@ -334,7 +334,7 @@ class SignConsistencyLRs(TensorTransform):
 
     ```python
 
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Mul(tz.m.SignConsistencyLRs()),
         tz.m.LR(1e-2)

torchzero/modules/adaptive/sam.py

@@ -31,7 +31,7 @@ class SAM(Transform):
     SAM-SGD:
 
     ```py
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.SAM(),
         tz.m.LR(1e-2)
@@ -41,7 +41,7 @@ class SAM(Transform):
     SAM-Adam:
 
     ```
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.SAM(),
         tz.m.Adam(),
@@ -149,7 +149,7 @@ class ASAM(SAM):
     ASAM-SGD:
 
     ```py
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.ASAM(),
         tz.m.LR(1e-2)
@@ -159,7 +159,7 @@ class ASAM(SAM):
     ASAM-Adam:
 
     ```
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.ASAM(),
         tz.m.Adam(),

torchzero/modules/adaptive/shampoo.py

@@ -1,4 +1,4 @@
-from collections.abc import Sequence
+from collections.abc import Sequence, Iterable
 
 import numpy as np
 import torch
@@ -82,6 +82,31 @@ def _unmerge_small_dims(tensor: torch.Tensor, flat_sizes: Sequence[int] | None,
     tensor = tensor.unflatten(0, flat_sizes)
     return tensor.permute(*np.argsort(sort_idxs).tolist())
 
+def diagonal_memory(params: torch.nn.Module | torch.Tensor | Iterable[torch.Tensor]):
+    """computes number of parameters"""
+    if isinstance(params, torch.nn.Module): params = params.parameters()
+    if isinstance(params, torch.Tensor): params = [params,]
+    params = list(params)
+    return sum(p.numel() for p in params)
+
+def kronecker_memory(params: torch.nn.Module | torch.Tensor | Iterable[torch.Tensor], merge_small:bool=True, max_dim:int=10_000):
+    """computes total size of tensors required to store shampoo preconditioner"""
+    if isinstance(params, torch.nn.Module): params = params.parameters()
+    if isinstance(params, torch.Tensor): params = [params,]
+    params = list(params)
+
+    memory = 0
+    for p in params:
+        if merge_small:
+            p, _, _ = _merge_small_dims(p, max_dim)
+        for dim in p.size():
+            if dim > max_dim: memory += dim
+            else: memory += dim**2
+
+    return memory
+
+
+
 
 class Shampoo(TensorTransform):
     """Shampoo from Preconditioned Stochastic Tensor Optimization (https://arxiv.org/abs/1802.09568).
@@ -112,7 +137,7 @@ class Shampoo(TensorTransform):
     Shampoo grafted to Adam
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.GraftModules(
             direction = tz.m.Shampoo(),
@@ -125,7 +150,7 @@ class Shampoo(TensorTransform):
     Adam with Shampoo preconditioner
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Shampoo(beta=0.999, inner=tz.m.EMA(0.9)),
         tz.m.Debias(0.9, 0.999),

torchzero/modules/adaptive/soap.py

@@ -132,7 +132,7 @@ class SOAP(TensorTransform):
     SOAP:
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.SOAP(),
         tz.m.LR(1e-3)
@@ -141,7 +141,7 @@ class SOAP(TensorTransform):
     Stabilized SOAP:
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.SOAP(),
         tz.m.NormalizeByEMA(max_ema_growth=1.2),
@@ -156,7 +156,7 @@ class SOAP(TensorTransform):
         shampoo_beta: float | None = 0.95,
         precond_freq: int = 10,
         merge_small: bool = True,
-        max_dim: int = 10_000,
+        max_dim: int = 4096,
         precondition_1d: bool = True,
         eps: float = 1e-8,
         debias: bool = True,

torchzero/modules/adaptive/sophia_h.py

@@ -50,7 +50,7 @@ class SophiaH(Transform):
 
     ```python
 
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.SophiaH(),
         tz.m.LR(0.1)
@@ -63,7 +63,7 @@ class SophiaH(Transform):
 
     ```python
 
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.SophiaH(beta1=0, inner=tz.m.NAG(0.96)),
         tz.m.LR(0.1)

torchzero/modules/clipping/clipping.py

@@ -161,7 +161,7 @@ class ClipValue(TensorTransform):
 
     Gradient clipping:
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.ClipValue(1),
         tz.m.Adam(),
@@ -171,7 +171,7 @@ class ClipValue(TensorTransform):
 
     Update clipping:
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Adam(),
         tz.m.ClipValue(1),
@@ -211,7 +211,7 @@ class ClipNorm(TensorTransform):
 
     Gradient norm clipping:
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.ClipNorm(1),
         tz.m.Adam(),
@@ -221,7 +221,7 @@ class ClipNorm(TensorTransform):
 
     Update norm clipping:
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Adam(),
         tz.m.ClipNorm(1),
@@ -277,7 +277,7 @@ class Normalize(TensorTransform):
     Examples:
     Gradient normalization:
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Normalize(1),
         tz.m.Adam(),
@@ -288,7 +288,7 @@ class Normalize(TensorTransform):
     Update normalization:
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Adam(),
         tz.m.Normalize(1),
@@ -378,7 +378,7 @@ class Centralize(TensorTransform):
 
     Standard gradient centralization:
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Centralize(dim=0),
         tz.m.LR(1e-2),

torchzero/modules/conjugate_gradient/cg.py

@@ -7,7 +7,7 @@ from ...core import Chainable, TensorTransform
 
 from ...utils import TensorList, safe_dict_update_, unpack_dicts, unpack_states
 from ..quasi_newton.quasi_newton import HessianUpdateStrategy
-from ..functional import safe_clip
+from ..opt_utils import safe_clip
 
 
 class ConguateGradientBase(TensorTransform, ABC):
@@ -68,7 +68,7 @@ class ConguateGradientBase(TensorTransform, ABC):
         self.increment_counter("step", start=0)
 
         # initialize on first step
-        if self.global_state.get('stage', "first step") == "first update":
+        if self.global_state.get('stage', "first update") == "first update":
             g_prev, d_prev = unpack_states(states, tensors, 'g_prev', 'd_prev', cls=TensorList)
             d_prev.copy_(tensors)
             g_prev.copy_(tensors)

torchzero/modules/experimental/__init__.py

@@ -1,8 +1,13 @@
 """Those are various ideas of mine plus some other modules that I decided not to move to other sub-packages for whatever reason. This is generally less tested and shouldn't be used."""
+from .adanystrom import AdaNystrom
+from .common_directions_whiten import CommonDirectionsWhiten
 from .coordinate_momentum import CoordinateMomentum
+from .cubic_adam import CubicAdam, SubspaceCubicAdam
 from .curveball import CurveBall
+from .eigen_sr1 import EigenSR1
 
 # from dct import DCTProjection
+from .eigengrad import Eigengrad
 from .fft import FFTProjection
 from .gradmin import GradMin
 from .higher_order_newton import HigherOrderNewton

torchzero/modules/experimental/adanystrom.py

@@ -0,0 +1,258 @@
+# pylint: disable = non-ascii-name
+import torch
+
+from ...core import Chainable, TensorTransform
+from ...linalg import (
+    OrthogonalizeMethod,
+    orthogonalize,
+    regularize_eigh,
+    torch_linalg,
+)
+from ...linalg.linear_operator import Eigendecomposition
+from ..adaptive.lre_optimizers import LREOptimizerBase
+from .eigengrad import _eigengrad_update_state_, eigengrad_apply
+
+
+def weighted_eigen_plus_rank1_mm(
+    # A1 = Q1 @ diag(L1) @ Q1.T
+    L1: torch.Tensor,
+    Q1: torch.Tensor,
+
+    # K2 = v2 @ v2.T
+    v2: torch.Tensor,
+
+    # second matrix
+    B: torch.Tensor,
+
+    # weights
+    w1: float,
+    w2: float,
+
+) -> torch.Tensor:
+    """
+    Computes ``(w1 * A1 + w2 * A2) @ B``, where ``A1`` is an eigendecomposition, ``A2`` is symmetric rank 1.
+
+    Returns ``(n, k)``
+
+    Args:
+        L1 (torch.Tensor): eigenvalues of A1, shape ``(rank,)``.
+        Q1 (torch.Tensor): eigenvectors of A1, shape ``(n, rank)``.
+        v2 (torch.Tensor): vector such that ``v v^T = A2``, shape ``(n,)``.
+        B (torch.Tensor): shape ``(n, k)``.
+        w1 (float): weight for A1.
+        w2 (float): weight for A2.
+
+    """
+    # sketch A1
+    QTB = Q1.T @ B # (rank, k)
+    LQTB = L1.unsqueeze(1) * QTB  # (rank, k)
+    sketch1 = Q1 @ LQTB  # (n, k)
+
+    # skecth A2
+    vB = v2 @ B
+    sketch2 = v2.outer(vB)
+
+    return w1 * sketch1 + w2 * sketch2
+
+
+def adanystrom_update(
+    L1: torch.Tensor,
+    Q1: torch.Tensor,
+    v2: torch.Tensor,
+    w1: float,
+    w2: float,
+    oversampling_p: int,
+    rank: int,
+    eig_tol: float,
+    damping: float,
+    rdamping: float,
+    orthogonalize_method: OrthogonalizeMethod,
+
+) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    """computes the Nyström approximation of ``(w1 * A1 + w2 * A2)``,
+    where ``A1`` is an eigendecomposition, ``A2`` is symmetric rank 1.
+
+    returns L of shape ``(k, )`` and Q of shape ``(n, k)``.
+
+    Args:
+        L1 (torch.Tensor): eigenvalues of A1, shape ``(rank,)``.
+        Q1 (torch.Tensor): eigenvectors of A1, shape ``(n, rank)``.
+        v2 (torch.Tensor): vector such that ``v v^T = A2``, shape ``(n,)`` or ``(n, 1)``.
+        w1 (float): weight for A1.
+        w2 (float): weight for A2.
+    """
+    n = Q1.shape[0]
+    device = Q1.device
+    dtype = Q1.dtype
+    l = rank + oversampling_p
+
+    # gaussian test matrix
+    Omega = torch.randn(n, l, device=device, dtype=dtype)
+
+    # sketch
+    AOmega = weighted_eigen_plus_rank1_mm(L1, Q1, v2, Omega, w1, w2)
+    Q = orthogonalize(AOmega, orthogonalize_method)
+
+    AQ = weighted_eigen_plus_rank1_mm(L1, Q1, v2, Q, w1, w2)
+    QTAQ = Q.T @ AQ
+
+    W = (QTAQ + QTAQ.T) / 2.0
+
+    # compute new L and Q
+    try:
+        L_prime, S = torch_linalg.eigh(W, retry_float64=True)
+    except torch.linalg.LinAlgError:
+        return L1, Q1
+
+    L_prime, S = regularize_eigh(L=L_prime, Q=S, truncate=rank, tol=eig_tol, damping=damping, rdamping=rdamping)
+
+    if L_prime is None or S is None:
+        return L1, Q1
+
+    return L_prime, Q @ S
+
+
+# def adanystrom_update2(
+#     L1: torch.Tensor,
+#     Q1: torch.Tensor,
+#     v2: torch.Tensor,
+#     w1: float,
+#     w2: float,
+#     rank: int,
+# ):
+#     def A_mm(X):
+#         return weighted_eigen_plus_rank1_mm(L1=L1, Q1=Q1, v2=v2, B=X, w1=w1, w2=w2)
+
+#     return nystrom_approximation(A_mm, A_mm=A_mm, ndim=v2.numel(), rank=rank, device=L1.device, dtype=L1.dtype)
+
+class AdaNystrom(TensorTransform):
+    """Adagrad/RMSprop/Adam with Nyström-approximated covariance matrix.
+
+    Args:
+        rank (_type_): rank of Nyström approximation.
+        w1 (float, optional): weight of current covariance matrix. Defaults to 0.95.
+        w2 (float, optional): weight of new gradient in covariance matrix. Defaults to 0.05.
+        oversampling (int, optional): number of extra random vectors (top rank eigenvalues are kept). Defaults to 10.
+        eig_tol (float, optional):
+            removes eigenvalues this much smaller than largest eigenvalue when updating the preconditioner. Defaults to 1e-7.
+        damping (float, optional):
+            added to eigenvalues when updating the preconditioner. Defaults to 1e-8.
+        rdamping (float, optional):
+            added to eigenvalues when updating the preconditioner, relative to largest eigenvalue. Defaults to 0.
+        mm_tol (float, optional):
+            removes eigenvalues this much smaller than largest eigenvalue when computing the update. Defaults to 1e-7.
+        mm_truncate (int | None, optional):
+            uses top k eigenvalues to compute the update. Defaults to None.
+        mm_damping (float, optional):
+            added to eigenvalues when computing the update. Defaults to 1e-4.
+        mm_rdamping (float, optional):
+            added to eigenvalues when computing the update, relative to largest eigenvalue. Defaults to 0.
+        id_reg (float, optional):
+            multiplier to identity matrix added to preconditioner before computing update
+            If this value is given, solution from Nyström sketch-and-solve will be used to compute the update.
+            This value can't be too small (i.e. less than 1e-5) or the solver will be very unstable. Defaults to None.
+        concat_params (bool, optional):
+            whether to precondition all parameters at once if True, or each separately if False. Defaults to True.
+        update_freq (int, optional): update frequency. Defaults to 1.
+        inner (Chainable | None, optional): inner modules. Defaults to None.
+    """
+    def __init__(
+        self,
+        rank:int = 100,
+        beta=0.95,
+        oversampling: int = 10,
+        eig_tol: float | None = 1e-32,
+        damping: float = 0,
+        rdamping: float = 0,
+        mm_tol: float = 0,
+        mm_truncate: int | None = None,
+        mm_damping: float = 0,
+        mm_rdamping: float = 0,
+        id_reg: float | None = None,
+        orthogonalize_method: OrthogonalizeMethod = 'qr',
+        eigenbasis_optimizer: LREOptimizerBase | None = None,
+        orthogonalize_interval: int | None = 100,
+
+        concat_params: bool = True,
+        update_freq: int = 1,
+        inner: Chainable | None = None,
+    ):
+        defaults = locals().copy()
+        for k in ["self", "concat_params", "inner", "update_freq"]:
+            del defaults[k]
+
+        super().__init__(defaults, concat_params=concat_params, inner=inner, update_freq=update_freq)
+
+    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
+        state["step"] = state.get("step", 0) + 1
+        rank = setting["rank"]
+        device = tensor.device
+        dtype = tensor.dtype
+        beta = setting["beta"]
+
+        try:
+            if "L" not in state:
+                # use just tensor and zero L and Q with zero weight
+
+                L, Q = adanystrom_update(
+                    L1=torch.zeros(rank, device=device, dtype=dtype),
+                    Q1=torch.zeros((tensor.numel(), rank), device=device, dtype=dtype),
+                    v2=tensor.ravel(),
+                    w1=0,
+                    w2=1-beta,
+                    rank=rank,
+                    oversampling_p=setting["oversampling"],
+                    eig_tol=setting["eig_tol"],
+                    damping=setting["damping"],
+                    rdamping=setting["rdamping"],
+                    orthogonalize_method=setting["orthogonalize_method"],
+                )
+
+                state["L"] = state["L_reg"] = L
+                state["Q"] = state["Q_reg"] = Q
+
+            else:
+                L = state["L"]
+                Q = state["Q"]
+
+                w1 = beta
+                w2 = 1 - w1
+
+                # compute new factors (this function truncates them)
+                L_new, Q_new = adanystrom_update(
+                    L1=L,
+                    Q1=Q,
+                    v2=tensor.ravel(),
+                    w1=w1,
+                    w2=w2,
+                    rank=rank,
+                    oversampling_p=setting["oversampling"],
+                    eig_tol=setting["eig_tol"],
+                    damping=setting["damping"],
+                    rdamping=setting["rdamping"],
+                    orthogonalize_method=setting["orthogonalize_method"],
+                )
+
+                _eigengrad_update_state_(state=state, setting=setting, L_new=L_new, Q_new=Q_new)
+
+        except torch.linalg.LinAlgError:
+            pass
+
+    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
+        if "L_reg" not in state:
+            return tensor.clip(-0.1, 0.1)
+
+        if "eigenbasis_state" not in state:
+            state["eigenbasis_state"] = {}
+
+        return eigengrad_apply(
+            tensor=tensor,
+            L_reg = state["L_reg"],
+            Q_reg = state["Q_reg"],
+            beta = setting["beta"],
+            step = state["step"],
+            debias = True,
+            id_reg = setting["id_reg"],
+            eigenbasis_optimizer = setting["eigenbasis_optimizer"],
+            eigenbasis_state = state["eigenbasis_state"]
+        )