JAXopt with original OSQP library

[mhr]

I'm finding that for my particular usecase, the original OSQP library is much faster on the CPU since it uses sparse arrays, while your implementation of the OSQP algorithm in JAX is slower because it uses dense arrays.

I'm struggling to use custom_root, though. How can I implicitly differentiate through the OSQP library's QP solver with JAXopt?

[Algue-Rythme]

Hi,

We plan on adding more uttorials on how to use custom_root in the future. In the mean time you can look at our implementation:

jaxopt/jaxopt/_src/osqp.py

Line 50 in 238fcb4

def _make_osqp_optimality_fun(matvec_Q, matvec_A):

As you can see, you need primal/dual variables and the matrices Q and A (or associated linear operator). Now, the issue is that we don't support sparse arrays yet, so even if you solve the quadratic program with original's OSQP on sparse arrays, the implicit differentiation will require dense arrays for the computation.

Nonetheless, I think the most straighforward workaround would be the following: solve you problem on CPU with original's OSQP implementation and sparse arrays. Retrive primal and dual variables, and finally take advantage of the fact that the problem is convex.

Indeed, if you have primal/dual variables you can populate the KKTSolution tuple yourself (the one ordinarily returned by init_params) and pass it to init_state. Now, since the problem is convex, if you call the run method on this high quality initialization, OSQP should be able to detect that it is at convergence and will return within few iterations (maybe a single one ?) and the support for implicit differentiation will be automatic. So the total cost would be the compilation time (when applicable), the computation of termination criterion (not too expensive) and finally the solving of the linear system with dense matrices in implicit differentiation.

Note that if your sparse matrices have some appealing structure you want to use a matvec matvec_A(x) to compute Ax efficiently, without requiring to store the whole A in memory. This covers a lot of different sparse operators (but not those for which the sparsity is arbitrary and requires CSR format).

Note that a year ago I started to create a wrapper for sparse operator on my own repo: https://github.com/Algue-Rythme/jaxopt/blob/38ea8faa7dd4c02c65bcb538673cb0f0db3f7128/jaxopt/_src/linear_operator.py#L59

Maybe you can take a look at from jax.experimental import sparse to define your own matvecs based on sparse matrices.

[mhr]

Thank you!

I went ahead and made a notebook to demonstrate how it seems that custom_root works and also your suggestion of using vanilla OSQP to warm-start JAXopt's OSQP implementation.

Your warm-start suggestion seems to be the winner by performance. Would you mind taking a look to let me know if I did everything correctly?

https://colab.research.google.com/drive/1p-ICCI3MafF7lHZa34LlJtukt4o_ufLn#scrollTo=DHCNLs3mu3_l

[Algue-Rythme]

Hi,

You almost got it right!

I would replace qp.l2_optimality_error with qp.optimality_fun in your function optimality for Vanilla OSQP, and use @im_diff.custom_root(optimality), not @im_diff.custom_root(jax.grad(optimality)) as here:

jaxopt/jaxopt/_src/base.py

Line 245 in 238fcb4

decorator = idf.custom_root(

Also, no need to do jax.jit(qp.run), setting jit=True (or implicit_diff=True) in constructor is enough.

Finally, note that you shouldn't be able to differentiate wrt k: the initialization KKT solution has no influence the solution for convex problems (you end up in the same place regarless initialization). Hence mathematically the derivative is zero, and numerically proper implementation of Implicit Diff returns a very obscure error (something that we need to document). Hence you should use argnums=(1, 2, 3, 4, 5) instead of argnums=(0, 1, 2, 3, 4, 5) in value_and_grad.