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

torchzero/modules/adaptive/lre_optimizers.py

@@ -0,0 +1,299 @@
+"""subspace optimizers to be used in a low rank eigenbasis
+
+three opts support this - GGT and experimental AdaNystrom and Eigengrad
+
+I could define repoject on a module but because most opts use per-parameter state that is complicated"""
+
+import math
+from abc import ABC, abstractmethod
+from typing import Any, cast
+
+import torch
+
+from ...linalg import matrix_power_eigh, torch_linalg
+from .lion import lion_
+
+class LREOptimizerBase(ABC):
+    """Optimizer to run in a low rank eigenbasis.
+
+    notes:
+
+    1. it shouldn't store any states in self, everything should be in state.
+    This is because this may be called on multiple parameters in a sequence
+
+    2. apply is always called first, than reproject whenever eigenbasis gets updated
+
+    3. L is variance in the eigenbasis.
+    """
+    @abstractmethod
+    def step(self, g: torch.Tensor, L: torch.Tensor, Q: torch.Tensor, state: dict) -> torch.Tensor:
+        ...
+
+    @abstractmethod
+    def reproject(self, L_old: torch.Tensor, Q_old: torch.Tensor,
+                  L_new: torch.Tensor, Q_new: torch.Tensor, state: dict) -> None:
+        ...
+
+class Whiten(LREOptimizerBase):
+    """This simply applies whitening and is equivalent to not running an optimizer in the eigenbasis"""
+    def step(self, g, L, Q, state): return (Q * L.rsqrt()) @ (Q.T @ g)
+    def reproject(self, L_old, Q_old, L_new, Q_new, state): pass
+
+class EMA(LREOptimizerBase):
+    """Maintains exponential moving average of gradients in the low rank eigenbasis. Nesterov setting is experimental"""
+    def __init__(self, beta=0.9, nesterov:bool=False, cautious:bool=False, whiten:bool=True):
+        self.beta = beta
+        self.nesterov = nesterov
+        self.whiten = whiten
+        self.cautious = cautious
+
+    def step(self, g, L, Q, state):
+        g = Q.T @ g
+
+        if "exp_avg" not in state:
+            state["exp_avg"] = torch.zeros_like(g)
+
+        exp_avg = state["exp_avg"]
+        exp_avg.lerp_(g, 1-self.beta)
+
+        if self.nesterov:
+            dir = (g + exp_avg * self.beta) / (1 + self.beta)
+        else:
+            dir = exp_avg
+
+        if self.cautious:
+            mask = (g * dir) > 0
+            dir *= mask
+
+        if self.whiten: return (Q * L.rsqrt()) @ dir
+        return Q @ dir
+
+    def reproject(self, L_old, Q_old, L_new, Q_new, state):
+        if  "exp_avg" not in state: return
+        C = Q_new.T @ Q_old
+        state["exp_avg"] = C @ state["exp_avg"]
+
+
+def adam(g:torch.Tensor, state:dict, beta1, beta2, eps):
+
+    if "exp_avg" not in state:
+        state["exp_avg"] = torch.zeros_like(g)
+        state["exp_avg_sq"] = torch.zeros_like(g)
+        state["current_step"] = 1
+
+    exp_avg = state["exp_avg"]
+    exp_avg_sq = state["exp_avg_sq"]
+    current_step = state["current_step"]
+
+    exp_avg.lerp_(g, 1-beta1)
+    exp_avg_sq.mul_(beta2).addcmul_(g, g, value=1-beta2)
+    denom = exp_avg_sq.sqrt().add_(eps)
+
+    bias_correction1 = 1.0 - (beta1 ** current_step)
+    bias_correction2 = 1.0 - (beta2 ** current_step)
+    alpha = math.sqrt(bias_correction2) / bias_correction1
+    state["current_step"] = current_step + 1
+
+    return (exp_avg * alpha) / denom
+
+def _squared_reproject(C: torch.Tensor, sq: torch.Tensor, exact: bool):
+    if exact:
+        return (C @ sq.diag_embed() @ C.T).diagonal()
+
+    return C.square() @ sq
+
+class Adam(LREOptimizerBase):
+    """Runs Adam in low rank eigenbasis."""
+    def __init__(self, beta1=0.9, beta2=0.95, cautious:bool=False, eps=1e-8, exact_reproject:bool=True):
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.eps = eps
+        self.cautious = cautious
+        self.exact_reproject = exact_reproject
+
+    def step(self, g, L, Q, state):
+        g = Q.T @ g
+
+        dir = adam(g, state, self.beta1, self.beta2, self.eps)
+
+        if self.cautious:
+            mask = (g * dir) > 0
+            dir *= mask
+
+        return Q @ dir
+
+    def reproject(self, L_old, Q_old, L_new, Q_new, state):
+        if  "exp_avg" not in state: return
+        C = Q_new.T @ Q_old
+
+        state["exp_avg"] = C @ state["exp_avg"]
+        state["exp_avg_sq"] = _squared_reproject(C, state["exp_avg_sq"], self.exact_reproject)
+
+
+class FullMatrixAdam(LREOptimizerBase):
+    """Runs full-matrix Adam in low rank eigenbasis.
+    The preconditioner is updated whenever basis is updated"""
+    def __init__(self, beta1=0.9, beta2=0.95, eps=1e-8, matrix_power=-1/2, abs=True, cautious:bool=False):
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.eps = eps
+        self.matrix_power = matrix_power
+        self.abs = abs
+        self.cautious = cautious
+
+    def step(self, g, L, Q, state):
+        g = Q.T @ g
+
+        # initialize
+        if "exp_avg" not in state:
+            state["exp_avg"] = torch.zeros_like(g)
+            state["covariance"] = torch.eye(g.numel(), device=g.device, dtype=g.dtype)
+            state["preconditioner"] = torch.eye(g.numel(), device=g.device, dtype=g.dtype)
+            state["reprojected"] = True
+            state["current_step"] = 1
+
+        exp_avg = state["exp_avg"]
+        covariance = state["covariance"]
+        current_step = state["current_step"]
+
+        # update buffers
+        exp_avg.lerp_(g, 1-self.beta1)
+        covariance.lerp_(g.outer(g), weight=1-self.beta2)
+
+        # correct bias
+        bias_correction1 = 1.0 - (self.beta1 ** current_step)
+        exp_avg = exp_avg / bias_correction1
+
+        # after reprojecting update the preconditioner
+        if state["reprojected"]:
+            state["reprojected"] = False
+
+            bias_correction2 = 1.0 - (self.beta2 ** current_step)
+            covariance = covariance / bias_correction2
+
+            reg = torch.eye(covariance.size(0), device=covariance.device, dtype=covariance.dtype).mul_(self.eps)
+            covariance = covariance + reg
+
+            # compute matrix power
+            try:
+                state["preconditioner"] = matrix_power_eigh(covariance, self.matrix_power, abs=self.abs)
+
+            except torch.linalg.LinAlgError:
+
+                # fallback to diagonal
+                state["preconditioner"] = covariance.diagonal().rsqrt().diag_embed()
+
+        # compute the update
+        state["current_step"] = current_step + 1
+        preconditioner = state["preconditioner"]
+        dir = preconditioner @ exp_avg
+
+        if self.cautious:
+            mask = (g * dir) > 0
+            dir *= mask
+
+        return Q @ dir
+
+    def reproject(self, L_old, Q_old, L_new, Q_new, state):
+        if  "exp_avg" not in state: return
+
+        state["reprojected"] = True
+
+        C = Q_new.T @ Q_old
+        state["exp_avg"] = C @ state["exp_avg"]
+        state["covariance"] = C @ state["covariance"] @ C.T
+
+class Lion(LREOptimizerBase):
+    """Runs Lion in the low rank eigenbasis."""
+    def __init__(self, beta1=0.9, beta2=0.99, cautious:bool=False):
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.cautious = cautious
+
+    def step(self, g, L, Q, state):
+        g = Q.T @ g
+
+        if "exp_avg" not in state:
+            state["exp_avg"] = torch.zeros_like(g)
+
+        dir = cast(torch.Tensor, lion_(g, state["exp_avg"], beta1=self.beta1, beta2=self.beta2))
+
+        if self.cautious:
+            mask = (g * dir) > 0
+            dir *= mask
+
+        return Q @ dir
+
+    def reproject(self, L_old, Q_old, L_new, Q_new, state):
+        if "exp_avg" not in state: return
+        C = Q_new.T @ Q_old
+        state["exp_avg"] = C @ state["exp_avg"]
+
+
+class Grams(LREOptimizerBase):
+    """Runs Grams in low rank eigenbasis."""
+    def __init__(self, beta1=0.9, beta2=0.95, eps=1e-8, exact_reproject=True):
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.eps = eps
+        self.exact_reproject = exact_reproject
+
+    def step(self, g, L, Q, state):
+        g = Q.T @ g
+        dir = adam(g, state, self.beta1, self.beta2, self.eps)
+        return Q @ dir.copysign(g)
+
+    def reproject(self, L_old, Q_old, L_new, Q_new, state):
+        if  "exp_avg" not in state: return
+        C = Q_new.T @ Q_old
+
+        state["exp_avg"] = C @ state["exp_avg"]
+        state["exp_avg_sq"] = _squared_reproject(C, state["exp_avg_sq"], self.exact_reproject)
+
+
+class LaProp(LREOptimizerBase):
+    """Runs LaProp in low rank eigenbasis."""
+    def __init__(self, beta1=0.9, beta2=0.95, eps=1e-8, cautious:bool=False, exact_reproject=True):
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.eps = eps
+        self.cautious = cautious
+        self.exact_reproject = exact_reproject
+
+    def step(self, g, L, Q, state):
+        g = Q.T @ g
+
+        if "exp_avg" not in state:
+            state["exp_avg"] = torch.zeros_like(g)
+            state["exp_avg_sq"] = torch.zeros_like(g)
+            state["current_step"] = 1
+
+        exp_avg = state["exp_avg"]
+        exp_avg_sq = state["exp_avg_sq"]
+        current_step = state["current_step"]
+
+        # update second moments
+        exp_avg_sq.mul_(self.beta2).addcmul_(g, g, value=1-self.beta2)
+        bias_correction2 = 1.0 - (self.beta2 ** current_step)
+
+        # divide by bias corrected second moments
+        dir = g / (exp_avg_sq / bias_correction2).sqrt().add_(self.eps)
+
+        # update first moments and bias correct
+        exp_avg.lerp_(dir, 1-self.beta1)
+        bias_correction1 = 1.0 - (self.beta1 ** current_step)
+        dir = exp_avg / bias_correction1
+
+        if self.cautious:
+            mask = (g * dir) > 0
+            dir *= mask
+
+        state["current_step"] = current_step + 1
+        return Q @ dir
+
+    def reproject(self, L_old, Q_old, L_new, Q_new, state):
+        if  "exp_avg" not in state: return
+        C = Q_new.T @ Q_old
+
+        state["exp_avg"] = C @ state["exp_avg"]
+        state["exp_avg_sq"] = _squared_reproject(C, state["exp_avg_sq"], self.exact_reproject)

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,
@@ -35,7 +35,7 @@ class MARSCorrection(Transform):
 
     Mars-AdamW
     ```python
-    optimizer = tz.Modular(
+    optimizer = tz.Optimizer(
         model.parameters(),
         tz.m.MARSCorrection(beta=0.95),
         tz.m.Adam(beta1=0.95, beta2=0.99),
@@ -46,7 +46,7 @@ class MARSCorrection(Transform):
 
     Mars-Lion
     ```python
