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

torchzero/modules/experimental/adanystrom.py

@@ -1,258 +0,0 @@
-# pylint: disable = non-ascii-name
-import torch
-
-from ...core import Chainable, TensorTransform
-from ...linalg import (
-    OrthogonalizeMethod,
-    orthogonalize,
-    regularize_eigh,
-    torch_linalg,
-)
-from ...linalg.linear_operator import Eigendecomposition
-from ..adaptive.lre_optimizers import LREOptimizerBase
-from .eigengrad import _eigengrad_update_state_, eigengrad_apply
-
-
-def weighted_eigen_plus_rank1_mm(
-    # A1 = Q1 @ diag(L1) @ Q1.T
-    L1: torch.Tensor,
-    Q1: torch.Tensor,
-
-    # K2 = v2 @ v2.T
-    v2: torch.Tensor,
-
-    # second matrix
-    B: torch.Tensor,
-
-    # weights
-    w1: float,
-    w2: float,
-
-) -> torch.Tensor:
-    """
-    Computes ``(w1 * A1 + w2 * A2) @ B``, where ``A1`` is an eigendecomposition, ``A2`` is symmetric rank 1.
-
-    Returns ``(n, k)``
-
-    Args:
-        L1 (torch.Tensor): eigenvalues of A1, shape ``(rank,)``.
-        Q1 (torch.Tensor): eigenvectors of A1, shape ``(n, rank)``.
-        v2 (torch.Tensor): vector such that ``v v^T = A2``, shape ``(n,)``.
-        B (torch.Tensor): shape ``(n, k)``.
-        w1 (float): weight for A1.
-        w2 (float): weight for A2.
-
-    """
-    # sketch A1
-    QTB = Q1.T @ B # (rank, k)
-    LQTB = L1.unsqueeze(1) * QTB  # (rank, k)
-    sketch1 = Q1 @ LQTB  # (n, k)
-
-    # skecth A2
-    vB = v2 @ B
-    sketch2 = v2.outer(vB)
-
-    return w1 * sketch1 + w2 * sketch2
-
-
-def adanystrom_update(
-    L1: torch.Tensor,
-    Q1: torch.Tensor,
-    v2: torch.Tensor,
-    w1: float,
-    w2: float,
-    oversampling_p: int,
-    rank: int,
-    eig_tol: float,
-    damping: float,
-    rdamping: float,
-    orthogonalize_method: OrthogonalizeMethod,
-
-) -> tuple[torch.Tensor | None, torch.Tensor | None]:
-    """computes the Nyström approximation of ``(w1 * A1 + w2 * A2)``,
-    where ``A1`` is an eigendecomposition, ``A2`` is symmetric rank 1.
-
-    returns L of shape ``(k, )`` and Q of shape ``(n, k)``.
-
-    Args:
-        L1 (torch.Tensor): eigenvalues of A1, shape ``(rank,)``.
-        Q1 (torch.Tensor): eigenvectors of A1, shape ``(n, rank)``.
-        v2 (torch.Tensor): vector such that ``v v^T = A2``, shape ``(n,)`` or ``(n, 1)``.
-        w1 (float): weight for A1.
-        w2 (float): weight for A2.
-    """
-    n = Q1.shape[0]
-    device = Q1.device
-    dtype = Q1.dtype
-    l = rank + oversampling_p
-
-    # gaussian test matrix
-    Omega = torch.randn(n, l, device=device, dtype=dtype)
-
-    # sketch
-    AOmega = weighted_eigen_plus_rank1_mm(L1, Q1, v2, Omega, w1, w2)
-    Q = orthogonalize(AOmega, orthogonalize_method)
-
-    AQ = weighted_eigen_plus_rank1_mm(L1, Q1, v2, Q, w1, w2)
-    QTAQ = Q.T @ AQ
-
-    W = (QTAQ + QTAQ.T) / 2.0
-
-    # compute new L and Q
-    try:
-        L_prime, S = torch_linalg.eigh(W, retry_float64=True)
-    except torch.linalg.LinAlgError:
-        return L1, Q1
-
-    L_prime, S = regularize_eigh(L=L_prime, Q=S, truncate=rank, tol=eig_tol, damping=damping, rdamping=rdamping)
-
-    if L_prime is None or S is None:
-        return L1, Q1
-
-    return L_prime, Q @ S
-
-
-# def adanystrom_update2(
-#     L1: torch.Tensor,
-#     Q1: torch.Tensor,
-#     v2: torch.Tensor,
-#     w1: float,
-#     w2: float,
-#     rank: int,
-# ):
-#     def A_mm(X):
-#         return weighted_eigen_plus_rank1_mm(L1=L1, Q1=Q1, v2=v2, B=X, w1=w1, w2=w2)
-
-#     return nystrom_approximation(A_mm, A_mm=A_mm, ndim=v2.numel(), rank=rank, device=L1.device, dtype=L1.dtype)
-
-class AdaNystrom(TensorTransform):
-    """Adagrad/RMSprop/Adam with Nyström-approximated covariance matrix.
-
-    Args:
-        rank (_type_): rank of Nyström approximation.
-        w1 (float, optional): weight of current covariance matrix. Defaults to 0.95.
-        w2 (float, optional): weight of new gradient in covariance matrix. Defaults to 0.05.
-        oversampling (int, optional): number of extra random vectors (top rank eigenvalues are kept). Defaults to 10.
-        eig_tol (float, optional):
-            removes eigenvalues this much smaller than largest eigenvalue when updating the preconditioner. Defaults to 1e-7.
-        damping (float, optional):
-            added to eigenvalues when updating the preconditioner. Defaults to 1e-8.
-        rdamping (float, optional):
-            added to eigenvalues when updating the preconditioner, relative to largest eigenvalue. Defaults to 0.
-        mm_tol (float, optional):
-            removes eigenvalues this much smaller than largest eigenvalue when computing the update. Defaults to 1e-7.
-        mm_truncate (int | None, optional):
-            uses top k eigenvalues to compute the update. Defaults to None.
-        mm_damping (float, optional):
-            added to eigenvalues when computing the update. Defaults to 1e-4.
-        mm_rdamping (float, optional):
-            added to eigenvalues when computing the update, relative to largest eigenvalue. Defaults to 0.
-        id_reg (float, optional):
-            multiplier to identity matrix added to preconditioner before computing update
-            If this value is given, solution from Nyström sketch-and-solve will be used to compute the update.
-            This value can't be too small (i.e. less than 1e-5) or the solver will be very unstable. Defaults to None.
-        concat_params (bool, optional):
-            whether to precondition all parameters at once if True, or each separately if False. Defaults to True.
-        update_freq (int, optional): update frequency. Defaults to 1.
-        inner (Chainable | None, optional): inner modules. Defaults to None.
-    """
-    def __init__(
-        self,
-        rank:int = 100,
-        beta=0.95,
-        oversampling: int = 10,
-        eig_tol: float | None = 1e-32,
-        damping: float = 0,
-        rdamping: float = 0,
-        mm_tol: float = 0,
-        mm_truncate: int | None = None,
-        mm_damping: float = 0,
-        mm_rdamping: float = 0,
-        id_reg: float | None = None,
-        orthogonalize_method: OrthogonalizeMethod = 'qr',
-        eigenbasis_optimizer: LREOptimizerBase | None = None,
-        orthogonalize_interval: int | None = 100,
-
-        concat_params: bool = True,
-        update_freq: int = 1,
-        inner: Chainable | None = None,
-    ):
-        defaults = locals().copy()
-        for k in ["self", "concat_params", "inner", "update_freq"]:
-            del defaults[k]
-
-        super().__init__(defaults, concat_params=concat_params, inner=inner, update_freq=update_freq)
-
-    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
-        state["step"] = state.get("step", 0) + 1
-        rank = setting["rank"]
-        device = tensor.device
-        dtype = tensor.dtype
-        beta = setting["beta"]
-
-        try:
-            if "L" not in state:
-                # use just tensor and zero L and Q with zero weight
-
-                L, Q = adanystrom_update(
-                    L1=torch.zeros(rank, device=device, dtype=dtype),
-                    Q1=torch.zeros((tensor.numel(), rank), device=device, dtype=dtype),
-                    v2=tensor.ravel(),
-                    w1=0,
-                    w2=1-beta,
-                    rank=rank,
-                    oversampling_p=setting["oversampling"],
-                    eig_tol=setting["eig_tol"],
-                    damping=setting["damping"],
-                    rdamping=setting["rdamping"],
-                    orthogonalize_method=setting["orthogonalize_method"],
-                )
-
-                state["L"] = state["L_reg"] = L
-                state["Q"] = state["Q_reg"] = Q
-
-            else:
-                L = state["L"]
-                Q = state["Q"]
-
-                w1 = beta
-                w2 = 1 - w1
-
-                # compute new factors (this function truncates them)
-                L_new, Q_new = adanystrom_update(
-                    L1=L,
-                    Q1=Q,
-                    v2=tensor.ravel(),
-                    w1=w1,
-                    w2=w2,
-                    rank=rank,
-                    oversampling_p=setting["oversampling"],
-                    eig_tol=setting["eig_tol"],
-                    damping=setting["damping"],
-                    rdamping=setting["rdamping"],
-                    orthogonalize_method=setting["orthogonalize_method"],
-                )
-
-                _eigengrad_update_state_(state=state, setting=setting, L_new=L_new, Q_new=Q_new)
-
-        except torch.linalg.LinAlgError:
-            pass
-
-    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
-        if "L_reg" not in state:
-            return tensor.clip(-0.1, 0.1)
-
-        if "eigenbasis_state" not in state:
-            state["eigenbasis_state"] = {}
-
-        return eigengrad_apply(
-            tensor=tensor,
-            L_reg = state["L_reg"],
-            Q_reg = state["Q_reg"],
-            beta = setting["beta"],
-            step = state["step"],
-            debias = True,
-            id_reg = setting["id_reg"],
-            eigenbasis_optimizer = setting["eigenbasis_optimizer"],
-            eigenbasis_state = state["eigenbasis_state"]
-        )

