Box projection with zoom line search

[jaxengodfrey]

I have an objective function that I am trying to minimize with optax.lbfgs(), which I believe chains optax.scale_by_lbfgs() and optax.scale_by_zoom_linesearch() under the hood. I use optax.projections.projection_box() to restrict the function parameters to a specified range. In particular, the parameters need to be strictly positive otherwise the function evaluation fails. This is where my problem arises: it seems like the zoom linesearch will sometimes try to evaluate the function at negative values that fall outside of the box projection, causing it to throw an error and thus the optimizer fails. Is there a way to restrict the parameter values during the zoom linesearch? Or is there a linesearch better suited for constrained problems?

I am not able to create a minimum example of the problem, but I am using the example LBFGS solver from the docs page with one line added to do the box projection:

def run_opt(init_params, fun, opt, max_iter, tol):
  value_and_grad_fun = optax.value_and_grad_from_state(fun)

  def step(carry):
    params, state = carry
    params = optax.projections.projection_box(params, lower = jnp.array([0.01, 0.01]), upper = jnp.array([10., 10.]))
    value, grad = value_and_grad_fun(params, state=state)
    updates, state = opt.update(
        grad, state, params, value=value, grad=grad, value_fn=fun
    )
    params = optax.apply_updates(params, updates)
    #I've also tried putting the box projection call here
    return params, state

  def continuing_criterion(carry):
    _, state = carry
    iter_num = optax.tree.get(state, 'count')
    grad = optax.tree.get(state, 'grad')
    err = optax.tree.norm(grad)
    return (iter_num == 0) | ((iter_num < max_iter) & (err >= tol))

  init_carry = (init_params, opt.init(init_params))
  final_params, final_state = jax.lax.while_loop(
      continuing_criterion, step, init_carry
  )
  return final_params, final_state

[rdyro]

I don't believe it's something we support now, but we can look into it.

You could try a proxy objective with an exact penalty function or writing your own linesearch implementation.

[emilyfertig]

You could also try chaining optax.scale_by_lbfgs(), a transform that scales the gradients so that a step stays in bounds, and optax.scale_by_zoom_linesearch(max_learning_rate=1.). You could define the middle transform following this pattern, e.g. passing the lower/upper bounds to the scale_by... function you define and scaling the updates according to the distance between params and bounds in the update_fn. That might not have good convergence properties but it could still be worth trying.

Another option is to optimize in an unconstrained space and pass the unconstrained values through a sigmoid before evaluating the loss, like this:

def modified_loss(unconstrained_params, ...):
  params = jax.nn.sigmoid(unconstrained_params) * (ub - lb) + lb
  return loss(params, ...)

though that will warp the loss surface and the gradients will approach 0 near the bounds.