torchzero 0.3.10__py3-none-any.whl → 0.3.11__py3-none-any.whl

torchzero/modules/optimizers/esgd.py

@@ -0,0 +1,171 @@
+import math
+from collections.abc import Callable
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, Target, Transform, apply_transform
+from ...utils import NumberList, TensorList, as_tensorlist
+
+
+def esgd_(
+    tensors_: TensorList,
+    D: TensorList | None,
+    D_sq_acc_: TensorList,
+    damping: float | NumberList,
+    update_freq: int,
+    step: int,
+    i: int,
+):
+    # update preconditioner
+    if step % update_freq == 0:
+        assert D is not None
+        D_sq_acc_.addcmul_(D, D)
+        i += 1
+    else:
+        assert D is None
+
+    denom = (D_sq_acc_ / max(i, 1)).sqrt_().add_(damping)
+    return tensors_.div_(denom), i
+
+
+class ESGD(Module):
+    """Equilibrated Gradient Descent (https://arxiv.org/abs/1502.04390)
+
+    This is similar to Adagrad, but the accumulates squared randomized hessian diagonal estimates instead of squared gradients.
+
+    .. note::
+        In most cases Adagrad should be the first module in the chain because it relies on autograd. Use the :code:`inner` argument if you wish to apply Adagrad preconditioning to another module's output.
+
+    .. note::
+        If you are using gradient estimators or reformulations, set :code:`hvp_method` to "forward" or "central".
+
+    .. note::
+        This module requires a closure passed to the optimizer step,
+        as it needs to re-evaluate the loss and gradients for calculating HVPs.
+        The closure must accept a ``backward`` argument (refer to documentation).
+
+    Args:
+        damping (float, optional): added to denominator for stability. Defaults to 1e-4.
+        update_freq (int, optional):
+            frequency of updating hessian diagonal estimate via a hessian-vector product.
+            This value can be increased to reduce computational cost. Defaults to 20.
+        hvp_method (str, optional):
+            Determines how Hessian-vector products are evaluated.
+
+            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
+              This requires creating a graph for the gradient.
+            - ``"forward"``: Use a forward finite difference formula to
+              approximate the HVP. This requires one extra gradient evaluation.
+            - ``"central"``: Use a central finite difference formula for a
+              more accurate HVP approximation. This requires two extra
+              gradient evaluations.
+            Defaults to "autograd".
+        h (float, optional): finite difference step size if :code:`hvp_method` is "forward" or "central". Defaults to 1e-3.
+        n_samples (int, optional):
+            number of hessian-vector products with random vectors to evaluate each time when updating
+            the preconditioner. Larger values may lead to better hessian diagonal estimate. Defaults to 1.
+        seed (int | None, optional): seed for random vectors. Defaults to None.
+        inner (Chainable | None, optional):
+            Inner module. If this is specified, operations are performed in the following order.
+            1. compute hessian diagonal estimate.
+            2. pass inputs to :code:`inner`.
+            3. momentum and preconditioning are applied to the ouputs of :code:`inner`.
+
+    Examples:
+        Using ESGD:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.ESGD(),
+                tz.m.LR(0.1)
+            )
+
+        ESGD preconditioner can be applied to any other module by passing it to the :code:`inner` argument. Here is an example of applying
+        ESGD preconditioning to nesterov momentum (:code:`tz.m.NAG`):
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.ESGD(beta1=0, inner=tz.m.NAG(0.9)),
+                tz.m.LR(0.1)
+            )
+
+    """
+    def __init__(
+        self,
+        damping: float = 1e-4,
+        update_freq: int = 20,
+        hvp_method: Literal['autograd', 'forward', 'central'] = 'autograd',
+        fd_h: float = 1e-3,
+        n_samples = 1,
+        seed: int | None = None,
+        inner: Chainable | None = None
+    ):
+        defaults = dict(damping=damping, update_freq=update_freq, hvp_method=hvp_method, n_samples=n_samples, fd_h=fd_h, seed=seed)
+        super().__init__(defaults)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, var):
+        params = var.params
+        settings = self.settings[params[0]]
+        hvp_method = settings['hvp_method']
+        fd_h = settings['fd_h']
+        update_freq = settings['update_freq']
+        n_samples = settings['n_samples']
+
+        seed = settings['seed']
+        generator = None
+        if seed is not None:
+            if 'generator' not in self.global_state:
+                self.global_state['generator'] = torch.Generator(params[0].device).manual_seed(seed)
+            generator = self.global_state['generator']
+
+        damping = self.get_settings(params, 'damping', cls=NumberList)
+        D_sq_acc = self.get_state(params, 'D_sq_acc', cls=TensorList)
+        i = self.global_state.get('i', 0)
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        closure = var.closure
+        assert closure is not None
+
+        D = None
+        if step % update_freq == 0:
+
+            rgrad=None
+            for j in range(n_samples):
+                u = [torch.randn(p.size(), generator=generator, device=p.device, dtype=p.dtype) for p in params]
+
+                Hvp, rgrad = self.Hvp(u, at_x0=True, var=var, rgrad=rgrad, hvp_method=hvp_method,
+                                     h=fd_h, normalize=True, retain_grad=j < n_samples-1)
+
+                if D is None: D = Hvp
+                else: torch._foreach_add_(D, Hvp)
+
+            assert D is not None
+            if n_samples > 1: torch._foreach_div_(D, n_samples)
+
+            D = TensorList(D)
+
+        update = var.get_update()
+        if 'inner' in self.children:
+            update = apply_transform(self.children['inner'], tensors=update, params=params, grads=var.grad, var=var)
+
+        var.update, self.global_state['i'] = esgd_(
+            tensors_=TensorList(update),
+            D=TensorList(D) if D is not None else None,
+            D_sq_acc_=D_sq_acc,
+            damping=damping,
+            update_freq=update_freq,
+            step=step,
+            i=i,
+        )
+        return var