torchzero/modules/experimental/common_directions_whiten.py

@@ -1,142 +0,0 @@
-from collections import deque
-from typing import Literal
-
-import torch
-
-from torchzero.core import Chainable, TensorTransform
-from torchzero.linalg import matrix_power_eigh, torch_linalg, orthogonalize, OrthogonalizeMethod, regularize_eigh
-from torchzero.utils import TensorList, vec_to_tensors_
-
-
-def update_subspace_preconditioner_(
-    grad: torch.Tensor, # store grads and basis as vectors for matmul
-    basis: torch.Tensor, # ndim, k
-    accumulator_: torch.Tensor, # k, k
-    beta: float | None,
-):
-    projected = basis.T @ grad # k
-    outer = torch.outer(projected, projected)
-
-    if beta is None: accumulator_.add_(outer)
-    else: accumulator_.lerp_(outer, 1-beta)
-
-# yeah so I can also run subspace opts in this basis
-def apply_subspace_preconditioner(
-    tensor: torch.Tensor,
-    basis: torch.Tensor, # ndim, k
-    accumulator: torch.Tensor,
-    tol: float,
-    truncate: int | None,
-    damping: float,
-    rdamping: float,
-):
-    L, Q = torch_linalg.eigh(accumulator, retry_float64=True)
-    L, Q = regularize_eigh(L=L, Q=Q, truncate=truncate, tol=tol, damping=damping, rdamping=rdamping)
-
-    if L is None or Q is None:
-        return tensor.clip(-0.1, 0.1)
-
-    preconditioner = (Q * L.rsqrt().unsqueeze(-2)) @ Q.mH
-
-    tensor_projected = basis.T @ tensor # k
-    update_projected = preconditioner @ tensor_projected # k
-    return basis @ update_projected # d
-
-
-class CommonDirectionsWhiten(TensorTransform):
-    """Whitens in subspace spanned by history of gradient differences.
-
-    Args:
-        beta - for preconditioner itself in the basis.
-        basis_beta - how much basis is allowed to change.
-    """
-
-    def __init__(
-        self,
-        k: int = 100,
-        beta: float | None = 0.95,
-        basis_beta=0.95,
-        tol: float = 1e-7,
-        truncate: int | None = None,
-        damping: float = 1e-4,
-        rdamping: float = 0,
-        basis_type: Literal["gradients", "differences"] = "differences",
-        orthogonalize_method: OrthogonalizeMethod | None = 'newtonschulz',
-
-        concat_params: bool = True,
-        inner: Chainable | None = None,
-    ):
-        defaults = locals().copy()
-        for key in ["self", "inner", "concat_params"]:
-            del defaults[key]
-
-        super().__init__(defaults, concat_params=concat_params, inner=inner)
-
-    @torch.no_grad
-    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
-        g = tensor.ravel()
-        k = setting['k']
-        beta = setting['beta']
-        basis_beta = setting['basis_beta']
-        step = state.get("step", 0)
-        state["step"] = step + 1
-
-        # initialize history
-        if 'history' not in state:
-            state['history'] = deque(maxlen=k)
-            state['accumulator'] = torch.eye(k, device=g.device, dtype=g.dtype)
-            state['basis'] = torch.zeros(g.numel(), k, device=g.device, dtype=g.dtype)
-
-        history: deque = state['history']
-        accumulator = state['accumulator']
-        basis = state['basis']
-        history.append(g)
-
-        # stack history to new basis term, if history isn't full, fill with random vecs
-        if len(history) < k:
-            basis_t = torch.randn(g.numel(), k, device=g.device, dtype=g.dtype)
-            history_basis = torch.stack(tuple(history), -1)
-            basis_t[:, -len(history):] = history_basis
-
-        else:
-            basis_t = torch.stack(tuple(history), -1)
-
-        # in this case basis uses differences in gradients except last entry is the gradient
-        if setting["basis_type"] == "differences":
-            basis_t[:,:-1] = basis_t[:, :-1] - basis_t[:, 1:]
-
-        # normalize or orthonormalize new basis term
-        if setting["orthogonalize_method"] is not None:
-            basis_t = orthogonalize(basis_t, method = setting["orthogonalize_method"])
-        else:
-            basis_t = (basis_t - basis_t.mean()) / basis_t.std().clip(min=torch.finfo(g.dtype).tiny * 2)
-
-        # lerp basis
-        basis.lerp_(basis_t, 1-basis_beta)
-        basis = basis /  (1 - basis_beta ** (step+1)) # correct bias on basis EMA
-        update_subspace_preconditioner_(g, basis, accumulator, beta)
-
-    @torch.no_grad
-    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
-        g = tensor.ravel()
-
-        basis = state['basis']
-        accumulator = state['accumulator']
-        step = state["step"]
-        accumulator = accumulator / (1 - setting["beta"] ** (step+1)) # correct bias on accumulator EMA
-
-        try:
-            preconditioned = apply_subspace_preconditioner(
-                g,
-                basis,
-                accumulator,
-                tol=setting["tol"],
-                truncate=setting["truncate"],
-                damping=setting["damping"],
-                rdamping=setting["rdamping"],
-            )
-        except torch.linalg.LinAlgError:
-            preconditioned = g.clip(-0.1, 0.1)
-
-        return preconditioned.view_as(tensor)
-

