jax scipy optimizer (BFGS) slower compared to scipy

[mgbukov]

I'm comparing the BFGS solver in scipy and jax.scipy and I am seeing a huge slowdown in jax.

Is the jax implementation of the algorithm fundamentally different (see also the output of the function and jacobian evaluations res.nfev, res.njev below)? Or am I using it inappropriately?

Here's a minimal example:

import numpy as np 
import scipy as sp
import scipy.optimize as sp_opt

import jax.numpy as jnp 
import jax.scipy as jsp
import jax.scipy.optimize as jsp_opt


import time


H = [[ 1.+0.j,  0.+0.j,  0.+0.j,  0.+0.j],
	 [ 0.+0.j, -1.+0.j,  0.+0.j,  0.+0.j],
	 [ 0.+0.j,  0.+0.j, -1.+0.j,  0.+0.j],
	 [ 0.+0.j,  0.+0.j,  0.+0.j,  1.+0.j]]

psi = [0.45220585, 0.58929456, 0.49665892, 0.44896738]


H_jnp, psi_jnp = jnp.array(H, dtype=jnp.complex64), jnp.array(psi, dtype=jnp.complex64)
H_np, psi_np   =  np.array(H, dtype=np.complex64) ,  np.array(psi, dtype=np.complex64)


#@jit
def compute_energy_jnp(angle,  H,psi):
	U=jsp.linalg.expm(-1j*angle*H)
	psi_new = U @ psi
	return (psi_new.conj() @ H @ psi_new).real


def compute_energy_np(angle,  H,psi):
	U=sp.linalg.expm(-1j*angle*H)
	psi_new = U @ psi
	return (psi_new.conj() @ H @ psi_new).real



# initial solver condition
angle_0 = np.pi/np.exp(1.0)
angle_0_jnp, angle_0_np = jnp.array([angle_0,], dtype=jnp.float32), np.array([angle_0,], dtype=np.float32)


# scipy optimization

t_i = time.time()
res = sp_opt.minimize(compute_energy_np, angle_0_np, args=(H_np,psi_np,),  method='BFGS', tol=1e-3)
t_f = time.time()

print(res.x, res.fun, res.success, res.nfev, res.njev)
print('scipy calculation took {:0.6f} secs.\n'.format(t_f-t_i))


# jax.scipy optimization

compute_energy_jnp(angle_0_jnp, H_jnp, psi_jnp)


t_i = time.time()
res = jsp_opt.minimize(compute_energy_jnp, angle_0_jnp, args=(H_jnp,psi_jnp,),  method='BFGS', tol=1e-3)
t_f = time.time()

print(res.x, res.fun, res.success, res.nfev, res.njev)
print('JAX calculation took {:0.6f} secs.'.format(t_f-t_i))

which gives the output:

[1.15572735] -0.18787633820120936 True 2 1
scipy calculation took 0.001452 secs.

[1.1557274] -0.18787628 True 1 1
JAX calculation took 4.770535 secs.

[lukepfister]

I've noticed the same, but for sufficiently large problems the jax BFGS outperforms scipy. Could be "initialization" overhead?

[mgbukov]

I thought so too, but even if you call the jax optimization a few times consecutively, the time required for each calculation remains roughly the constant.

I wonder if jax for some reason tries to compute the derivative, while scipy does everything numerically.

[jakevdp]

For operations on very small arrays, JAX tends to be orders of magnitude slower than numpy, and this is expected because JAX's computation model involves some per-computation overhead that is non-negligible for small operations. Here's a simpler example:

In [1]: import numpy as np                                                                                                                            

In [2]: import jax.numpy as jnp                                                                                                                       

In [3]: x = np.random.rand(10, 10)                                                                                                                    

In [4]: %timeit np.dot(x, x)                                                                                                                          
1.05 µs ± 21.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [5]: y = jnp.asarray(x)                                                                                                                            
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

In [6]: %timeit jnp.dot(y, y).block_until_ready()                                                                                                                  ``       
129 µs ± 5.6 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

If you do operations on larger arrays, with proportionally more time spent on the actual computation, JAX is often faster. For example:

In [7]: x = np.random.rand(2000, 2000)                                                                                                               

In [8]: %timeit np.dot(x, x)                                                                                                                         
96.3 ms ± 6.68 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [9]: y = jnp.asarray(x)                                                                                                                           

In [10]: %timeit jnp.dot(y, y).block_until_ready()                                                                                                    
45.9 ms ± 602 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In your case, your operations are on very small arrays, so I wouldn't expect JAX to out-perform numpy/scipy in terms of total computation time.

[mgbukov]

I see, that makes sense! Thanks, @jakevdp and @lukepfister!