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

torchzero/modules/adaptive/esgd.py

@@ -1,49 +1,20 @@
-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
+from ...core import Chainable, HVPMethod, Transform
+from ...utils import Distributions, NumberList, TensorList, unpack_dicts, unpack_states
 
 
-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):
+class ESGD(Transform):
     """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.
+    Notes:
+        - In most cases ESGD should be the first module in the chain because it relies on autograd. Use the ``inner`` argument if you wish to apply ESGD 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).
+        - 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.
@@ -51,17 +22,17 @@ class ESGD(Module):
             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".
-        fd_h (float, optional): finite difference step size if :code:`hvp_method` is "forward" or "central". Defaults to 1e-3.
+            Determines how hessian-vector products are computed.
+
+            - ``"batched_autograd"`` - uses autograd with batched hessian-vector products. If a single hessian-vector is evaluated, equivalent to ``"autograd"``. Faster than ``"autograd"`` but uses more memory.
+            - ``"autograd"`` - uses autograd hessian-vector products. If multiple hessian-vector products are evaluated, uses a for-loop. Slower than ``"batched_autograd"`` but uses less memory.
+            - ``"fd_forward"`` - uses gradient finite difference approximation with a less accurate forward formula which requires one extra gradient evaluation per hessian-vector product.
+            - ``"fd_central"`` - uses gradient finite difference approximation with a more accurate central formula which requires two gradient evaluations per hessian-vector product.
+
+            Defaults to ``"autograd"``.
+        h (float, optional):
+            The step size for finite difference if ``hvp_method`` is
+            ``"fd_forward"`` or ``"fd_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.
@@ -72,100 +43,108 @@ class ESGD(Module):
             2. pass inputs to :code:`inner`.
             3. momentum and preconditioning are applied to the ouputs of :code:`inner`.
 
-    Examples:
-        Using ESGD:
-
-        .. code-block:: python
+    ### Examples:
 
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.ESGD(),
-                tz.m.LR(0.1)
-            )
+    Using ESGD:
+```python
 
-        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`):
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.ESGD(),
+        tz.m.LR(0.1)
+    )
+    ```
 
-        .. code-block:: python
+    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`):
 
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.ESGD(beta1=0, inner=tz.m.NAG(0.9)),
-                tz.m.LR(0.1)
-            )
+    ```python
+    opt = tz.Optimizer(
+        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,
+        distribution: Distributions = 'gaussian',
+        hvp_method: HVPMethod = 'autograd',
+        h: float = 1e-3,
         n_samples = 1,
+        zHz: bool = False,
         seed: int | None = None,
-        inner: Chainable | None = None
+        beta: float | None = None,
+        beta_debias: bool = True,
+
+        inner: Chainable | None = None,
+        Hz_sq_acc_tfm: 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)
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner'], defaults["Hz_sq_acc_tfm"]
+        super().__init__(defaults, inner=inner)
 
-        if inner is not None:
-            self.set_child('inner', inner)
+        self.set_child("Hz_sq_acc", Hz_sq_acc_tfm)
 
     @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
+    def update_states(self, objective, states, settings):
+        params = objective.params
+
+        fs = settings[0]
+        update_freq = fs['update_freq']
+
+        # ------------------------------- accumulate Hz ------------------------------ #
+        step = self.increment_counter("step", start=0)
+
         if step % update_freq == 0:
+            self.increment_counter("num_Hzs", start=1)
+
+            Hz, _ = objective.hutchinson_hessian(
+                rgrad = None,
+                at_x0 = True,
+                n_samples = fs['n_samples'],
+                distribution = fs['distribution'],
+                hvp_method = fs['hvp_method'],
+                h = fs['h'],
+                zHz = fs["zHz"], # default is False, so it returns Hz, not z⊙Hz
+                generator = self.get_generator(params[0].device, fs["seed"]),
+            )
 
-            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]
+            Hz = TensorList(Hz)
+            Hz_sq_acc = unpack_states(states, params, 'Hz_sq_acc', cls=TensorList)
+
+            beta = fs["beta"]
+            if beta is None:
+                Hz_sq_acc.addcmul_(Hz, Hz)
+
+            else:
+                Hz_sq_acc.mul_(beta).addcmul_(Hz, Hz, value=1-beta)
+
+    @torch.no_grad
+    def apply_states(self, objective, states, settings):
+        tensors = TensorList(objective.get_updates())
+        Hz_sq_acc = unpack_states(states, tensors, 'Hz_sq_acc', cls=TensorList)
+        num_Hzs = self.global_state["num_Hzs"]
+        fs = settings[0]
 
-                Hvp, rgrad = var.hessian_vector_product(u, at_x0=True, rgrad=rgrad, hvp_method=hvp_method,
-                                     h=fd_h, normalize=True, retain_graph=j < n_samples-1)
+        # ---------------------------------- debias ---------------------------------- #
+        beta = fs["beta"]
+        beta_debias = fs["beta_debias"]
 
-                if D is None: D = Hvp
-                else: torch._foreach_add_(D, Hvp)
+        if beta_debias and beta is not None:
+            bias_correction = 1.0 - beta ** num_Hzs
+            Hz_sq_acc = Hz_sq_acc / bias_correction
 
-            assert D is not None
-            if n_samples > 1: torch._foreach_div_(D, n_samples)
+        else:
+            Hz_sq_acc = Hz_sq_acc / num_Hzs
 
-            D = TensorList(D)
+        # ---------------------------------- update ---------------------------------- #
+        damping = [s["damping"] for s in settings]
 
-        update = var.get_update()
-        if 'inner' in self.children:
-            update = apply_transform(self.children['inner'], tensors=update, params=params, grads=var.grad, var=var)
+        denom = (Hz_sq_acc / num_Hzs).sqrt_().add_(damping)
 
-        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
+        objective.updates = tensors.div_(denom)
+        return objective

torchzero/modules/adaptive/ggt.py

@@ -0,0 +1,186 @@
+from collections import deque
+from typing import Literal, Any
+import warnings
+
+import torch
+from ...core import Chainable, TensorTransform
+from ...linalg import torch_linalg, regularize_eigh
+from .lre_optimizers import LREOptimizerBase
+
+def ggt_update(history: deque[torch.Tensor] | torch.Tensor, damping, rdamping, truncate, eig_tol):
+    """returns U ``(ndim, rank)``, L ``(rank, )``"""
+    if isinstance(history, torch.Tensor):
+        M = history
+    else:
+        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:
+        L, Q = torch_linalg.eigh(MtM, retry_float64=True)
+
+        # damping is already added to MTM, rdamping is added afterwards
+        L, Q = regularize_eigh(L, Q, truncate=truncate, tol=eig_tol, damping=0, rdamping=0)
+
+        if L is None or Q is None: # this means there are no finite eigenvalues
+            return None, None
+
+        U = (M @ Q) * L.rsqrt()
+
+        # this damping is added after computing U, this is why I didn't use one in linalg.regularize_eig
+        # that's because we damp singular values this way
+        if rdamping != 0:
+            L.add_(rdamping * L[-1]) # L is sorted in ascending order
+
+        return L, U
+
+    except torch.linalg.LinAlgError:
+        return None, None
+
+
+class GGT(TensorTransform):
+    """
+    GGT method from https://arxiv.org/pdf/1806.02958
+
+    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.
+        beta (float, optional): beta for momentum maintained in whitened space. Defaults to 0.0.
+        update_freq (int, optional): frequency of updating the preconditioner (U and S). Defaults to 1.
+        eig_tol (float, optional): removes eigenvalues this much smaller than largest eigenvalue. Defaults to 1e-7.
+        truncate (int, optional): number of larges eigenvalues to keep. None to disable. Defaults to None.
+        damping (float, optional): damping value. Defaults to 1e-4.
+        rdamping (float, optional): value of damping relative to largest eigenvalue. Defaults to 0.
+        concat_params (bool, optional): if True, treats all parameters as a single vector. Defaults to True.
+        inner (Chainable | None, optional): preconditioner will be applied to output of this module. Defaults to None.
+
+    ## Examples:
+
+    Limited-memory Adagrad
+
+    ```python
+    optimizer = tz.Optimizer(
+        model.parameters(),
+        tz.m.GGT(),
+        tz.m.LR(0.1)
+    )
+    ```
+    Adam with L-Adagrad preconditioner (for debiasing second beta is 0.999 arbitrarily)
+
+    ```python
+    optimizer = tz.Optimizer(
+        model.parameters(),
+        tz.m.GGT(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)
+
+    ```python
+    optimizer = tz.Optimizer(
+        model.parameters(),
+        tz.m.GGT(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 = 100,
+        update_freq: int = 1,
+        eig_tol: float = 1e-7,
+        truncate: int | None = None,
+        damping: float = 1e-4,
+        rdamping: float = 0,
+        eigenbasis_optimizer: LREOptimizerBase | None = None,
+        concat_params: bool = True,
+
+        inner: Chainable | None = None,
+    ):
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner'], defaults['concat_params']
+
+        super().__init__(defaults, concat_params=concat_params, inner=inner)
+
+    @torch.no_grad
+    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
+        history_size = setting['history_size']
+        update_freq = setting['update_freq']
+
+        if 'history' not in state: state['history'] = deque(maxlen=history_size)
+        history = state['history']
+
+        t = tensor.clone().view(-1)
+        history.append(t)
+
+        step = state.get('step', 0)
+        state['step'] = step + 1
+
+        if step % update_freq == 0 :
+
+            # compute new factors
+            L = state.get("L", None)
+            U = state.get("U", None)
+
+            L_new, U_new = ggt_update(
+                history,
+                damping=setting["damping"],
+                rdamping=setting["rdamping"],
+                truncate=setting["truncate"],
+                eig_tol=setting["eig_tol"],
+            )
+
+            # reproject eigenbasis optimizer
+            eigenbasis_optimizer: LREOptimizerBase | None = setting["eigenbasis_optimizer"]
+            if eigenbasis_optimizer is not None:
+                if (L is not None) and (U is not None) and (L_new is not None) and (U_new is not None):
+                    eigenbasis_state = state["eigenbasis_state"]
+                    eigenbasis_optimizer.reproject(L_old=L, Q_old=U, L_new=L_new, Q_new=U_new, state=eigenbasis_state)
+
+
+            # store new factors
+            if L_new is not None: state["L"] = L_new
+            if U_new is not None: state["U"] = U_new
+
+
+    @torch.no_grad
+    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
+        g = tensor.view(-1)
+        U = state.get('U', None)
+
+        if U is None:
+            # fallback to element-wise preconditioning
+            history = torch.stack(tuple(state["history"]), 0)
+            g /= history.square().mean(0).sqrt().add(1e-8)
+            return g.view_as(tensor)
+
+        L = state['L']
+
+        # step with eigenbasis optimizer
+        eigenbasis_optimizer: LREOptimizerBase | None = setting["eigenbasis_optimizer"]
+        if eigenbasis_optimizer is not None:
+
+            if "eigenbasis_state" not in state: state["eigenbasis_state"] = {}
+            eigenbasis_state = state["eigenbasis_state"]
+
+            update = eigenbasis_optimizer.step(g, L=L, Q=U, state=eigenbasis_state)
+            return update.view_as(tensor)
+
+        # or just whiten
+        z = U.T @ g
+        update = (U * L.rsqrt()) @ z
+        return update.view_as(tensor)
+

torchzero/modules/adaptive/lion.py

@@ -1,21 +1,17 @@
+from typing import Any
 import torch
 
-from ...core import Module, Target, Transform
+from ...core import TensorTransform
 from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 
 
-def lion_(tensors: TensorList, exp_avg_: TensorList, beta1, beta2,):
-    """
-    Lion update rule.
-
-    Returns new tensors.
-    """
+def lion_(tensors: TensorList | Any, exp_avg_: TensorList | Any, beta1, beta2,):
     update = exp_avg_.lerp(tensors, 1-beta1).sign_()
     exp_avg_.lerp_(tensors, 1-beta2)
     return update
 
 
-class Lion(Transform):
+class Lion(TensorTransform):
     """Lion (EvoLved Sign Momentum) optimizer from https://arxiv.org/abs/2302.06675.
 
     Args:
@@ -25,11 +21,11 @@ class Lion(Transform):
 
     def __init__(self, beta1: float = 0.9, beta2: float = 0.99):
         defaults = dict(beta1=beta1, beta2=beta2)
-        super().__init__(defaults, uses_grad=False)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(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)
+        return lion_(TensorList(tensors), exp_avg, beta1, beta2)