Will constrained optimization be supported?

[zhouxian]

Hi, thank you for the amazing library!

I was wondering if minimizing with user-specified bounds, or algorithms like projected gradient descent are supported?
If not, what would be the best practice you suggest if we are trying to solve things like argmin(norm(Ax+b)) s.t. x >=0?

Thanks in advance!

[patrick-kidger]

So right now we do support hypercube constraints in our root-finding methods (e.g. Newton or Bisection). And moreover it is sort-of-possible to cast a minimisation problem as a root-finding problem, by seeking a root of the gradient of the function. (You have to be careful as this may also converge to any other external point.) So I think the particular example you've given could be tackled today, although I'm not sure how efficient it would be as compared to a dedicated solver.

Other than that we don't currently support any kind of constraints in our solves. We may do at some point, but as always developer time is the limiting factor!

[mjo22]

Hello! Was also curious on this front. In particular, I am wondering if there are plans to support box constraints in BFGS, or if this can already be done in the current framework. jaxopt seems to have an implementation of this, but I’m not sure how well maintained it is.

[patrick-kidger]

For something simple like box constraints, then it should be fairly easy to add that in. We'd be happy to take a PR adding that. (Note that box constraints are already implemented for the root-finders, so you could follow a similar pattern to there.)

[f0uriest]

If box constraints are supported it shouldn't be too much work to add an augmented Lagrangian type method for general constraints. Could maybe be like a "meta solver" where you can wrap another solver to manage the unconstrained sub problems? I can hopefully take a look at it once I finish up a few other projects.

[billtubbs]

Constrained nonlinear optimization would be great.

I'm not an optimization expert, so this may be a lot to ask, but what would it take to get a solver like IPOPT integrated with Jax? IPOPT is a tried and tested solver for large scale nonlinear optimization (with nonlinear constraints) used widely in engineering fields like chemical engineering, control, and robotics (see for example CasADi, GEKKO, Pyomo). Would be great to see more people in these fields trying out Jax.

From my limited understanding, solving large scale problems requires the solver to take advantage of the sparsity of the constraint Jacobian and Hessian matrices. Is this a fundamental problem in Jax?

[johannahaffner]

Hi @billtubbs,

since box constraints (mentioned above) are just a special case of inequality constraints, I have been looking into ways to make some of the "ingredients" used for their implementation in various methods available. It is still early days though!

Thank you for the solver suggestions, I will look into them :)

[f0uriest]

If I remember correctly, IPOPT uses a filter to decide on accepting steps, which might be difficult to implement in JAX with compile time memory bounds? It can also switch between different "modes" dynamically which might make vmap super inefficient.

Something like LANCELOT/GALAHAD or SLSQP would probably be easier.

As for sparsity, that is usually what lets you solve large problems efficiently, though you don't necessarily need actual sparse data structures. It should be easy to just use some of the iterative solvers from lineax on the KKT system, though you may need a good preconditioner. @patrick-kidger is it possible to pass linear solver options like preconditoners through optimistix to lineax?

[patrick-kidger]

I think frequently "no" but happy to consider a PR changing that fact :)