Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add zoom line search. #255

Merged
merged 1 commit into from
Jun 27, 2022
Merged

Add zoom line search. #255

merged 1 commit into from
Jun 27, 2022

Conversation

mblondel
Copy link
Collaborator

@mblondel mblondel commented Jun 23, 2022

This PR is the first of a series of PRs to integrate zoom line search in JAXopt. It uses the code of @Joshuaalbert currently in JAX core. It supports only 1d arrays for parameters at the moment.

To be done in other PRs:

  1. write more correctness tests, e.g. tests for isolated functions such as quadmin or cubicmin
  2. simplify the code when possible
  3. add pytree support, either natively or via flattening
  4. benchmarks: backtracking vs. zoom, SciPy vs. JAXopt on various problems (both convex and nonconvex)
  5. use our iterative solver API instead of using a lax.while_loop (this will give access to individual steps of the line search)

CC: @shoyer @jakevdp @junpenglao

@mblondel
Copy link
Collaborator Author

This is now ready for review. I'll tackle the aforementioned improvements separately, as I think it's important to have a copy of the original code in the commit history. A benchmark on a small synthetic dataset (10,000 samples, 200 features) looks good:

image

Copy link
Collaborator

@fabianp fabianp left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks good to me, thanks!

@copybara-service copybara-service bot merged commit fdaf31d into google:main Jun 27, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants