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

torchzero/modules/smoothing/sampling.py

@@ -0,0 +1,300 @@
+import math
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Sequence
+from contextlib import nullcontext
+from functools import partial
+from typing import Literal, cast
+
+import torch
+
+from ...core import Chainable, Modular, Module, Var
+from ...core.reformulation import Reformulation
+from ...utils import Distributions, NumberList, TensorList
+from ..termination import TerminationCriteriaBase, make_termination_criteria
+
+
+def _reset_except_self(optimizer: Modular, var: Var, self: Module):
+    for m in optimizer.unrolled_modules:
+        if m is not self:
+            m.reset()
+
+
+class GradientSampling(Reformulation):
+    """Samples and aggregates gradients and values at perturbed points.
+
+    This module can be used for gaussian homotopy and gradient sampling methods.
+
+    Args:
+        modules (Chainable | None, optional):
+            modules that will be optimizing the modified objective.
+            if None, returns gradient of the modified objective as the update. Defaults to None.
+        sigma (float, optional): initial magnitude of the perturbations. Defaults to 1.
+        n (int, optional): number of perturbations per step. Defaults to 100.
+        aggregate (str, optional):
+            how to aggregate values and gradients
+            - "mean" - uses mean of the gradients, as in gaussian homotopy.
+            - "max" - uses element-wise maximum of the gradients.
+            - "min" - uses element-wise minimum of the gradients.
+            - "min-norm" - picks gradient with the lowest norm.
+
+            Defaults to 'mean'.
+        distribution (Distributions, optional): distribution for random perturbations. Defaults to 'gaussian'.
+        include_x0 (bool, optional): whether to include gradient at un-perturbed point. Defaults to True.
+        fixed (bool, optional):
+            if True, perturbations do not get replaced by new random perturbations until termination criteria is satisfied. Defaults to True.
+        pre_generate (bool, optional):
+            if True, perturbations are pre-generated before each step.
+            This requires more memory to store all of them,
+            but ensures they do not change when closure is evaluated multiple times.
+            Defaults to True.
+        termination (TerminationCriteriaBase | Sequence[TerminationCriteriaBase] | None, optional):
+            a termination criteria module, sigma will be multiplied by ``decay`` when termination criteria is satisfied,
+            and new perturbations will be generated if ``fixed``. Defaults to None.
+        decay (float, optional): sigma multiplier on termination criteria. Defaults to 2/3.
+        reset_on_termination (bool, optional): whether to reset states of all other modules on termination. Defaults to True.
+        sigma_strategy (str | None, optional):
+            strategy for adapting sigma. If condition is satisfied, sigma is multiplied by ``sigma_nplus``,
+            otherwise it is multiplied by ``sigma_nminus``.
+            - "grad-norm" - at least ``sigma_target`` gradients should have lower norm than at un-perturbed point.
+            - "value" - at least ``sigma_target`` values (losses) should be lower than at un-perturbed point.
+            - None - doesn't use adaptive sigma.
+
+            This introduces a side-effect to the closure, so it should be left at None of you use
+            trust region or line search to optimize the modified objective.
+            Defaults to None.
+        sigma_target (int, optional):
+            number of elements to satisfy the condition in ``sigma_strategy``. Defaults to 1.
+        sigma_nplus (float, optional): sigma multiplier when ``sigma_strategy`` condition is satisfied. Defaults to 4/3.
+        sigma_nminus (float, optional): sigma multiplier when ``sigma_strategy`` condition is not satisfied. Defaults to 2/3.
+        seed (int | None, optional): seed. Defaults to None.
+    """
+    def __init__(
+        self,
+        modules: Chainable | None = None,
+        sigma: float = 1.,
+        n:int = 100,
+        aggregate: Literal['mean', 'max', 'min', 'min-norm', 'min-value'] = 'mean',
+        distribution: Distributions = 'gaussian',
+        include_x0: bool = True,
+
+        fixed: bool=True,
+        pre_generate: bool = True,
+        termination: TerminationCriteriaBase | Sequence[TerminationCriteriaBase] | None = None,
+        decay: float = 2/3,
+        reset_on_termination: bool = True,
+
+        sigma_strategy: Literal['grad-norm', 'value'] | None = None,
+        sigma_target: int | float = 0.2,
+        sigma_nplus: float = 4/3,
+        sigma_nminus: float = 2/3,
+
+        seed: int | None = None,
+    ):
+
+        defaults = dict(sigma=sigma, n=n, aggregate=aggregate, distribution=distribution, seed=seed, include_x0=include_x0, fixed=fixed, decay=decay, reset_on_termination=reset_on_termination, sigma_strategy=sigma_strategy, sigma_target=sigma_target, sigma_nplus=sigma_nplus, sigma_nminus=sigma_nminus, pre_generate=pre_generate)
+        super().__init__(defaults, modules)
+
+        if termination is not None:
+            self.set_child('termination', make_termination_criteria(extra=termination))
+
+    @torch.no_grad
+    def pre_step(self, var):
+        params = TensorList(var.params)
+
+        fixed = self.defaults['fixed']
+
+        # check termination criteria
+        if 'termination' in self.children:
+            termination = cast(TerminationCriteriaBase, self.children['termination'])
+            if termination.should_terminate(var):
+
+                # decay sigmas
+                states = [self.state[p] for p in params]
+                settings = [self.settings[p] for p in params]
+
+                for state, setting in zip(states, settings):
+                    if 'sigma' not in state: state['sigma'] = setting['sigma']
+                    state['sigma'] *= setting['decay']
+
+                # reset on sigmas decay
+                if self.defaults['reset_on_termination']:
+                    var.post_step_hooks.append(partial(_reset_except_self, self=self))
+
+                # clear perturbations
+                self.global_state.pop('perts', None)
+
+        # pre-generate perturbations if not already pre-generated or not fixed
+        if self.defaults['pre_generate'] and (('perts' not in self.global_state) or (not fixed)):
+            states = [self.state[p] for p in params]
+            settings = [self.settings[p] for p in params]
+
+            n = self.defaults['n'] - self.defaults['include_x0']
+            generator = self.get_generator(params[0].device, self.defaults['seed'])
+
+            perts = [params.sample_like(self.defaults['distribution'], generator=generator) for _ in range(n)]
+
+            self.global_state['perts'] = perts
+
+    @torch.no_grad
+    def closure(self, backward, closure, params, var):
+        params = TensorList(params)
+        loss_agg = None
+        grad_agg = None
+
+        states = [self.state[p] for p in params]
+        settings = [self.settings[p] for p in params]
+        sigma_inits = [s['sigma'] for s in settings]
+        sigmas = [s.setdefault('sigma', si) for s, si in zip(states, sigma_inits)]
+
+        include_x0 = self.defaults['include_x0']
+        pre_generate = self.defaults['pre_generate']
+        aggregate: Literal['mean', 'max', 'min', 'min-norm', 'min-value'] = self.defaults['aggregate']
+        sigma_strategy: Literal['grad-norm', 'value'] | None = self.defaults['sigma_strategy']
+        distribution = self.defaults['distribution']
+        generator = self.get_generator(params[0].device, self.defaults['seed'])
+
+
+        n_finite = 0
+        n_good = 0
+        f_0 = None; g_0 = None
+
+        # evaluate at x_0
+        if include_x0:
+            f_0 = cast(torch.Tensor, var.get_loss(backward=backward))
+
+            isfinite = math.isfinite(f_0)
+            if isfinite:
+                n_finite += 1
+                loss_agg = f_0
+
+            if backward:
+                g_0 = var.get_grad()
+                if isfinite: grad_agg = g_0
+
+        # evaluate at x_0 + p for each perturbation
+        if pre_generate:
+            perts = self.global_state['perts']
+        else:
+            perts = [None] * (self.defaults['n'] - include_x0)
+
+        x_0 = [p.clone() for p in params]
+
+        for pert in perts:
+            loss = None; grad = None
+
+            # generate if not pre-generated
+            if pert is None:
+                pert = params.sample_like(distribution, generator=generator)
+
+            # add perturbation and evaluate
+            pert = pert * sigmas
+            torch._foreach_add_(params, pert)
+
+            with torch.enable_grad() if backward else nullcontext():
+                loss = closure(backward)
+
+            if math.isfinite(loss):
+                n_finite += 1
+
+                # add loss
+                if loss_agg is None:
+                    loss_agg = loss
+                else:
+                    if aggregate == 'mean':
+                        loss_agg += loss
+
+                    elif (aggregate=='min') or (aggregate=='min-value') or (aggregate=='min-norm' and not backward):
+                        loss_agg = loss_agg.clamp(max=loss)
+
+                    elif aggregate == 'max':
+                        loss_agg = loss_agg.clamp(min=loss)
+
+                # add grad
+                if backward:
+                    grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+                    if grad_agg is None:
+                        grad_agg = grad
+                    else:
+                        if aggregate == 'mean':
+                            torch._foreach_add_(grad_agg, grad)
+
+                        elif aggregate == 'min':
+                            grad_agg_abs = torch._foreach_abs(grad_agg)
+                            torch._foreach_minimum_(grad_agg_abs, torch._foreach_abs(grad))
+                            grad_agg = [g_abs.copysign(g) for g_abs, g in zip(grad_agg_abs, grad_agg)]
+
+                        elif aggregate == 'max':
+                            grad_agg_abs = torch._foreach_abs(grad_agg)
+                            torch._foreach_maximum_(grad_agg_abs, torch._foreach_abs(grad))
+                            grad_agg = [g_abs.copysign(g) for g_abs, g in zip(grad_agg_abs, grad_agg)]
+
+                        elif aggregate == 'min-norm':
+                            if TensorList(grad).global_vector_norm() < TensorList(grad_agg).global_vector_norm():
+                                grad_agg = grad
+                                loss_agg = loss
+
+                        elif aggregate == 'min-value':
+                            if loss < loss_agg:
+                                grad_agg = grad
+                                loss_agg = loss
+
+            # undo perturbation
+            torch._foreach_copy_(params, x_0)
+
+            # adaptive sigma
+            # by value
+            if sigma_strategy == 'value':
+                if f_0 is None:
+                    with torch.enable_grad() if backward else nullcontext():
+                        f_0 = closure(False)
+
+                if loss < f_0:
+                    n_good += 1
+
+            # by gradient norm
+            elif sigma_strategy == 'grad-norm' and backward and math.isfinite(loss):
+                assert grad is not None
+                if g_0 is None:
+                    with torch.enable_grad() if backward else nullcontext():
+                        closure()
+                        g_0 = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+
+                if TensorList(grad).global_vector_norm() < TensorList(g_0).global_vector_norm():
+                    n_good += 1
+
+        # update sigma if strategy is enabled
+        if sigma_strategy is not None:
+
+            sigma_target = self.defaults['sigma_target']
+            if isinstance(sigma_target, float):
+                sigma_target = int(max(1, n_finite * sigma_target))
+
+            if n_good >= sigma_target:
+                key = 'sigma_nplus'
+            else:
+                key = 'sigma_nminus'
+
+            for p in params:
+                self.state[p]['sigma'] *= self.settings[p][key]
+
+        # if no finite losses, just return inf
+        if n_finite == 0:
+            assert loss_agg is None and grad_agg is None
+            loss = torch.tensor(torch.inf, dtype=params[0].dtype, device=params[0].device)
+            grad = [torch.full_like(p, torch.inf) for p in params]
+            return loss, grad
+
+        assert loss_agg is not None
+
+        # no post processing needed when aggregate is 'max', 'min', 'min-norm', 'min-value'
+        if aggregate != 'mean':
+            return loss_agg, grad_agg
+
+        # on mean divide by number of evals
+        loss_agg /= n_finite
+
+        if backward:
+            assert grad_agg is not None
+            torch._foreach_div_(grad_agg, n_finite)
+
+        return loss_agg, grad_agg

