Penalize different posterior estimates

Hi All,
I am looking for a way to penalize different posterior distributions arising from different training sets.
To be more precise, assume we are given the Pyro (shortened) tutorial example:

def model(is_cont_africa, log_gdp=None):
    a = pyro.sample("a", dist.Normal(0.0, 10.0))
    b_a = pyro.sample("b_a", dist.Normal(0.0, 1.0))
    sigma = pyro.sample("sigma", dist.Uniform(0.0, 10.0))

    mean = a + b_a * is_cont_africa

    with pyro.plate("data", len(is_cont_africa)):
        pyro.sample("obs", dist.Normal(mean, sigma), obs=log_gdp)

Now, in each epoch I want to penalize a difference in the posterior distributions over b_a given two different input data sets, x and y, where y includes additional samples compared to x. x is used for ELBO gradient estimate, though.
Hence, instead of only maximizing the ELBO only I want to optimize the following objective:

max ELBO - KL[p(b_a | x) || p(b_a | y)]

Is there any way to achieve something like this in pyro?
If this is not implemented yet, but intended, I am happy to contribute!

that should be pretty straightforward to do. you can either follow some of the approaches described here or you can add an appropriate pyro.factor statement to your model. note that a given log_factor that is in your model appears with a positive sign in the elbo (and a negative sign in the elbo loss). vice versa for a log_factor that appears in your guide

though iā€™m not really sure what your x and y are and how they are related so you would presumably need to do inference twice e.g. once w.r.t. y and then subsequently w.r.t. x with the additional regularizer

1 Like