torchzero/modules/{experimental/spectral.py → optimizers/ladagrad.py}

@@ -1,55 +1,53 @@
-from abc import ABC, abstractmethod
-import math
 from collections import deque
 from typing import Literal, Any
-import itertools
+import warnings
 
 import torch
 from ...core import Chainable, TensorwiseTransform
-from ...utils.linalg.matrix_funcs import matrix_power_eigh
-from ...utils.linalg.svd import randomized_svd
-from ...utils.linalg.qr import qr_householder
 
-def spectral_update(history, damping, rdamping, true_damping: bool):
-    M_hist = torch.stack(tuple(history), dim=1)
-    device = M_hist.device
-    M_hist = M_hist.cuda()
+def lm_adagrad_update(history: deque[torch.Tensor], damping, rdamping):
+    M = torch.stack(tuple(history), dim=1)# / len(history)
+    MTM = M.T @ M
+    if damping != 0:
+        MTM.add_(torch.eye(MTM.size(0), device=MTM.device, dtype=MTM.dtype).mul_(damping))
 
     try:
-        U, S, _ = torch.linalg.svd(M_hist, full_matrices=False, driver='gesvda') # pylint:disable=not-callable
-        U = U.to(device); S = S.to(device)
+        L, Q = torch.linalg.eigh(MTM) # pylint:disable=not-callable
 
-        if damping != 0 or rdamping != 0:
-            if rdamping != 0: rdamping *= torch.linalg.vector_norm(S) # pylint:disable=not-callable
-            Iu = damping + rdamping
-            if true_damping:
-                S.pow_(2)
-                Iu **= 2
-            S.add_(Iu)
-            if true_damping: S.sqrt_()
+        tol = torch.finfo(M.dtype).eps * L.amax() # remove small eigenvalues
+        indices = L > tol
+        L = L[indices]
+        Q = Q[:, indices]
 
-        return U, 1/S
+        U = (M @ Q) * L.rsqrt()
+
+        if rdamping != 0:
+            rdamping *= torch.linalg.vector_norm(L) # pylint:disable=not-callable
+            L.add_(rdamping)
+
+        return U, L
 
     except torch.linalg.LinAlgError:
         return None, None
 
