[WIP] [ENH] add faster jvp computation for lasso type problems

[QB3]

Algorithmically implicit differentiation is composed of 2 steps:

• 1 compute the solution of the optimization problem
• 2 solve a linear system.

This linear system can be large, and expensive to solve.
When one differentiates the solution of sparse optimization problems (such as the Lasso), it is possible to reduce the size of the linear system to solve (http://proceedings.mlr.press/v119/bertrand20a/bertrand20a.pdf, https://arxiv.org/pdf/2105.01637.pdf).

The goal of this PR is to implement such an acceleration for sparse optimization problems.

To this aim, I implemented a sparse_root_vjp function (https://github.com/QB3/jaxopt/blob/c9b0daea3f90dec392fbcb73097eb88d39010aa2/jaxopt/_src/implicit_diff.py#L125) where I solve a smaller linear system (https://github.com/QB3/jaxopt/blob/c9b0daea3f90dec392fbcb73097eb88d39010aa2/jaxopt/_src/implicit_diff.py#L169).
The correctness of the implementation is tested here https://github.com/QB3/jaxopt/blob/c9b0daea3f90dec392fbcb73097eb88d39010aa2/tests/implicit_diff_test.py#L88.
Note that this implementation is very brutal, and not general at all, the goal is to see if we observe any speedups.

To check the speed of the implementation I created a sparse_vjp.py file in the benchmarks directory (https://github.com/QB3/jaxopt/blob/add_sparse_vjp/benchmarks/sparse_vjp.py).
I differentiate the solution of a Lasso with (n_samples=10, n_features=1000), lam = lam_max / 2.

I have the following results:
Time taken to solve the Lasso optimization problem 0.008
Time taken to compute the Jacobian 2.168
Time taken to compute the Jacobian with the sparse implementation 2.168

This benchmark tells us 2 things:

• there is room for improvement when computing the Jacobian (it takes 1000 more times than solving the optimization problem)
• the sparse implementation does not provide speedups.

I do not understand why we do not observe any speedups, does someone have a lead?

[google-cla]

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it! and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[google-cla]

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it! and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[google-cla]

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it! and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[google-cla]

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it! and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[google-cla]

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it! and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[mblondel]

@QB3 updated the pull request description with the current state of this pull request. Currently, he doesn't observe any speed ups despite restricting the linear system to the support.

@shoyer @froystig @fabianp If you have any idea what could make things faster, let us know.

[shoyer]

@QB3 updated the pull request description with the current state of this pull request. Currently, he doesn't observe any speed ups despite restricting the linear system to the support.

I believe the number of iterations required for CG to converge typically depends on the condition number of the linear operator, rather than size of the system. It's possible that restricting the support of the system does not change that. More generally, it seems like it would be a good idea to collect some metrics (e.g., via host_callback) about the rate of convergence.

I would test with direct solvers jax.scipy.linalg.solve (on small systems), which likely depend more directly on the size of the system. You also might try calculating the condition number of the linear operators, e.g., via the approximate eigenvalue calcualtions from scipy.sparse.linalg.

[froystig]

Might it also be useful to profile (tutorial, docs) for any lower-level surprises?

[mblondel]

@shoyer In the reference NumPy-based implementation of our paper, that also uses CG, @QB3 apparently observed speed ups by restricting the support (please confirm @QB3).

[QB3]

Thanks a lot for your answers, I tried on a larger example and I observed some nice speedups!
(n=10, p=10_000)

Time taken to solve the Lasso optimization problem 0.015
Time taken to compute the Jacobian 24.673
Time taken to compute the Jacobian with the sparse implementation 1.435