How to assign different learning rates to different parameters

numpyro
1

Hi,

I am doing SVI and was wondering if there is a way to choose different learning rates for different parameters?

Currently I am just doing the following:

...
...
optimizer = numpyro.optim.Adam(step_size=learning_rate)
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())
...
...

Thanks,
Atharva

2

Hi @atharvahans, you can use optax for this. See this example.

1 Like
3

Thank you @fehiepsi! This worked great.