-def spectral_apply(g: torch.Tensor, U: torch.Tensor, S_inv: torch.Tensor):
-    Utg = (U.T @ g)*S_inv
-    return U @ Utg
-
+def lm_adagrad_apply(g: torch.Tensor, U: torch.Tensor, L: torch.Tensor):
+    Z = U.T @ g
+    return (U * L.rsqrt()) @ Z
 
 def maybe_lerp_(state_: dict, beta: float | None, key, value: Any):
     if (key not in state_) or (beta is None) or (not isinstance(value, torch.Tensor)): state_[key] = value
     else:
-        if state_[key].shape != value.shape: state_[key] = value
+        if state_[key] is None or state_[key].shape != value.shape: state_[key] = value
         else: state_[key].lerp_(value, 1-beta)
 
-class SpectralPreconditioner(TensorwiseTransform):
+class LMAdagrad(TensorwiseTransform):
     """
-    The update rule is to stack recent gradients into M, compute U, S <- SVD(M), then calculate U (Uᵀg)/S.
-    This is equivalent to full matrix Adagrad with accumulator initialized to zeros,
-    except only recent :code:`history_size` gradients are used.
-    However this doesn't require N^2 memory and is computationally less expensive than Shampoo.
+    Limited-memory full matrix Adagrad.
+
+    The update rule is to stack recent gradients into M, compute U, S <- SVD(M), then calculate update as U S^-1 Uᵀg.
+    But it uses eigendecomposition on MᵀM to get U and S^2 because that is faster when you don't neeed V.
+
+    This is equivalent to full-matrix Adagrad on recent gradients.
 
     Args:
         history_size (int, optional): number of past gradients to store. Defaults to 10.
@@ -60,55 +58,84 @@ class SpectralPreconditioner(TensorwiseTransform):
             order=2 means gradient differences are used in place of gradients. Higher order uses higher order differences. Defaults to 1.
         true_damping (bool, optional):
             If True, damping is added to squared singular values to mimic Adagrad. Defaults to True.
+        eigh (bool, optional): uses a more efficient way to calculate U and S. Defaults to True.
         U_beta (float | None, optional): momentum for U (too unstable, don't use). Defaults to None.
-        S_beta (float | None, optional): momentum for 1/S (too unstable, don't use). Defaults to None.
+        S_beta (float | None, optional): momentum for S (too unstable, don't use). Defaults to None.
         interval (int, optional): Interval between gradients that are added to history (2 means every second gradient is used). Defaults to 1.
-        concat_params (bool, optional): if True, treats all parameters as a single vector, meaning it will also whiten inter-parameters. Defaults to False.
-        normalize (bool, optional): whether to normalize gradients, this doesn't work well so don't use it. Defaults to False.
-        centralize (bool, optional): whether to centralize gradients, this doesn't work well so don't use it. Defaults to False.
+        concat_params (bool, optional): if True, treats all parameters as a single vector, meaning it will also whiten inter-parameters. Defaults to True.
         inner (Chainable | None, optional): preconditioner will be applied to output of this module. Defaults to None.
+
+    Examples:
+        Limited-memory Adagrad
+
+        .. code-block:: python
+
+            optimizer = tz.Modular(
+                model.parameters(),
+                tz.m.LMAdagrad(),
+                tz.m.LR(0.1)
+            )
+
+        Adam with L-Adagrad preconditioner (for debiasing second beta is 0.999 arbitrarily)
+
+        .. code-block:: python
+
+            optimizer = tz.Modular(
+                model.parameters(),
+                tz.m.LMAdagrad(inner=tz.m.EMA()),
+                tz.m.Debias(0.9, 0.999),
+                tz.m.LR(0.01)
+            )
+
+        Stable Adam with L-Adagrad preconditioner (this is what I would recommend)
+
+        .. code-block:: python
+
+            optimizer = tz.Modular(
+                model.parameters(),
+                tz.m.LMAdagrad(inner=tz.m.EMA()),
+                tz.m.Debias(0.9, 0.999),
+                tz.m.ClipNormByEMA(max_ema_growth=1.2),
+                tz.m.LR(0.01)
+            )
+
+    Reference:
+        Agarwal N. et al. Efficient full-matrix adaptive regularization //International Conference on Machine Learning. – PMLR, 2019. – С. 102-110.
     """
 
     def __init__(
         self,
-        history_size: int = 10,
+        history_size: int = 100,
         update_freq: int = 1,
         damping: float = 1e-4,
         rdamping: float = 0,
         order: int = 1,
         true_damping: bool = True,
         U_beta: float | None = None,
-        S_beta: float | None = None,
+        L_beta: float | None = None,
         interval: int = 1,
-        concat_params: bool = False,
-        normalize: bool=False,
-        centralize:bool = False,
+        concat_params: bool = True,
         inner: Chainable | None = None,
     ):
         # history is still updated each step so Precondition's update_freq has different meaning
