torchzero 0.3.14__py3-none-any.whl → 0.4.0__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.Modular(
+        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.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,
+        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 = 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)
+        # ---------------------------------- 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/lion.py

@@ -1,21 +1,16 @@
 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.
-    """
     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 +20,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)
 

torchzero/modules/adaptive/lmadagrad.py

@@ -3,9 +3,11 @@ from typing import Literal, Any
 import warnings
 
 import torch
-from ...core import Chainable, TensorwiseTransform
+from ...core import Chainable, TensorTransform
+from ...linalg import torch_linalg
 
-def lm_adagrad_update(history: deque[torch.Tensor] | torch.Tensor, damping, rdamping):
+def lm_adagrad_update(history: deque[torch.Tensor] | torch.Tensor, damping, rdamping, truncate, tol):
+    """returns U ``(ndim, rank)``, L ``(rank, )``"""
     if isinstance(history, torch.Tensor):
         M = history
     else:
@@ -16,35 +18,49 @@ def lm_adagrad_update(history: deque[torch.Tensor] | torch.Tensor, damping, rdam
         MTM.add_(torch.eye(MTM.size(0), device=MTM.device, dtype=MTM.dtype).mul_(damping))
 
     try:
-        L, Q = torch.linalg.eigh(MTM) # pylint:disable=not-callable
+        L, Q = torch_linalg.eigh(MTM, retry_float64=True)
 
-        tol = torch.finfo(M.dtype).eps * L.amax() # remove small eigenvalues
-        indices = L > tol
-        L = L[indices]
-        Q = Q[:, indices]
+        # truncate to top n largest eigenvalues
+        if truncate is not None and truncate > 0:
+            # L is ordered in ascending order
+            L = L[-truncate:]
+            Q = Q[:, -truncate:]
+
+        # remove small eigenvalues relative to largest
+        L_max = L.amax()
+        indices = L > tol * L_max
+        if indices.any():
+            L = L[indices]
+            Q = Q[:, indices]
 
         U = (M @ Q) * L.rsqrt()
 
         if rdamping != 0:
-            rdamping *= torch.linalg.vector_norm(L) # pylint:disable=not-callable
-            L.add_(rdamping)
+            L.add_(rdamping * L_max)
 
         return U, L
 
     except torch.linalg.LinAlgError:
         return None, None
 
-def lm_adagrad_apply(g: torch.Tensor, U: torch.Tensor, L: torch.Tensor):
-    Z = U.T @ g
-    return (U * L.rsqrt()) @ Z
+def lm_adagrad_apply(g: torch.Tensor, U: torch.Tensor, L: torch.Tensor, exp_avg_proj: torch.Tensor | None, beta:float):
+    z = U.T @ g
+
+    if beta != 0:
+        if exp_avg_proj is None: exp_avg_proj = torch.zeros_like(z)
+        exp_avg_proj.lerp_(z, weight=1-beta)
+        z = exp_avg_proj
+
+    return (U * L.rsqrt()) @ z, exp_avg_proj
 
 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
+    if value is None: return
+    if (key not in state_) or (beta is None): state_[key] = value
     else:
         if state_[key] is None or state_[key].shape != value.shape: state_[key] = value
         else: state_[key].lerp_(value, 1-beta)
 
-class LMAdagrad(TensorwiseTransform):
+class LMAdagrad(TensorTransform):
     """
     Limited-memory full matrix Adagrad.
 
@@ -55,17 +71,18 @@ class LMAdagrad(TensorwiseTransform):
 
     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.
         damping (float, optional): damping value. Defaults to 1e-4.
         rdamping (float, optional): value of damping relative to singular values norm. Defaults to 0.
