torchzero 0.1.8__py3-none-any.whl → 0.3.1__py3-none-any.whl

torchzero/modules/experimental/adasoap.py

@@ -0,0 +1,282 @@
+from operator import itemgetter
+
+import torch
+
+from ...core import Chainable, Transform, apply
+from ...modules.optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
+
+@torch.no_grad
+def update_soap_covariances_(
+    grad: torch.Tensor,
+    GGs_: list[torch.Tensor | None],
+    GG_sqs: list[torch.Tensor | None],
+    beta: float | None,
+    precond_beta: float | None,
+):
+    for i, (GG, GG_sq) in enumerate(zip(GGs_, GG_sqs)):
+        if GG is None: continue
+        assert GG_sq is not None
+
+        if precond_beta is None: GG_sq.addcmul_(GG, GG)
+        else: GG_sq.mul_(precond_beta).addcmul_(GG, GG, value=1-precond_beta)
+
+        axes = list(range(i)) + list(range(i + 1, grad.ndim)) # this works fine with 1d params
+        if beta is None: GG.add_(torch.tensordot(grad, grad, (axes, axes))) # pyright:ignore[reportArgumentType]
+        else: GG.lerp_(torch.tensordot(grad, grad, (axes, axes)), 1-beta) # pyright:ignore[reportArgumentType]
+
+@torch.no_grad
+def project(tensors: torch.Tensor, Q: list[torch.Tensor | None]):
+    """
+    Projects the gradient to the eigenbases of the preconditioner.
+    """
+    for mat in Q:
+        if mat is None: continue
+        if len(mat) > 0:
+            tensors = torch.tensordot(tensors, mat, dims=[[0], [0]]) # pyright:ignore[reportArgumentType]
+        else:
+            # I don't understand this part but it is in https://github.com/nikhilvyas/SOAP/blob/main/soap.py
+            permute_order = list(range(1, len(tensors.shape))) + [0]
+            tensors = tensors.permute(permute_order)
+
+    return tensors
+
+@torch.no_grad
+def project_back(tensors: torch.Tensor, Q: list[torch.Tensor| None]):
+    """
+    Projects the gradient back to the original space.
+    """
+    for mat in Q:
+        if mat is None: continue
+        if len(mat) > 0:
+            tensors = torch.tensordot(tensors, mat,dims=[[0], [1]]) # pyright:ignore[reportArgumentType]
+        else:
+            permute_order = list(range(1, len(tensors.shape))) + [0]
+            tensors = tensors.permute(permute_order)
+
+    return tensors
+
+# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py
+@torch.no_grad
+def get_orthogonal_matrix(mat: list[torch.Tensor | None]):
+    """
+    Computes the eigenbases of the preconditioner using torch.linalg.eigh decomposition.
+    """
+    matrix = []
+    float_data = False
+    original_type = original_device = None
+    for m in mat:
+        if m is None: continue
+        if len(m) == 0:
+            matrix.append([])
+            continue
+        if m.dtype != torch.float:
+            original_type = m.dtype
+            original_device = m.device
+            matrix.append(m.float())
+        else:
+            float_data = True
+            matrix.append(m)
+
+    final = []
+    for m in matrix:
+        if len(m) == 0:
+            final.append([])
+            continue
+        try:
+            _, Q = torch.linalg.eigh(m+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
+        except Exception:
+            _, Q = torch.linalg.eigh(m.to(torch.float64)+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
+            Q = Q.to(m.dtype)
+        Q = torch.flip(Q, [1])
+
+        if not float_data:
+            Q = Q.to(original_device).type(original_type)
+        final.append(Q)
+    return final
+
+# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py#L240
+@torch.no_grad
+def get_orthogonal_matrix_QR(exp_avg_sq: torch.Tensor, GG: list[torch.Tensor | None], Q_list: list[torch.Tensor | None]):
+    """
+    Computes the eigenbases of the preconditioner using one round of power iteration
+    followed by torch.linalg.qr decomposition.
+    """
+    matrix = []
+    orth_matrix = []
+    float_data = False
+    original_type = original_device = None
+    for m,o in zip(GG, Q_list):
+        if m is None: continue
+        assert o is not None
+
+        if len(m) == 0:
+            matrix.append([])
+            orth_matrix.append([])
+            continue
+        if m.data.dtype != torch.float:
+            original_type = m.data.dtype
+            original_device = m.data.device
+            matrix.append(m.data.float())
+            orth_matrix.append(o.data.float())
+        else:
+            float_data = True
+            matrix.append(m.data.float())
+            orth_matrix.append(o.data.float())
+
+    final = []
+    for ind, (m,o) in enumerate(zip(matrix, orth_matrix)):
+        if len(m)==0:
+            final.append([])
+            continue
+        est_eig = torch.diag(o.T @ m @ o)
+        sort_idx = torch.argsort(est_eig, descending=True)
+        exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
+        o = o[:,sort_idx]
+        power_iter = m @ o
+        Q, _ = torch.linalg.qr(power_iter) # pylint:disable=not-callable
+
+        if not float_data:
+            Q = Q.to(original_device).type(original_type)
+        final.append(Q)
+
+    return final, exp_avg_sq
+
+class AdaSOAP(Transform):
+    """SOAP with diagonally preconditioned GG^Ts
+
+    precond_beta - beta for GG^T squares
+    """
+    def __init__(
+        self,
+        beta1: float = 0.95,
+        beta2: float = 0.95,
+        shampoo_beta: float | None = 0.95,
+        precond_beta: float | None = 0.95,
+        precond_freq: int = 10,
+        merge_small: bool = True,
+        max_dim: int = 2_000,
+        precondition_1d: bool = True,
+        eps: float = 1e-8,
+        decay: float | None = None,
+        alpha: float = 1,
+        unprojected_exp_avg: bool = True,
+        bias_correction: bool = True,
+    ):
+        defaults = dict(
+            beta1=beta1,
+            beta2=beta2,
+            shampoo_beta=shampoo_beta,
+            precond_beta=precond_beta,
+            precond_freq=precond_freq,
+            merge_small=merge_small,
+            max_dim=max_dim,
+            precondition_1d=precondition_1d,
+            eps=eps,
+            decay=decay,
+            unprojected_exp_avg=unprojected_exp_avg,
+            bias_correction=bias_correction,
+            alpha=alpha,
+        )
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        updates = []
+        # update preconditioners
+        for i,(p,t) in enumerate(zip(params, tensors)):
+            state = self.state[p]
+            settings = self.settings[p]
+            beta1, beta2, shampoo_beta, merge_small, max_dim, precondition_1d, eps, unprojected_exp_avg,alpha = itemgetter(
+                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'unprojected_exp_avg','alpha')(settings)
+            precond_beta = settings['precond_beta']
+
+            if merge_small:
+                t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
+
+            # initialize state on 1st step
+            if 'GG' not in state:
+                state["exp_avg"] = torch.zeros_like(t)
+                state["exp_avg_sq"] = torch.zeros_like(t)
+
+                if not precondition_1d and t.ndim <= 1:
+                    state['GG'] = []
+                    state['GG_sq'] = []
+
+                else:
+                    state['GG'] = [torch.zeros(s, s, dtype=t.dtype, device=t.device) if 1<s<max_dim else None for s in t.shape]
+                    state['GG_sq'] = [torch.zeros(s, s, dtype=t.dtype, device=t.device) if 1<s<max_dim else None for s in t.shape]
+
+                # either scalar parameter, 1d with precondition_1d=False, or all dims are too big.
+                if len([i is not None for i in state['GG']]) == 0:
+                    state['GG'] = None
+                    state['GG_sq'] = None
+
+                if state['GG'] is not None:
+                    assert state['GG_sq'] is not None
+                    update_soap_covariances_(t, GGs_=state['GG'], GG_sqs=state['GG_sq'], beta=shampoo_beta, precond_beta=precond_beta)
+                    GG_precond = [GG / (GG_sq+1e-8) if GG is not None and GG_sq is not None else None for GG, GG_sq in zip(state['GG'], state['GG_sq'])]
+                    state['Q'] = get_orthogonal_matrix(GG_precond)
+
+                state['step'] = 0
+                updates.append(tensors[i].sign())
+                continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
+                # I use sign instead as to not mess up with next modules. 1st Adam step is always sign anyway.
+
+            # Projecting gradients to the eigenbases of Shampoo's preconditioner
+            # i.e. projecting to the eigenbases of matrices in state['GG']
+            t_projected = None
+            if state['GG'] is not None:
+                t_projected = project(t, state['Q'])
+
+            # exponential moving averages
+            # this part could be foreached but I will do that at some point its not a big difference compared to preconditioning
+            exp_avg: torch.Tensor = state["exp_avg"]
+            exp_avg_sq: torch.Tensor = state["exp_avg_sq"]
+
+            if unprojected_exp_avg or t_projected is None:
+                exp_avg.lerp_(t, 1-beta1)
+            else:
+                exp_avg.lerp_(t_projected, 1-beta1)
+
+            if t_projected is None:
+                exp_avg_sq.mul_(beta2).addcmul_(t, t, value=1-beta2)
+            else:
+                exp_avg_sq.mul_(beta2).addcmul_(t_projected, t_projected, value=1-beta2)
+
+            # project exponential moving averages if they are accumulated unprojected
+            exp_avg_projected = exp_avg
+            if unprojected_exp_avg and t_projected is not None:
+                exp_avg_projected = project(exp_avg, state['Q'])
+
+            exp_avg_sq_projected = exp_avg_sq
+
+            denom = exp_avg_sq_projected.sqrt().add_(eps)
+            # print(f'{t_projected = }, {exp_avg = }, {exp_avg_projected = }, {exp_avg_sq = }, {exp_avg_sq_projected = }, {denom = }')
+
+            # Projecting back the preconditioned (by Adam) exponential moving average of gradients
+            # to the original space
+            update = exp_avg_projected / denom
+            if t_projected is not None:
+                update = project_back(update, state["Q"])
+
+            if settings['bias_correction']:
+                bias_correction1 = 1.0 - beta1 ** (state["step"]+1)
+                bias_correction2 = 1.0 - beta2 ** (state["step"]+1)
+                update *= ((bias_correction2 ** .5) / bias_correction1) * alpha
+            elif alpha is not None:
+                update *= alpha
+
+            if merge_small:
+                update = _unmerge_small_dims(update, state['flat_sizes'], state['sort_idxs'])
+
+            updates.append(update)
+            state["step"] += 1
+
+            # Update is done after the gradient step to avoid using current gradients in the projection.
+            if state['GG'] is not None:
+                update_soap_covariances_(t, GGs_=state['GG'], GG_sqs=state['GG_sq'], beta=shampoo_beta, precond_beta=precond_beta)
+                GG_precond = [GG / (GG_sq+1e-8) if GG is not None and GG_sq is not None else None for GG, GG_sq in zip(state['GG'], state['GG_sq'])]
+                if state['step'] % settings['precond_freq'] == 0:
+                    state['Q'], state['exp_avg_sq'] = get_orthogonal_matrix_QR(exp_avg_sq, GG_precond, state['Q'])
+
+        return updates

torchzero/modules/experimental/algebraic_newton.py

@@ -0,0 +1,145 @@
+import warnings
+from functools import partial
+from typing import Literal
+from collections.abc import Callable
+import torch
+import torchalgebras as ta
+
+from ...core import Chainable, apply, Module
+from ...utils import vec_to_tensors, TensorList
+from ...utils.derivatives import (
+    hessian_list_to_mat,
+    hessian_mat,
+    jacobian_and_hessian_wrt,
+)
+
+class MaxItersReached(Exception): pass
+def tropical_lstsq(
+    H: torch.Tensor,
+    g: torch.Tensor,
+    solver,
+    maxiter,
+    tol,
+    algebra,
+    verbose,
+):
+    """it can run on any algebra with add despite it saying tropical"""
+    algebra = ta.get_algebra(algebra)
+
+    x = torch.zeros_like(g, requires_grad=True)
+    best_x = x.detach().clone()
+    best_loss = float('inf')
+    opt = solver([x])
+
+    niter = 0
+    def closure(backward=True):
+        nonlocal niter, best_x, best_loss
+        if niter == maxiter: raise MaxItersReached
+        niter += 1
+
+        g_hat = algebra.mm(H, x)
+        loss = torch.nn.functional.mse_loss(g_hat, g)
+        if loss < best_loss:
+            best_x = x.detach().clone()
+            best_loss = loss.detach()
+
+        if backward:
+            opt.zero_grad()
+            loss.backward()
+        return loss
+
+    loss = None
+    prev_loss = float('inf')
+    for i in range(maxiter):
+        try:
+            loss = opt.step(closure)
+            if loss == 0: break
+            if tol is not None and prev_loss - loss < tol: break
+            prev_loss = loss
+        except MaxItersReached:
+            break
+
+    if verbose: print(f'{best_loss = } after {niter} iters')
+    return best_x.detach()
+
+def tikhonov(H: torch.Tensor, reg: float, algebra: ta.Algebra = ta.TropicalSemiring()):
+    if reg!=0:
+        I = ta.AlgebraicTensor(torch.eye(H.size(-1), dtype=H.dtype, device=H.device), algebra)
+        I = I * reg
+        H = algebra.add(H, I.data)
+    return H
+
+
+class AlgebraicNewton(Module):
+    """newton in other algebras, not practical because solving linear system is very hard."""
+    def __init__(
+        self,
+        reg: float | None = None,
+        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
+        vectorize: bool = True,
+        solver=lambda p: torch.optim.LBFGS(p, line_search_fn='strong_wolfe'),
+        maxiter=1000,
+        tol: float | None = 1e-10,
+        algebra: ta.Algebra | str = 'tropical max',
+        verbose: bool = False,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(reg=reg, hessian_method=hessian_method, vectorize=vectorize)
+        super().__init__(defaults)
+
+        self.algebra = ta.get_algebra(algebra)
+        self.lstsq_args:dict = dict(solver=solver, maxiter=maxiter, tol=tol, algebra=algebra, verbose=verbose)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, vars):
+        params = TensorList(vars.params)
+        closure = vars.closure
+        if closure is None: raise RuntimeError('NewtonCG requires closure')
+
+        settings = self.settings[params[0]]
+        reg = settings['reg']
+        hessian_method = settings['hessian_method']
+        vectorize = settings['vectorize']
+
+        # ------------------------ calculate grad and hessian ------------------------ #
+        if hessian_method == 'autograd':
+            with torch.enable_grad():
+                loss = vars.loss = vars.loss_approx = closure(False)
+                g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
+                g_list = [t[0] for t in g_list] # remove leading dim from loss
+                vars.grad = g_list
+                H = hessian_list_to_mat(H_list)
+
+        elif hessian_method in ('func', 'autograd.functional'):
+            strat = 'forward-mode' if vectorize else 'reverse-mode'
+            with torch.enable_grad():
+                g_list = vars.get_grad(retain_graph=True)
+                H: torch.Tensor = hessian_mat(partial(closure, backward=False), params,
+                                method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
+
+        else:
+            raise ValueError(hessian_method)
+
+        # -------------------------------- inner step -------------------------------- #
+        if 'inner' in self.children:
+            g_list = apply(self.children['inner'], list(g_list), params=params, grads=list(g_list), vars=vars)
+        g = torch.cat([t.view(-1) for t in g_list])
+
+        # ------------------------------- regulazition ------------------------------- #
+        if reg is not None: H = tikhonov(H, reg)
+
+        # ----------------------------------- solve ---------------------------------- #
+        tropical_update = tropical_lstsq(H, g, **self.lstsq_args)
+        # what now? w - u is not defined, it is defined for max version if u < w
+        # w = params.to_vec()
+        # w_hat = self.algebra.sub(w, tropical_update)
+        # update = w_hat - w
+        # no
+        # it makes sense to solve tropical system and sub normally
+        # the only thing is that tropical system can have no solutions
+
+        vars.update = vec_to_tensors(tropical_update, params)
+        return vars

torchzero/modules/experimental/curveball.py

@@ -0,0 +1,89 @@
+from typing import Literal
+from collections.abc import Callable
+import torch
+
+from ...core import Module, Target, Transform, Chainable, apply
+from ...utils import NumberList, TensorList, as_tensorlist
+from ...utils.derivatives import hvp, hvp_fd_forward, hvp_fd_central
+
+def curveball(
+    tensors: TensorList,
+    z_: TensorList,
+    Hz: TensorList,
+    momentum: float | NumberList,
+    precond_lr: float | NumberList,
+):
+    """returns z_, clone it!!!"""
+    delta = Hz + tensors
+    z_.mul_(momentum).sub_(delta.mul_(precond_lr)) # z ← ρz − βΔ
+    return z_
+
+
+class CurveBall(Module):
+    """CurveBall method from https://arxiv.org/pdf/1805.08095#page=4.09.
+
+    For now this implementation does not include automatic ρ, α and β hyper-parameters in closed form, therefore it is expected to underperform compared to official implementation (https://github.com/jotaf98/pytorch-curveball/tree/master) so I moved this to experimental.
+
+    Args:
+        precond_lr (float, optional): learning rate for updating preconditioned gradients. Defaults to 1e-3.
+        momentum (float, optional): decay rate for preconditioned gradients. Defaults to 0.9.
+        hvp_method (str, optional): how to calculate hessian vector products. Defaults to "autograd".
+        h (float, optional): finite difference step size for when hvp_method is set to finite difference. Defaults to 1e-3.
+        reg (float, optional): hessian regularization. Defaults to 1.
+        inner (Chainable | None, optional): Inner modules. Defaults to None.
+    """
+    def __init__(
+        self,
+        precond_lr: float=1e-3,
+        momentum: float=0.9,
+        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
+        h: float = 1e-3,
+        reg: float = 1,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(precond_lr=precond_lr, momentum=momentum, hvp_method=hvp_method, h=h, reg=reg)
+        super().__init__(defaults)
+
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, vars):
+
+        params = vars.params
+        settings = self.settings[params[0]]
+        hvp_method = settings['hvp_method']
+        h = settings['h']
+
+        precond_lr, momentum, reg = self.get_settings('momentum', 'decay_rate', 'reg', params=params, cls=NumberList)
+
+
+        closure = vars.closure
+        assert closure is not None
+
+        z, Hz = self.get_state('z', 'Hz', params=params, cls=TensorList)
+
+        if hvp_method == 'autograd':
+            grad = vars.get_grad(create_graph=True)
+            Hvp = hvp(params, grad, z)
+
+        elif hvp_method == 'forward':
+            loss, Hvp = hvp_fd_forward(closure, params, z, h=h, g_0=vars.get_grad(), normalize=True)
+
+        elif hvp_method == 'central':
+            loss, Hvp = hvp_fd_central(closure, params, z, h=h, normalize=True)
+
+        else:
+            raise ValueError(hvp_method)
+
+
+        Hz.set_(Hvp + z*reg)
+
+
+        update = vars.get_update()
+        if 'inner' in self.children:
+            update = apply(self.children['inner'], update, params, grads=vars.grad, vars=vars)
+
+        z = curveball(TensorList(update), z, Hz, momentum=momentum, precond_lr=precond_lr)
+        vars.update = z.neg()
+
+        return vars