Is it possible to "continue" optimization from a full OptStep?

[DanPuzzuoli]

Hi!

I'm playing around with trying to build some convenience functionality for doing a bunch of optimizations in JAXopt with different initial seeds. The landscape of the optimization problems I'm dealing is pretty complicated, so the number of iterations that an optimizer needs to take for different seeds can be wildly different. What I'm trying to build is a function that will run n optimizations, doing k simultaneously by vectorizing the optimizer with vmap, but that will check in every m optimizer iterations to see if any of the seeds have finished. What I'm currently trying is setting maxiter=m when constructing the optimizer, then doing a loop in which:

• the vectorized optimizer is called
• the results are checked to see if any seeds took less than m iterations
• Any init_params taking less than m iterations are swapped out for new ones
• loop back until all are done

My first attempt at this is working, but one issue I'm running into is that the optimizer state is totally reset in each iteration of the above loop, and hence any history in the corresponding OptStep for seeds that need to continue running is lost between loops. This ends up resulting in a single optimization ultimately taking more iterations than is necessary as things like step size or LBFGS history are lost. Is there any way to do something equivalent to passing a full OptStep object to run, to continue an optimization for where it last left off?

[Algue-Rythme]

Unfortunately, with run method only, the answer is no. It is not meant to take the state in input. However by calling repeatidly the method update yourself you can pass a custom state. However you will have to monitore yourself the termination condition (by default state.iter_num > solver.maxiter or state.error < solver.tol. And you will have to choose between Python loop (not parallelizable), a jax.lax.scan or a jax.lax.while_loop.

For your own custom implementation it is worth reading the way we did it in jaxopt