Structured SVI

numpyro
1

Is it possible to model a custom dependency structure in the posterior guide between random variables with separate definitions in the model?

Here is an example of what I’d like to accomplish, but I’m not sure how to do that with the naming constraints in the guide.

def model(x, options):
  gamma_g = numpyro.sample("gamma_g", dist.Gamma(options.conc_g, options.rate_g))
  gamma_l = numpyro.sample("gamma_l", dist.Gamma(options.conc_l, gamma_g))
  ...
  return

def guide(x, options):
  w_df = numpyro.param("w_df",  1., constraint=constraints.positive)
  w_scale = numpyro.param("w_scale", jnp.eye(2), constraint=constraints.lower_cholesky)
  # i know the next line is incorrect, but this is where I'm unclear on how to proceed
  gamma_joint = numpyro.sample("[gamma_g, gamma_l]", tfpdist.WishartTriL(w_df, w_scale))

Could this be accomplished with a substitute call in the guide that replaces gamma_g and gamma_l by their respective sampled values in gamma_joint?

1 Like
2

Hi @nmancuso, it is certainly possible to introduce joint auxiliary sites and deterministically compute other latent variables from those, as is done in NumPyro’s AutoContinuous guides. The only trick is to mark the new auxiliary sites with infer={"is_auxiliary": True}. For example you could write:

def guide(x, options):
    w_df = numpyro.param("w_df",  1., constraint=constraints.positive)
    w_scale = numpyro.param("w_scale", jnp.eye(2), constraint=constraints.lower_cholesky)
    gamma_joint = numpyro.sample(
        "[gamma_g, gamma_l]",
        tfpdist.WishartTriL(w_df, w_scale),
        infer={"is_auxiliary": True},
    )
    numpyro.sample("gamma_g", dist.Delta(gamma_joint[..., 0]))
    numpyro.sample("gamma_l", dist.Delta(gamma_joint[..., 1]))
2 Likes
3

Nice, thanks @fritzo, this is great!

1 Like