BFGS with constraints

[jecampagne]

Hi,

I do not know if it is possible to perform with the current JAX lib something corresponding to the scipy.optimize.fmin_l_bfgs_b optimisation, ie. minimization with constraints on parameters, as well as computation of the Hessian at the minimum point found?

Thanks.

[shoyer]

JAXopt has support for constrained optimization: https://jaxopt.github.io/stable/constrained.html

[Algue-Rythme]

Oh yes, I should have precised: all Jaxopt functions and optimizers are meant to be applied on single example, not whole batch.
Thanks to jax.vmap this is not an issue since we can wrap .run() with it and optimize over batchs.
This code works:

from jaxopt.projection import projection_l2_ball
import jax
import matplotlib.pyplot as plt

def myproj(x, hyper_params):
    bounds, radius = hyper_params
    x = projection_l2_ball(x - jnp.array([1,0]), radius) + jnp.array([1,0])
    return x

bounds= (jnp.array([-10., -10.]), jnp.array([10., 10.]))
radius = 1.5
hyper_params = bounds, radius

rng_key = jax.random.PRNGKey(42)
x = jax.random.uniform(rng_key, (1000,2), minval=-5,maxval=5)
p_x=jax.vmap(myproj, in_axes=(0,None))(x,hyper_params)

fig, ax = plt.subplots()
ax.scatter(x[:,0],x[:,1])
ax.scatter(p_x[:,0],p_x[:,1])

If you are used to Numpy this might be surprising at first glance. However it some sense it is closer from the "mathematical approach": you think and write your code for a single example, and let vmap handle the rest.

[jecampagne]

Great!

here it is if radius = 1.5**2

Now I was in trouble with the projection_box, but I found :)

def myproj(x, hyper_params):
    bounds, radius = hyper_params
    x = projection_box(x+ jnp.array([1,1]), bounds) - jnp.array([1,1])
    return x

bounds= (jnp.array([-1, -1]), jnp.array([1., 1.]))

I get

which is ok !

And here is a nice combinaison

Bye the way, I'm always in trouble with in_axes of jax.vmap, do you have an hint to get right?

Is there a way to define a poly line shape (in 2D) of the projection set?

Anyway Thanks for your clear explanations. I hope that the other kinds of constraints (ie. those available by scipy) will be implemented.

[Algue-Rythme]

Nice figures @jecampagne !
I guess you could define a poly line shape manually with your own tools, but not with tools in Jaxopt currently.

[Algue-Rythme]

For vmap I pass a tuple of length 2 to in_axes because the function myproj takes two arguments. I want to batch wrt the first argument, but not wrt the second, so I use (0, None): as a result I project 1000 examples onto the same ball.

If I wanted a different ball with different radius for each x, I would use in_axes=(0,0) and I could pass an array of 1000 different radius (one for each x). Or I could use in_axes=(None,0) to have a single example projected onto 1000 different balls.

Or I could decide to create x as x = jax.random.uniform(rng_key, (2,1000), minval=-5,maxval=5) instead (notice that rows and columns are now permuted) so I would have used in_axes=(1,None) because the batch size woul be in axis 1, not 0.

[jecampagne]

Good ! Thanks a lot @Algue-Rythme. We can close the thread. I do not know if there is a jaxopt tag/forum for discussing further if needed.

[jecampagne]

Hi @Algue-Rythme
You mention jaxopt.ScipyLeastSquaresbut I do not manage to get the API in the doc ???

[Algue-Rythme]

It's not referenced in the doc yet.

You will have to read the source and the docstring: https://github.com/google/jaxopt/blob/782b0a918695b4c9f75717a3bb4e371b4884d98a/jaxopt/_src/scipy_wrappers.py#L605 good luck!

[jecampagne]

@Algue-Rythme have you an example of jaxopt.projection.projection_polyhedron in 2D, thanks.

Note also that in the doc

the w is certainly related to the params (ie. general L2 norm constraint, moreover I would have expected this use-case in the "minimimzation with constraint solved by Lagrangian method... it is a matter of taste), in any case this is not clear enough.

[Algue-Rythme]

Thanks for the typo, this is indeed params and not w!

You will find examples of usage in tests/ folder: https://github.com/google/jaxopt/blob/03d5d3f6ebefbab3f1986a715789071580fb03c9/tests/projection_test.py#L279
(in 5D) in addition to the doc. Let me know if it works!

I would have expected this use-case in the "minimimzation with constraint solved by Lagrangian method... it is a matter of taste

L2 regularization is very common in ML (one of the major target of Jaxopt) and is rarely formulated as a hard constraint in this context (here Ridge regression).
Consequently there is no need to introduce the Lagrangian here, even though the link is direct. l2reg is better understood as an hyper-parameter of the algorithm that we can tune/optimize (and that ML practitioners will recognize) rather than a Lagrangian multiplier of some constrained optimization problem.