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

torchzero/modules/conjugate_gradient/cg.py

@@ -3,21 +3,14 @@ from typing import Literal
 
 import torch
 
-from ...core import (
-    Chainable,
-    Modular,
-    Module,
-    Transform,
-    Var,
-    apply_transform,
-)
-from ...utils import TensorList, as_tensorlist, unpack_dicts, unpack_states
-from ..line_search import LineSearchBase
+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(Transform, ABC):
+class ConguateGradientBase(TensorTransform, ABC):
     """Base class for conjugate gradient methods. The only difference between them is how beta is calculated.
 
     This is an abstract class, to use it, subclass it and override `get_beta`.
@@ -52,13 +45,8 @@ class ConguateGradientBase(Transform, ABC):
     """
     def __init__(self, defaults, clip_beta: bool, restart_interval: int | None | Literal['auto'], inner: Chainable | None = None):
         if defaults is None: defaults = {}
-        defaults['restart_interval'] = restart_interval
-        defaults['clip_beta'] = clip_beta
-        super().__init__(defaults, uses_grad=False)
-
-        if inner is not None:
-            self.set_child('inner', inner)
-
+        safe_dict_update_(defaults, dict(restart_interval=restart_interval, clip_beta=clip_beta))
+        super().__init__(defaults, inner=inner)
 
     def reset_for_online(self):
         super().reset_for_online()
@@ -74,40 +62,38 @@ class ConguateGradientBase(Transform, ABC):
         """returns beta"""
 
     @torch.no_grad
-    def update_tensors(self, tensors, params, grads, loss, states, settings):
-        tensors = as_tensorlist(tensors)
-        params = as_tensorlist(params)
-
-        step = self.global_state.get('step', 0) + 1
-        self.global_state['step'] = step
+    def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
+        tensors = TensorList(tensors)
+        params = TensorList(params)
+        self.increment_counter("step", start=0)
 
         # initialize on first step
-        if self.global_state.get('stage', 0) == 0:
+        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)
             self.initialize(params, tensors)
-            self.global_state['stage'] = 1
+            self.global_state['stage'] = "first apply"
 
         else:
             # if `update_tensors` was called multiple times before `apply_tensors`,
             # stage becomes 2
-            self.global_state['stage'] = 2
+            self.global_state['stage'] = "initialized"
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
-        tensors = as_tensorlist(tensors)
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
+        tensors = TensorList(tensors)
         step = self.global_state['step']
 
-        if 'inner' in self.children:
-            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads))
+        assert self.global_state['stage'] != "first update"
 
-        assert self.global_state['stage'] != 0
-        if self.global_state['stage'] == 1:
-            self.global_state['stage'] = 2
+        # on 1st apply we don't have previous gradients
+        # so just return tensors
+        if self.global_state['stage'] == "first apply":
+            self.global_state['stage'] = "initialized"
             return tensors
 
-        params = as_tensorlist(params)
+        params = TensorList(params)
         g_prev, d_prev = unpack_states(states, tensors, 'g_prev', 'd_prev', cls=TensorList)
 
         # get beta
@@ -119,10 +105,13 @@ class ConguateGradientBase(Transform, ABC):
         dir = tensors.add_(d_prev.mul_(beta))
         d_prev.copy_(dir)
 
-        # resetting
+        # resetting every `reset_interval` steps, use step+1 to not reset on 1st step
+        # so if reset_interval=2, then 1st step collects g_prev and d_prev, then
+        # two steps will happen until reset.
         restart_interval = settings[0]['restart_interval']
         if restart_interval == 'auto': restart_interval = tensors.global_numel() + 1
-        if restart_interval is not None and step % restart_interval == 0:
+
+        if restart_interval is not None and (step + 1) % restart_interval == 0:
             self.state.clear()
             self.global_state.clear()
 

torchzero/modules/experimental/__init__.py

@@ -1,19 +1,20 @@
 """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
 from .l_infinity import InfinityNormTrustRegion
-from .momentum import (
-    CoordinateMomentum,
-    NesterovEMASquared,
-    PrecenteredEMASquared,
-    SqrtNesterovEMASquared,
-)
 from .newton_solver import NewtonSolver
 from .newtonnewton import NewtonNewton
 from .reduce_outward_lr import ReduceOutwardLR
 from .scipy_newton_cg import ScipyNewtonCG
+from .spsa1 import SPSA1
 from .structural_projections import BlockPartition, TensorizeProjection

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"]
+        )

torchzero/modules/experimental/common_directions_whiten.py

@@ -0,0 +1,142 @@
+from collections import deque
+from typing import Literal
+
+import torch
+
+from torchzero.core import Chainable, TensorTransform
+from torchzero.linalg import matrix_power_eigh, torch_linalg, orthogonalize, OrthogonalizeMethod, regularize_eigh
+from torchzero.utils import TensorList, vec_to_tensors_
+
+
+def update_subspace_preconditioner_(
+    grad: torch.Tensor, # store grads and basis as vectors for matmul
+    basis: torch.Tensor, # ndim, k
+    accumulator_: torch.Tensor, # k, k
+    beta: float | None,
+):
+    projected = basis.T @ grad # k
+    outer = torch.outer(projected, projected)
+
+    if beta is None: accumulator_.add_(outer)
+    else: accumulator_.lerp_(outer, 1-beta)
+
+# yeah so I can also run subspace opts in this basis
+def apply_subspace_preconditioner(
+    tensor: torch.Tensor,
+    basis: torch.Tensor, # ndim, k
+    accumulator: torch.Tensor,
+    tol: float,
+    truncate: int | None,
+    damping: float,
+    rdamping: float,
+):
+    L, Q = torch_linalg.eigh(accumulator, retry_float64=True)
+    L, Q = regularize_eigh(L=L, Q=Q, truncate=truncate, tol=tol, damping=damping, rdamping=rdamping)
+
+    if L is None or Q is None:
+        return tensor.clip(-0.1, 0.1)
+
+    preconditioner = (Q * L.rsqrt().unsqueeze(-2)) @ Q.mH
+
+    tensor_projected = basis.T @ tensor # k
+    update_projected = preconditioner @ tensor_projected # k
+    return basis @ update_projected # d
+
+
+class CommonDirectionsWhiten(TensorTransform):
+    """Whitens in subspace spanned by history of gradient differences.
+
+    Args:
+        beta - for preconditioner itself in the basis.
+        basis_beta - how much basis is allowed to change.
+    """
+
+    def __init__(
+        self,
+        k: int = 100,
+        beta: float | None = 0.95,
+        basis_beta=0.95,
+        tol: float = 1e-7,
+        truncate: int | None = None,
+        damping: float = 1e-4,
+        rdamping: float = 0,
+        basis_type: Literal["gradients", "differences"] = "differences",
+        orthogonalize_method: OrthogonalizeMethod | None = 'newtonschulz',
+
+        concat_params: bool = True,
+        inner: Chainable | None = None,
+    ):
+        defaults = locals().copy()
+        for key in ["self", "inner", "concat_params"]:
+            del defaults[key]
+
+        super().__init__(defaults, concat_params=concat_params, inner=inner)
+
+    @torch.no_grad
+    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
+        g = tensor.ravel()
+        k = setting['k']
+        beta = setting['beta']
+        basis_beta = setting['basis_beta']
+        step = state.get("step", 0)
+        state["step"] = step + 1
+
+        # initialize history
+        if 'history' not in state:
+            state['history'] = deque(maxlen=k)
+            state['accumulator'] = torch.eye(k, device=g.device, dtype=g.dtype)
+            state['basis'] = torch.zeros(g.numel(), k, device=g.device, dtype=g.dtype)
+
+        history: deque = state['history']
+        accumulator = state['accumulator']
+        basis = state['basis']
+        history.append(g)
+
+        # stack history to new basis term, if history isn't full, fill with random vecs
+        if len(history) < k:
+            basis_t = torch.randn(g.numel(), k, device=g.device, dtype=g.dtype)
+            history_basis = torch.stack(tuple(history), -1)
+            basis_t[:, -len(history):] = history_basis
+
+        else:
+            basis_t = torch.stack(tuple(history), -1)
+
+        # in this case basis uses differences in gradients except last entry is the gradient
+        if setting["basis_type"] == "differences":
+            basis_t[:,:-1] = basis_t[:, :-1] - basis_t[:, 1:]
+
+        # normalize or orthonormalize new basis term
+        if setting["orthogonalize_method"] is not None:
+            basis_t = orthogonalize(basis_t, method = setting["orthogonalize_method"])
+        else:
+            basis_t = (basis_t - basis_t.mean()) / basis_t.std().clip(min=torch.finfo(g.dtype).tiny * 2)
+
+        # lerp basis
+        basis.lerp_(basis_t, 1-basis_beta)
+        basis = basis /  (1 - basis_beta ** (step+1)) # correct bias on basis EMA
+        update_subspace_preconditioner_(g, basis, accumulator, beta)
+
+    @torch.no_grad
+    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
+        g = tensor.ravel()
+
+        basis = state['basis']
+        accumulator = state['accumulator']
+        step = state["step"]
+        accumulator = accumulator / (1 - setting["beta"] ** (step+1)) # correct bias on accumulator EMA
+
+        try:
+            preconditioned = apply_subspace_preconditioner(
+                g,
+                basis,
+                accumulator,
+                tol=setting["tol"],
+                truncate=setting["truncate"],
+                damping=setting["damping"],
+                rdamping=setting["rdamping"],
+            )
+        except torch.linalg.LinAlgError:
+            preconditioned = g.clip(-0.1, 0.1)
+
+        return preconditioned.view_as(tensor)
+

torchzero/modules/experimental/coordinate_momentum.py

@@ -0,0 +1,36 @@
+import torch
+
+from ...core import TensorTransform
+from ...utils import NumberList, TensorList, unpack_states
+
+
+def coordinate_momentum_(
+    tensors: TensorList,
+    velocity_: TensorList,
+    p: float | NumberList,
+):
+    """
+    sets `velocity_` to p% random values from `tensors`.
+
+    Returns `velocity_`
+    """
+    mask = tensors.bernoulli_like(p).as_bool()
+    velocity_.masked_set_(mask, tensors)
+    return velocity_
+
+
+class CoordinateMomentum(TensorTransform):
+    """Maintains a momentum buffer, on each step each value in the buffer has ``p`` chance to be updated with the new value.
+
+    Args:
+        p (float, optional): _description_. Defaults to 0.1.
+    """
+    def __init__(self, p: float = 0.1):
+        defaults = dict(p=p)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
+        p = NumberList(s['p'] for s in settings)
+        velocity = unpack_states(states, tensors, 'velocity', cls=TensorList)
+        return coordinate_momentum_(TensorList(tensors), velocity_=velocity, p=p).clone()