torchzero/modules/step_size/__init__.py

@@ -0,0 +1,2 @@
+from .lr import LR, StepSize, Warmup, WarmupNormClip, RandomStepSize
+from .adaptive import PolyakStepSize, BarzilaiBorwein, BBStab, AdGD

torchzero/modules/step_size/adaptive.py

@@ -0,0 +1,387 @@
+"""Various step size strategies"""
+import math
+from operator import itemgetter
+from typing import Any, Literal
+
+import torch
+
+from ...core import Chainable, Transform
+from ...utils import NumberList, TensorList, tofloat, unpack_dicts, unpack_states
+from ...utils.linalg.linear_operator import ScaledIdentity
+from ..functional import epsilon_step_size
+
+def _acceptable_alpha(alpha, param:torch.Tensor):
+    finfo = torch.finfo(param.dtype)
+    if (alpha is None) or (alpha < finfo.tiny*2) or (not math.isfinite(alpha)) or (alpha > finfo.max/2):
+        return False
+    return True
+
+def _get_H(self: Transform, var):
+    n = sum(p.numel() for p in var.params)
+    p = var.params[0]
+    alpha = self.global_state.get('alpha', 1)
+    if not _acceptable_alpha(alpha, p): alpha = 1
+
+    return ScaledIdentity(1 / alpha, shape=(n,n), device=p.device, dtype=p.dtype)
+
+
+class PolyakStepSize(Transform):
+    """Polyak's subgradient method with known or unknown f*.
+
+    Args:
+        f_star (float | Mone, optional):
+            minimal possible value of the objective function. If not known, set to ``None``. Defaults to 0.
+        y (float, optional):
+            when ``f_star`` is set to None, it is calculated as ``f_best - y``.
+        y_decay (float, optional):
+            ``y`` is multiplied by ``(1 - y_decay)`` after each step. Defaults to 1e-3.
+        max (float | None, optional): maximum possible step size. Defaults to None.
+        use_grad (bool, optional):
+            if True, uses dot product of update and gradient to compute the step size.
+            Otherwise, dot product of update with itself is used.
+        alpha (float, optional): multiplier to Polyak step-size. Defaults to 1.
+    """
+    def __init__(self, f_star: float | None = 0, y: float = 1, y_decay: float = 1e-3, max: float | None = None, use_grad=True, alpha: float = 1, inner: Chainable | None = None):
+
+        defaults = dict(alpha=alpha, max=max, f_star=f_star, y=y, y_decay=y_decay)
+        super().__init__(defaults, uses_grad=use_grad, uses_loss=True, inner=inner)
+
+    @torch.no_grad
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
+        assert grads is not None and loss is not None
+        tensors = TensorList(tensors)
+        grads = TensorList(grads)
+
+        # load variables
+        max, f_star, y, y_decay = itemgetter('max', 'f_star', 'y', 'y_decay')(settings[0])
+        y_val = self.global_state.get('y_val', y)
+        f_best = self.global_state.get('f_best', None)
+
+        # gg
+        if self._uses_grad: gg = tensors.dot(grads)
+        else: gg = tensors.dot(tensors)
+
+        # store loss
+        if f_best is None or loss < f_best: f_best = tofloat(loss)
+        if f_star is None: f_star = f_best - y_val
+
+        # calculate the step size
+        if gg <= torch.finfo(gg.dtype).tiny * 2: alpha = 0 # converged
+        else: alpha = (loss - f_star) / gg
+
+        # clip
+        if max is not None:
+            if alpha > max: alpha = max
+
+        # store state
+        self.global_state['f_best'] = f_best
+        self.global_state['y_val'] = y_val * (1 - y_decay)
+        self.global_state['alpha'] = alpha
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        alpha = self.global_state.get('alpha', 1)
+        if not _acceptable_alpha(alpha, tensors[0]): alpha = epsilon_step_size(TensorList(tensors))
+
+        torch._foreach_mul_(tensors, alpha * unpack_dicts(settings, 'alpha', cls=NumberList))
+        return tensors
+
+    def get_H(self, var):
+        return _get_H(self, var)
+
+
+def _bb_short(s: TensorList, y: TensorList, sy, eps):
+    yy = y.dot(y)
+    if yy < eps:
+        if sy < eps: return None # try to fallback on long
+        ss = s.dot(s)
+        return ss/sy
+    return sy/yy
+
+def _bb_long(s: TensorList, y: TensorList, sy, eps):
+    ss = s.dot(s)
+    if sy < eps:
+        yy = y.dot(y) # try to fallback on short
+        if yy < eps: return None
+        return sy/yy
+    return ss/sy
+
+def _bb_geom(s: TensorList, y: TensorList, sy, eps, fallback:bool):
+    short = _bb_short(s, y, sy, eps)
+    long = _bb_long(s, y, sy, eps)
+    if long is None or short is None:
+        if fallback:
+            if short is not None: return short
+            if long is not None: return long
+        return None
+    return (short * long) ** 0.5
+
+class BarzilaiBorwein(Transform):
+    """Barzilai-Borwein step size method.
+
+    Args:
+        type (str, optional):
+            one of "short" with formula sᵀy/yᵀy, "long" with formula sᵀs/sᵀy, or "geom" to use geometric mean of short and long.
+            Defaults to "geom".
+        fallback (float, optional): step size when denominator is less than 0 (will happen on negative curvature). Defaults to 1e-3.
+        inner (Chainable | None, optional):
+            step size will be applied to outputs of this module. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        type: Literal["long", "short", "geom", "geom-fallback"] = "geom",
+        alpha_0: float = 1e-7,
+        use_grad=True,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(type=type, alpha_0=alpha_0)
+        super().__init__(defaults, uses_grad=use_grad, inner=inner)
+
+    def reset_for_online(self):
+        super().reset_for_online()
+        self.clear_state_keys('prev_g')
+        self.global_state['reset'] = True
+
+    @torch.no_grad
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        prev_p, prev_g = unpack_states(states, tensors, 'prev_p', 'prev_g', cls=TensorList)
+        type = self.defaults['type']
+
+        g = grads if self._uses_grad else tensors
+        assert g is not None
+
+        reset = self.global_state.get('reset', False)
+        self.global_state.pop('reset', None)
+
+        if step != 0 and not reset:
+            s = params-prev_p
+            y = g-prev_g
+            sy = s.dot(y)
+            eps = torch.finfo(sy.dtype).tiny * 2
+
+            if type == 'short': alpha = _bb_short(s, y, sy, eps)
+            elif type == 'long': alpha = _bb_long(s, y, sy, eps)
+            elif type == 'geom': alpha = _bb_geom(s, y, sy, eps, fallback=False)
+            elif type == 'geom-fallback': alpha = _bb_geom(s, y, sy, eps, fallback=True)
+            else: raise ValueError(type)
+
+            # if alpha is not None:
+            self.global_state['alpha'] = alpha
+
+        prev_p.copy_(params)
+        prev_g.copy_(g)
+
+    def get_H(self, var):
+        return _get_H(self, var)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        alpha = self.global_state.get('alpha', None)
+
+        if not _acceptable_alpha(alpha, tensors[0]):
+            alpha = epsilon_step_size(TensorList(tensors), settings[0]['alpha_0'])
+
+        torch._foreach_mul_(tensors, alpha)
+        return tensors
+
+
+class BBStab(Transform):
+    """Stabilized Barzilai-Borwein method (https://arxiv.org/abs/1907.06409).
+
+    This clips the norm of the Barzilai-Borwein update by ``delta``, where ``delta`` can be adaptive if ``c`` is specified.
+
+    Args:
+        c (float, optional):
+            adaptive delta parameter. If ``delta`` is set to None, first ``inf_iters`` updates are performed
+            with non-stabilized Barzilai-Borwein step size. Then delta is set to norm of
+            the update that had the smallest norm, and multiplied by ``c``. Defaults to 0.2.
+        delta (float | None, optional):
+            Barzilai-Borwein update is clipped to this value. Set to ``None`` to use an adaptive choice. Defaults to None.
+        type (str, optional):
+            one of "short" with formula sᵀy/yᵀy, "long" with formula sᵀs/sᵀy, or "geom" to use geometric mean of short and long.
+            Defaults to "geom". Note that "long" corresponds to BB1stab and "short" to BB2stab,
+            however I found that "geom" works really well.
+        inner (Chainable | None, optional):
+            step size will be applied to outputs of this module. Defaults to None.
+
+    """
+    def __init__(
+        self,
+        c=0.2,
+        delta:float | None = None,
+        type: Literal["long", "short", "geom", "geom-fallback"] = "geom",
+        alpha_0: float = 1e-7,
+        use_grad=True,
+        inf_iters: int = 3,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(type=type,alpha_0=alpha_0, c=c, delta=delta, inf_iters=inf_iters)
+        super().__init__(defaults, uses_grad=use_grad, inner=inner)
+
+    def reset_for_online(self):
+        super().reset_for_online()
+        self.clear_state_keys('prev_g')
+        self.global_state['reset'] = True
+
+    @torch.no_grad
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        prev_p, prev_g = unpack_states(states, tensors, 'prev_p', 'prev_g', cls=TensorList)
+        type = self.defaults['type']
+        c = self.defaults['c']
+        delta = self.defaults['delta']
+        inf_iters = self.defaults['inf_iters']
+
+        g = grads if self._uses_grad else tensors
+        assert g is not None
+        g = TensorList(g)
+
+        reset = self.global_state.get('reset', False)
+        self.global_state.pop('reset', None)
+
+        if step != 0 and not reset:
+            s = params-prev_p
+            y = g-prev_g
+            sy = s.dot(y)
+            eps = torch.finfo(sy.dtype).tiny
+
+            if type == 'short': alpha = _bb_short(s, y, sy, eps)
+            elif type == 'long': alpha = _bb_long(s, y, sy, eps)
+            elif type == 'geom': alpha = _bb_geom(s, y, sy, eps, fallback=False)
+            elif type == 'geom-fallback': alpha = _bb_geom(s, y, sy, eps, fallback=True)
+            else: raise ValueError(type)
+
+            if alpha is not None:
+
+                # adaptive delta
+                if delta is None:
+                    niters = self.global_state.get('niters', 0) # this accounts for skipped negative curvature steps
+                    self.global_state['niters'] = niters + 1
+
+
+                    if niters == 0: pass # 1st iteration is scaled GD step, shouldn't be used to find s_norm_min
+                    elif niters <= inf_iters:
+                        s_norm_min = self.global_state.get('s_norm_min', None)
+                        if s_norm_min is None: s_norm_min = s.global_vector_norm()
+                        else: s_norm_min = min(s_norm_min, s.global_vector_norm())
+                        self.global_state['s_norm_min'] = s_norm_min
+                        # first few steps use delta=inf, so delta remains None
+
+                    else:
+                        delta = c * self.global_state['s_norm_min']
+
+                if delta is None: # delta is inf for first few steps
+                    self.global_state['alpha'] = alpha
+
+                # BBStab step size
+                else:
+                    a_stab = delta / g.global_vector_norm()
+                    self.global_state['alpha'] = min(alpha, a_stab)
+
+        prev_p.copy_(params)
+        prev_g.copy_(g)
+
+    def get_H(self, var):
+        return _get_H(self, var)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        alpha = self.global_state.get('alpha', None)
+
+        if not _acceptable_alpha(alpha, tensors[0]):
+            alpha = epsilon_step_size(TensorList(tensors), settings[0]['alpha_0'])
+
+        torch._foreach_mul_(tensors, alpha)
+        return tensors
+
+
+class AdGD(Transform):
+    """AdGD and AdGD-2 (https://arxiv.org/abs/2308.02261)"""
+    def __init__(self, variant:Literal[1,2]=2, alpha_0:float = 1e-7, sqrt:bool=True, use_grad=True, inner: Chainable | None = None,):
+        defaults = dict(variant=variant, alpha_0=alpha_0, sqrt=sqrt)
+        super().__init__(defaults, uses_grad=use_grad, inner=inner,)
+
+    def reset_for_online(self):
+        super().reset_for_online()
+        self.clear_state_keys('prev_g')
+        self.global_state['reset'] = True
+
+    @torch.no_grad
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
+        variant = settings[0]['variant']
+        theta_0 = 0 if variant == 1 else 1/3
+        theta = self.global_state.get('theta', theta_0)
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        p = TensorList(params)
+        g = grads if self._uses_grad else tensors
+        assert g is not None
+        g = TensorList(g)
+
+        prev_p, prev_g = unpack_states(states, tensors, 'prev_p', 'prev_g', cls=TensorList)
+
+        # online
+        if self.global_state.get('reset', False):
+            del self.global_state['reset']
+            prev_p.copy_(p)
+            prev_g.copy_(g)
+            return
+
+        if step == 0:
+            alpha_0 = settings[0]['alpha_0']
+            if alpha_0 is None: alpha_0 = epsilon_step_size(g)
+            self.global_state['alpha']  = alpha_0
+            prev_p.copy_(p)
+            prev_g.copy_(g)
+            return
+
+        sqrt = settings[0]['sqrt']
+        alpha = self.global_state.get('alpha', math.inf)
+        L = (g - prev_g).global_vector_norm() / (p - prev_p).global_vector_norm()
+        eps = torch.finfo(L.dtype).tiny * 2
+
+        if variant == 1:
+            a1 = math.sqrt(1 + theta)*alpha
+            val = math.sqrt(2) if sqrt else 2
+            if L > eps: a2 = 1 / (val*L)
+            else: a2 = math.inf
+
+        elif variant == 2:
+            a1 = math.sqrt(2/3 + theta)*alpha
+            a2 = alpha / math.sqrt(max(eps, 2 * alpha**2 * L**2 - 1))
+
+        else:
+            raise ValueError(variant)
+
+        alpha_new = min(a1, a2)
+        if alpha_new < 0: alpha_new = max(a1, a2)
+        if alpha_new > eps:
+            self.global_state['theta'] = alpha_new/alpha
+            self.global_state['alpha'] = alpha_new
+
+        prev_p.copy_(p)
+        prev_g.copy_(g)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        alpha = self.global_state.get('alpha', None)
+
+        if not _acceptable_alpha(alpha, tensors[0]):
+            # alpha isn't None on 1st step
+            self.state.clear()
+            self.global_state.clear()
+            alpha = epsilon_step_size(TensorList(tensors), settings[0]['alpha_0'])
+
+        torch._foreach_mul_(tensors, alpha)
+        return tensors
+
+    def get_H(self, var):
+        return _get_H(self, var)