Jax for ML estimation - Slower than scipy?

[saketkc]

Thanks for creating this library! I was very excited after hearing @jakevdp's wonderful talk.

I tried using JAX for MLE estimation of a Negative Binomial random variable. My baseline was scipy.minimize with BFGS method. I was hoping grad would lead to a speedup, but as you would see in the plot that follows, it seems not to be the case. LBFGS appears to be the fastest at the cost of being more biased, but that is not important. Surprisingly, BFGS with autograd appears to be slower than vanilla BFGS. It is also the case with jit enabled.

It's likely I missed something. I could use vmap for the multiple units to probably see a gain, but I am trying to wrap my around as to why does autograd lead too slower convergence?

Notebook that reproduces the above figure is on Colab

[shoyer]

I just wrote a new draft "FAQ" entry on "Benchmarking JAX code": https://github.com/google/jax/pull/5879/files

Please a take and let me know if it helps answer your question.

In this particular case, I think part of the issue is that your example using JIT only compiles the gradient function, not the entire thing. Ideally you would JIT the entire function that you're benchmarking, including the call to jax.scipy.optimize.minimize:

@jax.jit
def fit_nbinom_bfgs_autograd_jit(y, mu):
  ...

[saketkc]

Thanks @shoyer! The FAQ is indeed very helpful. I was able to write down a faster estimator using JIT following your notes. Thanks a lot!