Optimizer:Minimizer(BFGS) Not Working Properly as in Julia

Hello everyone,

I’m writing to you because I’m trying to run the Numpyro minimizer with the ‘BFGS’ mode to obtain best-fits but basically the output gives me the parameters at the border of the priors. That is very strange because when I run the same thing but with the Adam minimizer I obtain pretty good results.

I’m pretty sure that there is something that I’m doing wrong but I can’t see exactly where, on top of that I have a colleague that is performing basically the same analysis but in Julia and has very good results and both in good agreement.

def model():        
   b1 = numpyro.sample("b1", dist.Uniform(0.1, 5))         
   b2 = numpyro.sample("b2", dist.Uniform(-5, 5))         
   g2 = numpyro.sample("g2", dist.Uniform(-4, 4))         
   g21 = numpyro.sample("g21", dist.Uniform(-8, 8))         
   c0 = numpyro.sample("c0", dist.Uniform(-500, 500))         
   c2 = numpyro.sample("c2", dist.Uniform(-500, 500))         
   c4 = numpyro.sample("c4", dist.Uniform(-500, 500))         
   NP0 = numpyro.sample("NP0", dist.Uniform(-3, 3))         
   NP20 = numpyro.sample("NP20", dist.Uniform(-100, 100))         
   NP22 = numpyro.sample("NP22", dist.Uniform(-100, 100))         
   avir = numpyro.sample("avir", dist.Uniform(0, 20))         
   params_npyro = jnp.array([b1, b2, g2, g21, c0, c2, c4, NP0, NP20, NP22,avir])                                  

   chi2_val = chi2(params_npyro) #chi2 is a particular function that computes the chi2        
   numpyro.factor("chi2_factor", -0.5*chi2_val)     
   
for i in range(0,5):         
    guide = AutoDelta(model)         
    optimizer = Adam(0.001)#numpyro.optim.Minimize()#          
    svi = SVI(model, guide, optimizer, loss=TraceMeanField_ELBO())         
    key = jax.random.PRNGKey(i)         
    result = svi.run(key,50000)    
    state = result.state         
    params_new = svi.get_params(state)

If now I try to use the same with BFGS and Minimizer it doesn’t work :

for i in range(0,5):         
    guide = AutoDelta(model)         
    optimizer = =numpyro.optim.Minimize('BFGS')        
    svi = SVI(model, guide, optimizer, loss=TraceMeanField_ELBO())#Work bad in the     same way with Trace_ELBO()    
    key = jax.random.PRNGKey(i)         
    result = svi.run(key,50000)      
    state = result.state         
    params_new = svi.get_params(state) 

Is there anyone that can help me with that ?

Thank you in advance!!

There are a couple of hyperparameters that control the behavior: https://github.com/jax-ml/jax/blob/23db456a8acd01a04ed5a9f87f8265cc21703926/jax/\_src/scipy/optimize/bfgs.py#L79-L82 Could you try to play with them? You might also want to put numpyro.enable_x64() at the top of your script to enable float64.