torchzero 0.3.9__py3-none-any.whl → 0.3.10__py3-none-any.whl

torchzero/modules/experimental/newtonnewton.py

@@ -0,0 +1,88 @@
+import itertools
+import warnings
+from collections.abc import Callable
+from contextlib import nullcontext
+from functools import partial
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, apply_transform
+from ...utils import TensorList, vec_to_tensors
+from ...utils.derivatives import (
+    hessian_list_to_mat,
+    jacobian_wrt,
+)
+from ..second_order.newton import (
+    cholesky_solve,
+    eigh_solve,
+    least_squares_solve,
+    lu_solve,
+)
+
+
+class NewtonNewton(Module):
+    """
+    Method that I thought of and then it worked.
+
+    1. Calculate newton step by solving Hx=g
+
+    2. Calculate jacobian of x wrt parameters and call it H2
+
+    3. Solve H2 x2 = x for x2.
+
+    4. Optionally, repeat (if order is higher than 3.)
+
+    Memory is n^order. It tends to converge faster on convex functions, but can be unstable on non-convex. Orders higher than 3 are usually too unsable and have little benefit.
+    """
+    def __init__(
+        self,
+        reg: float = 1e-6,
+        order: int = 3,
+        search_negative: bool = False,
+        vectorize: bool = True,
+        eigval_tfm: Callable[[torch.Tensor], torch.Tensor] | None = None,
+    ):
+        defaults = dict(order=order, reg=reg, vectorize=vectorize, eigval_tfm=eigval_tfm, search_negative=search_negative)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        params = TensorList(var.params)
+        closure = var.closure
+        if closure is None: raise RuntimeError('NewtonCG requires closure')
+
+        settings = self.settings[params[0]]
+        reg = settings['reg']
+        vectorize = settings['vectorize']
+        order = settings['order']
+        search_negative = settings['search_negative']
+        eigval_tfm = settings['eigval_tfm']
+
+        # ------------------------ calculate grad and hessian ------------------------ #
+        with torch.enable_grad():
+            loss = var.loss = var.loss_approx = closure(False)
+            g_list = torch.autograd.grad(loss, params, create_graph=True)
+            var.grad = list(g_list)
+
+            xp = torch.cat([t.ravel() for t in g_list])
+            I = torch.eye(xp.numel(), dtype=xp.dtype, device=xp.device)
+
+            for o in range(2, order + 1):
+                is_last = o == order
+                H_list = jacobian_wrt([xp], params, create_graph=not is_last, batched=vectorize)
+                with torch.no_grad() if is_last else nullcontext():
+                    H = hessian_list_to_mat(H_list)
+                    if reg != 0: H = H + I * reg
+
+                    x = None
+                    if search_negative or (is_last and eigval_tfm is not None):
+                        x = eigh_solve(H, xp, eigval_tfm, search_negative=search_negative)
+                    if x is None: x = cholesky_solve(H, xp)
+                    if x is None: x = lu_solve(H, xp)
+                    if x is None: x = least_squares_solve(H, xp)
+                    xp = x.squeeze()
+
+        var.update = vec_to_tensors(xp, params)
+        return var
+

torchzero/modules/experimental/reduce_outward_lr.py

@@ -1,30 +1,33 @@
 import torch
 
 from ...core import Target, Transform
-from ...utils import TensorList
+from ...utils import TensorList, unpack_states, unpack_dicts
 
 class ReduceOutwardLR(Transform):
     """
     When update sign matches weight sign, the learning rate for that weight is multiplied by `mul`.
 
     This means updates that move weights towards zero have higher learning rates.
+
+    A note on this is that it sounded good but its really bad in practice.
     """
     def __init__(self, mul = 0.5, use_grad=False, invert=False, target: Target = 'update'):
         defaults = dict(mul=mul, use_grad=use_grad, invert=invert)
         super().__init__(defaults, uses_grad=use_grad, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         params = TensorList(params)
         tensors = TensorList(tensors)
 
-        mul = self.get_settings('mul', params=params)
-        s = self.settings[params[0]]
+        mul = [s['mul'] for s in settings]
+        s = settings[0]
         use_grad = s['use_grad']
         invert = s['invert']
 
-        if use_grad: cur = vars.get_grad()
+        if use_grad: cur = grads
         else: cur = tensors
+        assert cur is not None
 
         # mask of weights where sign matches with update sign (minus ascent sign), multiplied by `mul`.
         if invert: mask = (params * cur) > 0

torchzero/modules/experimental/soapy.py

@@ -2,147 +2,22 @@ from operator import itemgetter
 
 import torch
 
-from ...core import Chainable, Transform, apply
+from ...core import Chainable, Transform
 from ..optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
-
-@torch.no_grad
-def update_soap_covariances_(
-    grad: torch.Tensor,
-    GGs_: list[torch.Tensor | None],
-    beta: float | None,
-):
-    for i, GG in enumerate(GGs_):
-        if GG is None: continue
-
-        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
+from ..optimizers.soap import (
+    update_soap_covariances_,
+    get_orthogonal_matrix,
+    get_orthogonal_matrix_QR,
+    project,
+    project_back,
+)
 
 class SOAPY(Transform):
-    """SOAP but uses scaled gradient differences
-
-    new args
-
-    scale by s whether to scale gradient differences by parameter differences
+    """Adam but uses scaled gradient differences for GGᵀ. Please note that this is experimental and isn't guaranteed to work.
 
-    y_to_ema2 whether to use gradient differences for exponential moving average too
+    New args:
+        scale_by_s - whether to scale gradient differences by parameter differences
+        y_to_ema2 - whether to use gradient differences for exponential moving average too
     """
     def __init__(
         self,
@@ -178,16 +53,14 @@ class SOAPY(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         updates = []
         # update preconditioners
-        for i,(p,t) in enumerate(zip(params, tensors)):
-            state = self.state[p]
-            settings = self.settings[p]
+        for i,(p,t, state, setting) in enumerate(zip(params, tensors, states, settings)):
             beta1, beta2, shampoo_beta, merge_small, max_dim, precondition_1d, eps, alpha = itemgetter(
-                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'alpha')(settings)
-            scale_by_s = settings['scale_by_s']
-            y_to_ema2 = settings['y_to_ema2']
+                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'alpha')(setting)
+            scale_by_s = setting['scale_by_s']
+            y_to_ema2 = setting['y_to_ema2']
 
             if merge_small:
                 t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
@@ -268,7 +141,7 @@ class SOAPY(Transform):
             if z_projected is not None:
                 update = project_back(update, state["Q"])
 
-            if settings['bias_correction']:
+            if setting['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
@@ -284,7 +157,7 @@ class SOAPY(Transform):
             # Update is done after the gradient step to avoid using current gradients in the projection.
             if state['GG'] is not None:
                 update_soap_covariances_(y, state['GG'], shampoo_beta)
-                if state['step'] % settings['precond_freq'] == 0:
+                if state['step'] % setting['precond_freq'] == 0:
                     state['Q'], state['exp_avg_sq'] = get_orthogonal_matrix_QR(exp_avg_sq, state['GG'], state['Q'])
 
         return updates