How to assign different learning rates to different parameters


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())


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

1 Like

Thank you @fehiepsi! This worked great.