torchzero/modules/experimental/eigen_sr1.py

@@ -1,182 +0,0 @@
-import torch
-
-from ...core import Transform
-from ...linalg.orthogonalize import orthogonalize, OrthogonalizeMethod
-from ...linalg.eigh import eigh_plus_uuT, regularize_eigh
-from ...utils import TensorList, unpack_states, vec_to_tensors_
-from ..opt_utils import safe_clip
-from .eigengrad import _eigengrad_update_state_, eigengrad_apply
-
-
-def sr1_u(L: torch.Tensor, Q: torch.Tensor, s:torch.Tensor, y: torch.Tensor, tol:float):
-    """u from u u^T correction and its sign"""
-    r = y - torch.linalg.multi_dot([Q, L.diag_embed(), Q.T, s]) # pylint:disable=not-callable
-    rs = r.dot(s)
-
-    if rs.abs() < tol * torch.linalg.vector_norm(r) * torch.linalg.vector_norm(s): # pylint:disable=not-callable
-        return None, None
-
-    u = r / rs.abs().sqrt()
-    return u, torch.sign(rs)
-
-class EigenSR1(Transform):
-    def __init__(
-        self,
-        rank: int = 100,
-        tol: float = 1e-32,
-        eig_tol: float | None = None,
-        damping: float = 0,
-        rdamping: float = 0,
-        abs: bool = False,
-        mm_tol: float = 1e-7,
-        mm_truncate: int | None = None,
-        mm_damping: float = 1e-4,
-        mm_rdamping: float = 0,
-        mm_abs: bool = True,
-        id_reg: float | None = None,
-        column_space_tol=1e-9,
-        beta: float = 0.95,
-        balance_tol: float = 10,
-        balance_strength: float = 1e-1,
-
-        eigenbasis_optimizer = None,
-        update_freq: int = 1,
-        init_steps: int = 10,
-        orthogonalize_interval: int | None = 1,
-        orthogonalize_method: OrthogonalizeMethod = 'qr',
-
-        hvp_method = "autograd",
-        h = 1e-3,
-        inner = None,
-
-    ):
-        defaults = locals().copy()
-        for k in ["self", "inner"]:
-            del defaults[k]
-
-        super().__init__(defaults)
-
-    def update_states(self, objective, states, settings):
-        fs = settings[0]
-        step = self.increment_counter("step", 0)
-
-        if step % fs["update_freq"] == 0:
-
-            params = TensorList(objective.params)
-
-            # compute y as hessian-vector product with s (random vecs during init steps)
-            if ("p_prev" not in self.global_state) or (step < fs["init_steps"]):
-                s_list = params.sample_like('rademacher')
-
-            else:
-                p_prev = self.global_state["p_prev"]
-                s_list = params - p_prev
-
-            if s_list.dot(s_list) < torch.finfo(s_list[0].dtype).tiny * 2:
-                s_list = params.sample_like('rademacher')
-
-            self.global_state["p_prev"] = params
-
-            # compute y as hessian-vector product with s
-            Hz_list, _ = objective.hessian_vector_product(s_list, rgrad=None, at_x0=True, hvp_method=fs["hvp_method"], h=fs["h"])
-
-            s = torch.cat([t.ravel() for t in s_list])
-            y = torch.cat([t.ravel() for t in Hz_list])
-
-            # keep track of exponential moving average of hessian diagonal and balance eigenvalues
-            if (fs["balance_strength"] != 0) and (step > fs["init_steps"]) and ("L" in self.global_state):
-
-                D = s * y # hutchinson estimator
-                exp_avg = self.global_state.get("exp_avg", None)
-
-                if exp_avg is None:
-                    exp_avg = self.global_state["exp_avg"] = D
-
-                else:
-                    exp_avg.lerp_(D, weight=1-fs["beta"])
-
-                L = self.global_state["L"]
-                L_abs = L.abs()
-                tau = L_abs.amax() / exp_avg.abs().amax()
-
-                if tau > fs["balance_tol"]:
-                    L_balanced = L_abs.pow((1 / tau) ** (1 / fs["balance_strength"])).copysign(L)
-                    self.global_state["L"] = torch.where(L_abs > 1, L_balanced, L)
-
-            # initialize L and Q on 1st step
-            if "L" not in self.global_state:
-
-                L = torch.zeros(1, dtype=s.dtype, device=s.device) # rank, rank
-                Q = torch.zeros([s.numel(), 1], dtype=s.dtype, device=s.device) # ndim, rank
-
-                u, sign = sr1_u(L=L, Q=Q, s=s, y=y, tol=0)
-                assert u is not None and sign is not None
-
-                # for uu^T u is eigenvector and u^T u is eigenvalue
-                norm = torch.linalg.vector_norm(u).clip(min=torch.finfo(u.dtype).tiny * 2) # pylint:disable=not-callable
-
-                self.global_state["L"] = self.global_state["L_reg"] = (u.dot(u).unsqueeze(0) / norm) * sign # (rank,)
-                self.global_state["Q"] = self.global_state["Q_reg"] = u.unsqueeze(-1) / norm # (m, rank)
-
-            # update hessian
-            else:
-                try:
-                    L = self.global_state["L"]
-                    Q = self.global_state["Q"]
-
-                    H_step = self.increment_counter("H_step", start=0)
-                    if H_step % fs["orthogonalize_interval"] == 0:
-                        Q = orthogonalize(Q, method=fs["orthogonalize_method"])
-
-                    u, sign = sr1_u(L=L, Q=Q, s=s, y=y, tol=fs["tol"])
-
-                    if (u is not None) and (sign is not None):
-
-                        # compute new factors
-                        L_new, Q_new = eigh_plus_uuT(L, Q, u, tol=fs["column_space_tol"], alpha=sign.item(), retry_float64=True)
-
-                        # truncate/regularize new factors (those go into the accumulator)
-                        L_new, Q_new = regularize_eigh(L=L_new, Q=Q_new, truncate=min(fs["rank"], s.numel()),
-                                                      tol=fs["eig_tol"], damping=fs["damping"], rdamping=fs["rdamping"])
-
-                        _eigengrad_update_state_(state=self.global_state, setting=fs, L_new=L_new, Q_new=Q_new)
-
-                except torch.linalg.LinAlgError:
-                    pass
-
-
-
-    def apply_states(self, objective, states, settings):
-        fs = settings[0]
-        updates = objective.get_updates()
-
-        if "eigenbasis_state" not in self.global_state:
-            self.global_state["eigenbasis_state"] = {}
-
-        step = self.global_state["step"] # starts at 0
-        if step < fs["init_steps"]:
-
-            # skip update first init_steps to let hessian kick-start
-            objective.stop = True
-            objective.skip_update = True
-            return objective
-
-        if "L_reg" not in self.global_state:
-            TensorList(updates).clip_(-0.1, 0.1)
-            return objective
-
-        dir = eigengrad_apply(
-            tensor = torch.cat([t.ravel() for t in updates]),
-            L_reg = self.global_state["L_reg"],
-            Q_reg = self.global_state["Q_reg"],
-            beta = None,
-            step = None,
-            debias = False,
-            id_reg = fs["id_reg"],
-            eigenbasis_optimizer = fs["eigenbasis_optimizer"],
-            eigenbasis_state = self.global_state["eigenbasis_state"],
-            whiten_fn = lambda x: x
-        )
-
-        vec_to_tensors_(dir, updates)
-        return objective