python-somax 0.0.1__py3-none-any.whl

somax/gn/sgn.py

@@ -0,0 +1,237 @@
+"""
+Stochastic Gauss-Newton (SGN):
+- approximate solution to the linear system via the CG method
+- adaptive regularization (lambda)
+Paper: https://arxiv.org/abs/2006.02409
+"""
+
+from typing import Any
+from typing import Callable
+from typing import NamedTuple, Tuple
+from typing import Optional
+import dataclasses
+from functools import partial
+
+import jax
+import jax.lax as lax
+import jax.numpy as jnp
+from jax.flatten_util import ravel_pytree
+from jax.scipy.sparse.linalg import cg
+from optax import sigmoid_binary_cross_entropy
+from jaxopt.tree_util import tree_add_scalar_mul, tree_scalar_mul
+from jaxopt._src import base
+from jaxopt._src import loop
+
+
+class SGNState(NamedTuple):
+    """Named tuple containing state information."""
+    iter_num: int
+    # error: float
+    # value: float
+    stepsize: float
+    regularizer: float
+    # direction_inf_norm: float
+
+
+@dataclasses.dataclass(eq=False)
+class SGN(base.StochasticSolver):
+    predict_fun: Callable
+
+    loss_fun: Optional[Callable] = None
+    loss_type: str = 'mse'
+    maxcg: int = 10
+
+    learning_rate: float = 1.0  # default value recommended by Gargiani et al.
+
+    batch_size: Optional[int] = None
+
+    n_classes: Optional[int] = None
+
+    # Adaptive Regularization parameters
+    adaptive_lambda: bool = False
+    regularizer: float = 1e-3  # default value recommended by Gargiani et al.
+    lambda_decrease_factor: float = 0.99  # default value recommended by Kiros
+    lambda_increase_factor: float = 1.01  # default value recommended by Kiros
+
+    pre_update: Optional[Callable] = None
+
+    verbose: int = 0
+
+    jit: bool = True
+    unroll: base.AutoOrBoolean = "auto"
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        self.reference_signature = self.loss_fun
+
+        # Regression (MSE)
+        if self.loss_type == 'mse':
+            if self.loss_fun is None:  # mostly the case for supervised learning
+                self.loss_fun = self.mse
+
+            self.loss_fun_from_preds = self.mse_from_predictions
+
+        # Classification (Cross-Entropy)
+        elif self.loss_type == 'ce' or self.loss_type == 'xe':
+            if self.n_classes == 1:  # binary classification
+                if self.loss_fun is None:  # mostly the case for supervised learning
+                    self.loss_fun = self.ce_binary
+
+                self.loss_fun_from_preds = self.ce_binary_from_logits
+
+            else:
+                if self.loss_fun is None:  # mostly the case for supervised learning
+                    self.loss_fun = self.ce
+
+                self.loss_fun_from_preds = self.ce_from_logits
+
+        self.grad_fun = jax.grad(self.loss_fun)
+
+    def update(
+            self,
+            params: Any,
+            state: SGNState,
+            *args,
+            **kwargs,
+    ) -> base.OptStep:
+        """Performs one iteration of the solver.
+
+        Args:
+          params: pytree containing the parameters.
+          state: named tuple containing the solver state.
+          *args: additional positional arguments to be passed to ``fun``.
+          **kwargs: additional keyword arguments to be passed to ``fun``.
+        Returns:
+          (params, state)
+        """
+
+        # ---------- STEP 1: calculate direction with DG ---------- #
+        # TODO analyze *args and **kwargs
+        # split (x,y) pair into (x,) and (y,)
+        if 'targets' in kwargs:
+            targets = kwargs['targets']
+            nn_args = args
+        else:
+            targets = args[-1]
+            nn_args = args[:-1]
+
+        direction_tree, grad_loss_tree = self.calculate_direction(params, state, targets, *nn_args)
+
+        # # ---------- STEP 2: update (next step) lambda ---------- #
+        # TODO
+
+        next_params = tree_add_scalar_mul(params, state.stepsize, direction_tree)
+
+        # construct the next state
+        next_state = SGNState(
+            iter_num=state.iter_num + 1,  # Next Iteration
+            stepsize=state.stepsize,  # Current alpha
+            regularizer=state.regularizer,  # Next lambda
+        )
+
+        return base.OptStep(params=next_params, state=next_state)
+
+    def init_state(self,
+                   init_params: Any,
+                   *args,
+                   **kwargs) -> SGNState:
+        return SGNState(
+            iter_num=0,
+            stepsize=self.learning_rate,
+            regularizer=self.regularizer,
+        )
+
+    def optimality_fun(self, params, *args, **kwargs):
+        """Optimality function mapping compatible with ``@custom_root``."""
+        # return self._grad_fun(params, *args, **kwargs)[0]
+        raise NotImplementedError
+
+    def gnhvp(self, params, vec, targets, *args):
+        def inner_predict_fun(params):
+            return self.predict_fun(params, *args)
+
+        def hvp(network_outputs, vec, targets):
+            # prep grad function to accept only "network_outputs"
+            def inner_grad_fun(network_outputs):
+                return jax.grad(self.loss_fun_from_preds)(network_outputs, targets)
+
+            return jax.jvp(inner_grad_fun, (network_outputs,), (vec,))[1]
+
+        # H_GN ~ J^T Q J, s.t. hvp is J^T Q J v
+        network_outputs, Jv = jax.jvp(inner_predict_fun, (params,), (vec,))
+
+        if self.loss_type == 'mse':
+            _, JTJv_fun = jax.vjp(inner_predict_fun, params)
+            JTJv = JTJv_fun(Jv)[0]
+            return tree_scalar_mul(1 / self.batch_size, JTJv)
+        else:
+            QJv = hvp(network_outputs, Jv, targets)
+            _, JTQJv_fun = jax.vjp(inner_predict_fun, params)
+            JTQJv = JTQJv_fun(QJv)[0]
+            return JTQJv
+
+    def calculate_direction(self, params, state, targets, *args):
+        def mvp(vec):
+            # H_{GN}v
+            # hv = jax.jvp(jax.grad(stand_alone_loss_fn), (params,), (vec,))[1]
+            hv = self.gnhvp(params, vec, targets, *args)
+            # add regularization, works since (H + lambda*I) v = Hv + lambda*v
+            return tree_add_scalar_mul(hv, state.regularizer, vec)
+
+        # --------- Start Here --------- #
+        # calculate grad
+        grad_tree = self.grad_fun(params, *args, targets)
+
+        # CG iterations
+        # TODO initial guess and preconditioner
+        direction, _ = cg(
+            A=mvp,
+            b=tree_scalar_mul(-1, grad_tree),
+            maxiter=self.maxcg,
+            # x0=None,  # initial guess
+            # M=None,  # preconditioner
+        )
+
+        return direction, grad_tree
+
+    def mse(self, params, x, y):
+        return self.mse_from_predictions(self.predict_fun(params, x), y)
+
+    @staticmethod
+    def mse_from_predictions(preds, y):
+        return 0.5 * jnp.mean(jnp.square(y - preds))
+
+    def ce(self, params, x, y):
+        logits = self.predict_fun(params, x)
+        return self.ce_from_logits(logits, y)
+
+    @staticmethod
+    def ce_from_logits(logits, y):
+        # b x C
+        # jax.nn.log_softmax combines exp() and log() in a numerically stable way.
+        log_probs = jax.nn.log_softmax(logits)
+
+        # b x 1
+        # if y is one-hot encoded, this operation picks the log probability of the correct class
+        residuals = jnp.sum(y * log_probs, axis=-1)
+
+        # 1,
+        # average over the batch
+        return -jnp.mean(residuals)
+
+    def ce_binary(self, params, x, y):
+        logits = self.predict_fun(params, x)
+        return self.ce_binary_from_logits(logits, y)
+
+    def ce_binary_from_logits(self, logits, y):
+        # b x 1
+        loss = sigmoid_binary_cross_entropy(logits.ravel(), y)
+
+        # 1,
+        # average over the batch
+        return jnp.mean(loss)
+
+    def __hash__(self):
+        # We assume that the attribute values completely determine the solver.
+        return hash(self.attribute_values())