-        defaults = dict(history_size=history_size, update_freq=update_freq, damping=damping, rdamping=rdamping, true_damping=true_damping, order=order, U_beta=U_beta, S_beta=S_beta, normalize=normalize, centralize=centralize)
+        defaults = dict(history_size=history_size, update_freq=update_freq, damping=damping, rdamping=rdamping, true_damping=true_damping, order=order, U_beta=U_beta, L_beta=L_beta)
         super().__init__(defaults, uses_grad=False, concat_params=concat_params, inner=inner, update_freq=interval)
 
     @torch.no_grad
-    def update_tensor(self, tensor, param, grad, loss, state, settings):
-        order = settings['order']
-        history_size = settings['history_size']
-        update_freq = settings['update_freq']
-        damping = settings['damping']
-        rdamping = settings['rdamping']
-        true_damping = settings['true_damping']
-        U_beta = settings['U_beta']
-        S_beta = settings['S_beta']
-        normalize = settings['normalize']
-        centralize = settings['centralize']
+    def update_tensor(self, tensor, param, grad, loss, state, setting):
+        order = setting['order']
+        history_size = setting['history_size']
+        update_freq = setting['update_freq']
+        damping = setting['damping']
+        rdamping = setting['rdamping']
+        U_beta = setting['U_beta']
+        L_beta = setting['L_beta']
 
         if 'history' not in state: state['history'] = deque(maxlen=history_size)
         history = state['history']
 
         if order == 1:
             t = tensor.clone().view(-1)
-            if centralize: t -= t.mean()
-            if normalize: t /= torch.linalg.vector_norm(t).clip(min=1e-8) # pylint:disable=not-callable
             history.append(t)
         else:
 
@@ -122,42 +149,35 @@ class SpectralPreconditioner(TensorwiseTransform):
                     state[f'prev_g_{i}'] = cur_g
                     break
 
-                s_k = cur_p - state[f'prev_p_{i}']
-                y_k = cur_g - state[f'prev_g_{i}']
+                s = cur_p - state[f'prev_p_{i}']
+                y = cur_g - state[f'prev_g_{i}']
                 state[f'prev_p_{i}'] = cur_p
                 state[f'prev_g_{i}'] = cur_g
-                cur_p = s_k
-                cur_g = y_k
+                cur_p = s
+                cur_g = y
 
                 if i == order - 1:
-                    if centralize: cur_g = cur_g - cur_g.mean()
-                    if normalize: cur_g = cur_g / torch.linalg.vector_norm(cur_g).clip(min=1e-8) # pylint:disable=not-callable
-                    else: cur_g = cur_g / torch.linalg.norm(cur_p).clip(min=1e-8) # pylint:disable=not-callable
+                    cur_g = cur_g / torch.linalg.norm(cur_p).clip(min=1e-8) # pylint:disable=not-callable
                     history.append(cur_g.view(-1))
 
         step = state.get('step', 0)
         if step % update_freq == 0 and len(history) != 0:
-            U, S_inv = spectral_update(history, damping=damping, rdamping=rdamping, true_damping=true_damping)
+            U, L = lm_adagrad_update(history, damping=damping, rdamping=rdamping)
             maybe_lerp_(state, U_beta, 'U', U)
