search

google / jaxopt

Hardware accelerated, batchable and differentiable optimizers in JAX.
https://jaxopt.github.io
Apache License 2.0
903 stars 62 forks source link

Linesearch in Broyden #559

Open zaccharieramzi opened 7 months ago

zaccharieramzi commented 7 months ago

This PR follows an email thread with @vroulet (sorry for taking this long to address this). Basically the idea is that at the moment the Broyden tests are passing but with failed line searches. Ideally, they should pass without the linesearch failing.

This PR makes sure that at least one test is testing that condition, in addition to using the more modern linesearch API for Broyden.

At the moment, the test is failing because we cannot modify the condition parameter of the linesearch. Indeed, as pointed out by @vroulet , scipy uses the armijo line search for Broyden, and if we hack our way into using it, the test passes. @fabianp I noticed you deprecated condition in this commit, what was the reason for this?

This needs to be merged after #529

fabianp commented 7 months ago

@fabianp https://github.com/fabianp I noticed you deprecated condition in this commit https://github.com/google/jaxopt/commit/7238f760c717e1db534fae3890150b0bfed82a7d, what was the reason for this?

The idea was that the options passed with "condition" could now be passed through the stepsize , since this can be a callable now

Message ID: @.***>

zaccharieramzi commented 7 months ago

I am not sure I understand. If stepsize is a callable, then use_linesearch is False, i.e. it's the case where we use a step size scheduler rather than a line search. So when stepsize is a callable, nothing is passed to the line search because it's not used.

fabianp commented 7 months ago

the idea is that having a callable for a step-size could be used not just for simple pre-defined schedules, but also for line-searches like (say) a custom Armijo line-search (correct me @vroulet if I'm wrong)

zaccharieramzi commented 7 months ago

But this would mean having the step size callable take as input not just the iteration number but also all the things that the line search takes as input and change its output to return all the metadata returned by the line search, no?

fabianp commented 7 months ago

ok, scratch that, sorry for the confusion. I think that the custom line-search was meant to be passed as a callable linesearch function (eventually). But clearly we didn't properly finish the API, and "condition" should not have been deprecated so soon

zaccharieramzi commented 7 months ago

ok, I might reinstate it here as it's important for Broyden to be able to change the condition of the linesearch to armijo.

vroulet commented 7 months ago

the idea is that having a callable for a step-size could be used not just for simple pre-defined schedules, but also for line-searches like (say) a custom Armijo line-search (correct me @vroulet if I'm wrong)

The idea was to avoid making a list of potential linesearches with string options and let the user define his/her own linesearch, which would always be of type IterativeLineSearch. That would have let us avoid adding the arguments of the linesearch inside the solver, but yes it was not finished...