Potential issue with SVI using init_params argument

numpyro
1

I’m trying to utilize SVI with the AutoDelta guide. I run the optimization for a certain number of steps, and save the final parameters to a dictionary and check if the parameters have converged to my liking. If not, when I re-run the optimization (with a reduced learning rate), and pass in the last saved parameters as initial parameters for this new optimization stage, the init_loss printed out is the exact same as the init_loss printed out during the first optimization. This suggests to me that maybe the init_params are not being utilized when passed in to svi.run() as the init_loss should in principle be equivalent to the final loss from the first optimization (unless i’m misunderstanding). Any help would be appreciated.

2

not sure i follow but do you use pyro.clear_param_store()?

3

Thanks for the reply. I’m happy to clarify if there’s things that I have not articulated well, but the gist of it is that it seems when I use svi.run() with the init_params argument set equal to the optimized parameters from a previous run, these parameters are seemingly not utilized. From my understanding, these parameters should be passed in as a dictionary with keys equal to the parameter name, and values equal to the desired initialization values. I’m using the numpyro implementation for this

4

In my code it is something like this:

guide = AutoDelta(model_fn, prefix="")

svi = SVI(model_fn, guide, optim=adam(.000003), loss=Trace_ELBO(num_particles=10),
          )

result = svi.run(jax.random.PRNGKey(0), num_steps = num_steps,progress_bar = True)
fitted_params = result.params
svi2 = SVI(model_fn, guide, optim=adam(.00000003), loss=Trace_ELBO(num_particles=10),
          )
result2 = svi2.run(jax.random.PRNGKey(0), num_steps = num_steps, progress_bar = True, init_params= {{key.replace("__loc",""): fitted_params[key] for key in fitted_params}})

The issue is that in svi2.run(), the progress bar shows the same initial loss as svi.run(). My confusion comes from the fact that since I’m initializing svi2.run() with parameters recovered from the first svi.run(), the initial loss should be different (and equal to the final loss achieved by the initialization parameters). It almost seems as if the two runs are independent, and they are being initialized to the same starting parameters, despite init_params being passed to svi2.run()

5

are you sure this dictionary manipulation is correct? what if you do init_params=fitted_params?

6

Ah I see, yes. I removed the suffix because I thought the names had to match the param names given in the model_fn. Not doing this fixes the issue. Thanks!