+        rdamping (float, optional): value of damping relative to singular values norm. Defaults to 0.
+        truncate (int, optional): number of larges eigenvalues to keep. None to disable. Defaults to None.
+        tol (float, optional): removes eigenvalues this much smaller than largest eigenvalue. Defaults to 1e-7.
         order (int, optional):
             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.
         U_beta (float | None, optional): momentum for U (too unstable, don't use). Defaults to None.
         L_beta (float | None, optional): momentum for L (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 True.
+        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:
@@ -108,28 +125,35 @@ class LMAdagrad(TensorwiseTransform):
     def __init__(
         self,
         history_size: int = 100,
+        beta: float = 0.0,
         update_freq: int = 1,
         damping: float = 1e-4,
         rdamping: float = 0,
+        truncate: int | None = None,
+        tol: float = 1e-7,
         order: int = 1,
-        true_damping: bool = True,
         U_beta: float | None = None,
         L_beta: float | None = None,
-        interval: int = 1,
         concat_params: bool = True,
+
         inner: Chainable | None = None,
+        U_tfm: Chainable | None = None,
+        L_tfm: 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, L_beta=L_beta)
-        super().__init__(defaults, uses_grad=False, concat_params=concat_params, inner=inner, update_freq=interval)
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner'], defaults['concat_params'], defaults["U_tfm"], defaults["L_tfm"]
+
+        super().__init__(defaults, concat_params=concat_params, inner=inner)
+
+        self.set_child("U", U_tfm)
+        self.set_child("L", L_tfm)
+
 
     @torch.no_grad
-    def update_tensor(self, tensor, param, grad, loss, state, setting):
+    def single_tensor_update(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']
 
@@ -165,22 +189,53 @@ class LMAdagrad(TensorwiseTransform):
 
         step = state.get('step', 0)
         if step % update_freq == 0 and len(history) != 0:
-            U, L = lm_adagrad_update(history, damping=damping, rdamping=rdamping)
+
+            # if maintaining momentum, unproject exp_avg before updating factors and reproject
+            exp_avg_proj = state.get("exp_avg_proj", None)
+            exp_avg = None
+            if exp_avg_proj is not None and "U" in state:
+                exp_avg = state["U"] @ exp_avg_proj
+
+            # update factors
+            U, L = lm_adagrad_update(
+                history,
+                damping=setting["damping"],
+                rdamping=setting["rdamping"],
+                truncate=setting["truncate"],
+                tol=setting["tol"],
+            )
             maybe_lerp_(state, U_beta, 'U', U)
             maybe_lerp_(state, L_beta, 'L', L)
 
+            # re-project exp_avg with new factors
+            if U is not None and exp_avg_proj is not None:
+                assert exp_avg is not None
+                state["exp_avg_proj"] = U.T @ exp_avg
+
+
         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, setting):
+    def single_tensor_apply(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]
+            return tensor.clip_(-0.1, 0.1)
 
+        # -------------------------------- transforms -------------------------------- #
         L = state['L']
-        update = lm_adagrad_apply(tensor.view(-1), U, L).view_as(tensor)
-
-        return update
+        if "L" in self.children:
+            if not self._concat_params: raise RuntimeError("L/U transforms can only be used with concat_params=True")
+            L = self.inner_step_tensors("L", [L], clone=True)[0]
+
+        if "U" in self.children:
+            if not self._concat_params: raise RuntimeError("L/U transforms can only be used with concat_params=True")
+            U = self.inner_step_tensors("U", [U], clone=True)[0]
+
+        # ------------------------------- precondition ------------------------------- #
+        g = tensor.view(-1)
+        exp_avg_proj = state.get("exp_avg_proj", None)
+        update, state["exp_avg_proj"] = lm_adagrad_apply(g, U, L, exp_avg_proj, beta=setting["beta"])
+        return update.view_as(tensor)
 

torchzero/modules/adaptive/mars.py

@@ -1,6 +1,6 @@
 import torch
 
-from ...core import Transform
+from ...core import TensorTransform
 from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 
 
@@ -20,7 +20,7 @@ def mars_correction_(
 
     return c
 
-class MARSCorrection(Transform):
+class MARSCorrection(TensorTransform):
     """MARS variance reduction correction.
 
     Place any other momentum-based optimizer after this,
@@ -61,11 +61,11 @@ class MARSCorrection(Transform):
         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)
+        defaults = dict(beta=beta, scaling=scaling, max_norm=max_norm)
+        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):
         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']