-            maybe_lerp_(state, S_beta, 'S_inv', S_inv)
+            maybe_lerp_(state, L_beta, 'L', L)
 
         if len(history) != 0:
             state['step'] = step + 1 # do not increment if no history (gathering s_ks and y_ks)
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, settings):
-        history_size = settings['history_size']
-
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
         U = state.get('U', None)
         if U is None:
             # make a conservative step to avoid issues due to different GD scaling
             return tensor.clip_(-0.1, 0.1) # pyright:ignore[reportArgumentType]
 
-        S_inv = state['S_inv']
-        update = spectral_apply(tensor.view(-1), U, S_inv).view_as(tensor)
+        L = state['L']
+        update = lm_adagrad_apply(tensor.view(-1), U, L).view_as(tensor)
 
-        n = len(state['history'])
-        mh = min(history_size, 10)
-        if n <= mh: update.mul_(n/mh)
         return update
 

torchzero/modules/optimizers/lion.py

@@ -28,7 +28,7 @@ class Lion(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         beta1, beta2 = unpack_dicts(settings, 'beta1', 'beta2', cls=NumberList)
         exp_avg = unpack_states(states, tensors, 'ema', cls=TensorList)
         return lion_(TensorList(tensors),exp_avg,beta1,beta2)

torchzero/modules/optimizers/mars.py

@@ -0,0 +1,91 @@
+from operator import itemgetter
+from functools import partial
+
+import torch
+
+from ...core import Module, Target, Transform, apply_transform, Chainable
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+from ..functional import (
+    debias, debiased_step_size,
+    ema_,
+    sqrt_ema_sq_,
+)
+from ..step_size.lr import lazy_lr
+from ..momentum.experimental import sqrt_nag_ema_sq_
+from ..momentum.momentum import nag_
+
+
+def mars_correction_(
+    tensors_: TensorList,
+    prev_: TensorList,
+    beta: float | NumberList,
+    scaling: float | NumberList,
+    max_norm: float | NumberList |  None,
+):
+    dg = (tensors_ - prev_).mul_(scaling * beta / (1-beta))
+    prev_.copy_(tensors_)
+
+    c = tensors_.add_(dg)
+    if max_norm is not None:
+        c.clip_norm_(max=max_norm, tensorwise=False)
+
+    return c
+
+class MARSCorrection(Transform):
+    """MARS variance reduction correction.
+
+    Place any other momentum-based optimizer after this,
+    make sure :code:`beta` parameter matches with momentum in the optimizer.
+
+    Args:
+        beta (float, optional): use the same beta as you use in the momentum module. Defaults to 0.9.
+        scaling (float, optional): controls the scale of gradient correction in variance reduction. Defaults to 0.025.
+        max_norm (float, optional): clips norm of corrected gradients, None to disable. Defaults to 1.
+
+    Examples:
+        Mars-AdamW
+
+        .. code-block:: python
+
+            optimizer = tz.Modular(
+                model.parameters(),
+                tz.m.MARSCorrection(beta=0.95),
+                tz.m.Adam(beta1=0.95, beta2=0.99),
+                tz.m.WeightDecay(1e-3),
+                tz.m.LR(0.1)
+            )
+
+        Mars-Lion
+
+        .. code-block:: python
+
+            optimizer = tz.Modular(
+                model.parameters(),
+                tz.m.MARSCorrection(beta=0.9),
+                tz.m.Lion(beta1=0.9),
+                tz.m.LR(0.1)
+            )
+
+    """
+    def __init__(
+        self,
+        beta: float = 0.9,
+        scaling: float = 0.025,
+        max_norm: float | None = 1,
+    ):
+        defaults=dict(beta=beta, scaling=scaling, max_norm=max_norm)
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        prev = unpack_states(states, tensors, 'prev', init=tensors, cls=TensorList)
+        beta, scaling = unpack_dicts(settings, 'beta', 'scaling', cls=NumberList)
+        max_norm = settings[0]['max_norm']
+
+        return mars_correction_(
+            tensors_=TensorList(tensors),
+            prev_=prev,
+            beta=beta,
+            scaling=scaling,
+            max_norm=max_norm,
+        )