somax/hf/newton_cg.py

@@ -0,0 +1,339 @@
+"""
+Newton-CG Solver:
+- approximate solution to the linear system via the CG method
+- adaptive regularization (lambda)
+- line search
+"""
+
+from typing import Any
+from typing import Callable
+from typing import NamedTuple, Tuple
+from typing import Optional
+import dataclasses
+from functools import partial
+
+import jax
+import jax.lax as lax
+import jax.numpy as jnp
+from jax.flatten_util import ravel_pytree
+from jax.scipy.sparse.linalg import cg
+from jaxopt.tree_util import tree_add_scalar_mul, tree_scalar_mul
+from jaxopt._src import base
+from jaxopt._src import loop
+
+
+def wolfe_cond_violated(stepsize, coef, f_cur, f_next, direct_deriv):
+    eps = jnp.finfo(f_next.dtype).eps
+    loss_decrease = f_cur - f_next + eps
+    prescribed_decrease = -stepsize * coef * direct_deriv
+    return prescribed_decrease > loss_decrease
+
+
+def curvature_cond_violated(stepsize, coef, f_cur, f_next, direct_deriv):
+    loss_decrease = f_cur - f_next
+    prescribed_decrease = -stepsize * (1. - coef) * direct_deriv
+    return loss_decrease > prescribed_decrease
+
+
+def armijo_line_search(loss_fun, unroll, jit,
+                       goldstein, maxls,
+                       params, f_cur, stepsize,
+                       direction, direct_deriv,
+                       coef, decrease_factor, increase_factor, max_stepsize,
+                       args, targets):
+    # given direction calculate next params
+    next_params = tree_add_scalar_mul(params, stepsize, direction)
+
+    # calculate loss at next params
+    f_next = loss_fun(next_params, *args, targets)
+
+    # grad_sqnorm = tree_l2_norm(grad, squared=True)
+
+    def update_stepsize(t):
+        """Multiply stepsize per factor, return new params and new value."""
+        stepsize, factor = t
+        stepsize = stepsize * factor
+        stepsize = jnp.minimum(stepsize, max_stepsize)
+
+        next_params = tree_add_scalar_mul(params, stepsize, direction)
+
+        f_next = loss_fun(next_params, *args, targets)
+
+        return stepsize, next_params, f_next
+
+    def body_fun(t):
+        stepsize, next_params, f_next, _ = t
+
+        violated = wolfe_cond_violated(stepsize, coef, f_cur, f_next, direct_deriv)
+
+        stepsize, next_params, f_next = lax.cond(
+            violated,
+            update_stepsize,
+            lambda _: (stepsize, next_params, f_next),
+            operand=(stepsize, decrease_factor),
+        )
+
+        if goldstein:
+            goldstein_violated = curvature_cond_violated(
+                stepsize, coef, f_cur, f_next, direct_deriv)
+
+            stepsize, next_params, f_next = lax.cond(
+                goldstein_violated, update_stepsize,
+                lambda _: (stepsize, next_params, f_next),
+                operand=(stepsize, increase_factor),
+            )
+
+            violated = jnp.logical_or(violated, goldstein_violated)
+
+        return stepsize, next_params, f_next, violated
+
+    init_val = stepsize, next_params, f_next, jnp.array(True)
+
+    ret = loop.while_loop(cond_fun=lambda t: t[-1],  # check boolean violated
+                          body_fun=body_fun,
+                          init_val=init_val, maxiter=maxls,
+                          unroll=unroll, jit=jit)
+
+    return ret[:-1]  # remove boolean
+
+
+class NewtonCGState(NamedTuple):
+    """Named tuple containing state information."""
+    iter_num: int
+    stepsize: float
+    regularizer: float
+    cg_guess: Optional[Any]
+
+
+@dataclasses.dataclass(eq=False)
+class NewtonCG(base.StochasticSolver):
+    # Jacobian of the residual function
+    loss_fun: Callable
+    maxcg: int = 3
+
+    # Either fixed alpha if line_search=False or max_alpha if line_search=True
+    learning_rate: float = 1.0
+    decay_coef = 0.1
+
+    batch_size: Optional[int] = None
+
+    n_classes: Optional[int] = None
+
+    # Adaptive Regularization parameters
+    adaptive_lambda: bool = True
+    regularizer: float = 1.0
+    lambda_decrease_factor: float = 0.99  # default value recommended by Kiros
+    lambda_increase_factor: float = 1.01  # default value recommended by Kiros
+
+    # Line Search parameters
+    line_search: bool = True
+
+    aggressiveness: float = 0.9  # default value recommended by Vaswani et al.
+    decrease_factor: float = 0.8  # default value recommended by Vaswani et al.
+    increase_factor: float = 1.5  # default value recommended by Vaswani et al.
+    reset_option: str = 'increase'  # ['increase', 'goldstein', 'conservative']
+
+    max_stepsize: float = 1.0
+    maxls: int = 15
+
+    pre_update: Optional[Callable] = None
+
+    verbose: int = 0
+
+    jit: bool = True
+    unroll: base.AutoOrBoolean = "auto"
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        self.reference_signature = self.loss_fun
+        self.grad_fun = jax.grad(self.loss_fun)
+
+        # set up line search
+        if self.line_search:
+            # !! learning rate is the maximum step size in case of line search
+            self.max_stepsize = self.learning_rate
+
+            options = ['increase', 'goldstein', 'conservative']
+
+            if self.reset_option not in options:
+                raise ValueError(f"'reset_option' should be one of {options}")
+            if self.aggressiveness <= 0. or self.aggressiveness >= 1.:
+                raise ValueError(f"'aggressiveness' must belong to open interval (0,1)")
+
+            self._coef = 1 - self.aggressiveness
+
+            unroll = self._get_unroll_option()
+
+            armijo_with_fun = partial(armijo_line_search, self.loss_fun, unroll, self.jit)
+            if self.jit:
+                jitted_armijo = jax.jit(armijo_with_fun, static_argnums=(0, 1))
+                self._armijo_line_search = jitted_armijo
+            else:
+                self._armijo_line_search = armijo_with_fun
+
+    def update(
+            self,
+            params: Any,
+            state: NewtonCGState,
+            *args,
+            **kwargs,
+    ) -> base.OptStep:
+        """Performs one iteration of the solver.
+
+        Args:
+          params: pytree containing the parameters.
+          state: named tuple containing the solver state.
+          *args: additional positional arguments to be passed to ``fun``.
+          **kwargs: additional keyword arguments to be passed to ``fun``.
+        Returns:
+          (params, state)
+        """
+
+        # ---------- STEP 1: calculate direction with DG ---------- #
+        # TODO analyze *args and **kwargs
+        # split (x,y) pair into (x,) and (y,)
+        if 'targets' in kwargs:
+            targets = kwargs['targets']
+            nn_args = args
+        else:
+            targets = args[-1]
+            nn_args = args[:-1]
+
+        direction_tree, grad_loss_tree = self.calculate_direction(params, state, targets, *nn_args)
+
+        # ---------- STEP 2: line search for alpha ---------- #
+        f_cur = None
+        f_next = None
+        direct_deriv = None
+        if not self.line_search:
+            # constant learning rate
+            stepsize = state.stepsize
+            next_params = tree_add_scalar_mul(params, state.stepsize, direction_tree)
+        else:
+            stepsize = self.reset_stepsize(state.stepsize)
+
+            goldstein = self.reset_option == 'goldstein'
+
+            f_cur = self.loss_fun(params, *nn_args, targets)
+
+            # the directional derivative used for Armijo's line search
+            direction, _ = ravel_pytree(direction_tree)
+            grad_loss, _ = ravel_pytree(grad_loss_tree)
+            direct_deriv = grad_loss.T @ direction
+
+            stepsize, next_params, f_next = self._armijo_line_search(
+                goldstein, self.maxls, params, f_cur, stepsize,
+                direction_tree, direct_deriv, self._coef,
+                self.decrease_factor, self.increase_factor,
+                self.max_stepsize, nn_args, targets,
+            )
+
+        # ---------- STEP 3: update (next step) lambda ---------- #
+        if not self.adaptive_lambda:
+            # constant regularization
+            regularizer_next = state.regularizer
+        else:
+            # numerator: in a good scenario, should be large and negative
+            if f_cur is None:
+                f_cur = self.loss_fun(params, *nn_args, targets)
+            if f_next is None:
+                f_next = self.loss_fun(next_params, *nn_args, targets)
+
+            num = f_next - f_cur
+
+            # denominator
+            Hv_tree = self.hvp(params, direction_tree, targets, *nn_args)
+
+            # flattening stage
+            Hv, _ = ravel_pytree(Hv_tree)
+
+            if direct_deriv is None:
+                direction, _ = ravel_pytree(direction_tree)
+                grad_loss, _ = ravel_pytree(grad_loss_tree)
+                direct_deriv = grad_loss.T @ direction
+
+            denom = 0.5 * jnp.vdot(direction, Hv) + direct_deriv
+
+            # negative denominator means that the direction is a descent direction
+            rho = num / denom
+
+            regularizer_next = lax.cond(
+                rho < 0.25,
+                lambda _: self.lambda_increase_factor * state.regularizer,
+                lambda _: lax.cond(
+                    rho > 0.75,
+                    lambda _: self.lambda_decrease_factor * state.regularizer,
+                    lambda _: state.regularizer,
+                    None,
+                ),
+                None,
+            )
+
+        # construct the next state
+        next_state = NewtonCGState(
+            iter_num=state.iter_num + 1,  # Next Iteration
+            stepsize=stepsize,  # Current alpha
+            regularizer=regularizer_next,  # Next lambda
+            cg_guess=tree_scalar_mul(self.decay_coef, direction_tree),  # Next CG guess
+        )
+
+        return base.OptStep(params=next_params, state=next_state)
+
+    def init_state(self,
+                   init_params: Any,
+                   *args,
+                   **kwargs) -> NewtonCGState:
+        return NewtonCGState(
+            iter_num=0,
+            stepsize=self.learning_rate,
+            regularizer=self.regularizer,
+            cg_guess=None,
+        )
+
+    def optimality_fun(self, params, *args, **kwargs):
+        """Optimality function mapping compatible with ``@custom_root``."""
+        # return self._grad_fun(params, *args, **kwargs)[0]
+        raise NotImplementedError
+
+    def reset_stepsize(self, stepsize):
+        """Return new step size for current step, according to reset_option."""
+        if self.reset_option == 'goldstein':
+            return stepsize
+        if self.reset_option == 'conservative':
+            return stepsize
+        stepsize = stepsize * self.increase_factor
+        return jnp.minimum(stepsize, self.max_stepsize)
+
+    def hvp(self, params, vec, targets, *args):
+        # prep grad function to accept only "params"
+        def inner_grad_fun(params):
+            return self.grad_fun(params, *args, targets)
+
+        return jax.jvp(inner_grad_fun, (params,), (vec,))[1]
+
+    def calculate_direction(self, params, state, targets, *args):
+        def mvp(vec):
+            # Hv
+            hv = self.hvp(params, vec, targets, *args)
+            # add regularization, works since (H + lambda*I) v = Hv + lambda*v
+            return tree_add_scalar_mul(hv, state.regularizer, vec)
+
+        # --------- Start Here --------- #
+        # calculate grad
+        grad_tree = self.grad_fun(params, *args, targets)
+
+        # CG iterations
+        direction, _ = cg(
+            A=mvp,
+            b=tree_scalar_mul(-1, grad_tree),
+            maxiter=self.maxcg,
+            x0=state.cg_guess,  # initial guess
+            # M=None,  # preconditioner (see wiesler2013: preconditioner is not important)
+        )
+
+        return direction, grad_tree
+
+    def __hash__(self):
+        # We assume that the attribute values completely determine the solver.
+        return hash(self.attribute_values())