-    optimizer = tz.Modular(
+    optimizer = tz.Optimizer(
         model.parameters(),
         tz.m.MARSCorrection(beta=0.9),
         tz.m.Lion(beta1=0.9),
@@ -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']

torchzero/modules/adaptive/matrix_momentum.py

@@ -1,14 +1,13 @@
 from typing import Literal
-from collections.abc import Callable
+
 import torch
 
-from ...core import Module, apply_transform, Chainable
-from ...utils import NumberList, TensorList, as_tensorlist
-from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
-from ..functional import initial_step_size
+from ...core import Chainable, Transform, HVPMethod
+from ...utils import NumberList, TensorList, unpack_states, unpack_dicts
+from ..opt_utils import initial_step_size
 
 
-class MatrixMomentum(Module):
+class MatrixMomentum(Transform):
     """Second order momentum method.
 
     Matrix momentum is useful for convex objectives, also for some reason it has very really good generalization on elastic net logistic regression.
@@ -23,17 +22,17 @@ class MatrixMomentum(Module):
     Args:
         mu (float, optional): this has a similar role to (1 - beta) in normal momentum. Defaults to 0.1.
         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 hvp_method is set to finite difference. 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.
         hvp_tfm (Chainable | None, optional): optional module applied to hessian-vector products. Defaults to None.
 
     Reference:
@@ -44,51 +43,45 @@ class MatrixMomentum(Module):
         self,
         lr:float,
         mu=0.1,
-        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
+        hvp_method: HVPMethod = "autograd",
         h: float = 1e-3,
         adaptive:bool = False,
         adapt_freq: int | None = None,
-        hvp_tfm: Chainable | None = None,
+
+        inner: Chainable | None = None,
     ):
         defaults = dict(lr=lr, mu=mu, hvp_method=hvp_method, h=h, adaptive=adaptive, adapt_freq=adapt_freq)
-        super().__init__(defaults)
-
-        if hvp_tfm is not None:
-            self.set_child('hvp_tfm', hvp_tfm)
+        super().__init__(defaults, inner=inner)
 
     def reset_for_online(self):
         super().reset_for_online()
         self.clear_state_keys('p_prev')
 
     @torch.no_grad
-    def update(self, var):
-        assert var.closure is not None
-        p = TensorList(var.params)
-        p_prev = self.get_state(p, 'p_prev', init=var.params)
+    def update_states(self, objective, states, settings):
+        step = self.increment_counter("step", 0)
+        p = TensorList(objective.params)
+        p_prev = unpack_states(states, p, 'p_prev', init=p)
 
-        hvp_method = self.defaults['hvp_method']
-        h = self.defaults['h']
-        step = self.global_state.get("step", 0)
-        self.global_state["step"] = step + 1
+        fs = settings[0]
+        hvp_method = fs['hvp_method']
+        h = fs['h']
 
         if step > 0:
             s = p - p_prev
 
-            Hs, _ = var.hessian_vector_product(s, at_x0=True, rgrad=None, hvp_method=hvp_method, h=h, normalize=True, retain_graph=False)
+            Hs, _ = objective.hessian_vector_product(s, at_x0=True, rgrad=None, hvp_method=hvp_method, h=h, retain_graph=False)
             Hs = [t.detach() for t in Hs]
 
-            if 'hvp_tfm' in self.children:
-                Hs = TensorList(apply_transform(self.children['hvp_tfm'], Hs, params=p, grads=var.grad, var=var))
-
             self.store(p, ("Hs", "s"), (Hs, s))
 
             # -------------------------------- adaptive mu ------------------------------- #
-            if self.defaults["adaptive"]:
-                g = TensorList(var.get_grad())
+            if fs["adaptive"]:
+                g = TensorList(objective.get_grads())
 
-                if self.defaults["adapt_freq"] is None:
+                if fs["adapt_freq"] is None:
                     # ---------------------------- deterministic case ---------------------------- #
-                    g_prev = self.get_state(var.params, "g_prev", cls=TensorList)
+                    g_prev = unpack_states(states, p, "g_prev", cls=TensorList)
                     y = g - g_prev
                     g_prev.copy_(g)
                     denom = y.global_vector_norm()
@@ -101,14 +94,14 @@ class MatrixMomentum(Module):
 
                     # we start on 1nd step, and want to adapt when we start, so use (step - 1)
                     if (step - 1) % adapt_freq == 0:
-                        assert var.closure is not None
-                        params = TensorList(var.params)
+                        assert objective.closure is not None
+                        params = TensorList(objective.params)
                         p_cur = params.clone()
 
                         # move to previous params and evaluate p_prev with current mini-batch
-                        params.copy_(self.get_state(var.params, 'p_prev'))
+                        params.copy_(unpack_states(states, p, 'p_prev'))
                         with torch.enable_grad():
-                            var.closure()
+                            objective.closure()
                         g_prev = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
                         y = g - g_prev
 
@@ -119,12 +112,12 @@ class MatrixMomentum(Module):
                         denom = denom.clip(min=torch.finfo(denom.dtype).tiny * 2)
                         self.global_state["mu_mul"] = s.global_vector_norm() / denom
 
-        torch._foreach_copy_(p_prev, var.params)
+        torch._foreach_copy_(p_prev, objective.params)
 
     @torch.no_grad
-    def apply(self, var):
-        update = TensorList(var.get_update())
-        lr,mu = self.get_settings(var.params, "lr", 'mu', cls=NumberList)
+    def apply_states(self, objective, states, settings):
+        update = TensorList(objective.get_updates())
+        lr, mu = unpack_dicts(settings, "lr", 'mu', cls=NumberList)
 
         if "mu_mul" in self.global_state:
             mu = mu * self.global_state["mu_mul"]
@@ -133,14 +126,17 @@ class MatrixMomentum(Module):
         # p_prev is not available so make a small step
         step = self.global_state["step"]
         if step == 1:
-            if self.defaults["adaptive"]: self.get_state(var.params, "g_prev", init=var.get_grad())
+            if self.defaults["adaptive"]:
+                # initialize
+                unpack_states(states, objective.params, "g_prev", init=objective.get_grads())
+
             update.mul_(lr) # separate so that initial_step_size can clip correctly
             update.mul_(initial_step_size(update, 1e-7))
-            return var
+            return objective
 
         # -------------------------- matrix momentum update -------------------------- #
-        s, Hs = self.get_state(var.params, 's', 'Hs', cls=TensorList)
+        s, Hs = unpack_states(states, objective.params, 's', 'Hs', cls=TensorList)
 
         update.mul_(lr).sub_(s).add_(Hs*mu)
-        var.update = update
-        return var
+        objective.updates = update
+        return objective