how can I extract the log marginal likelihood from SMC?

[murphyk]

how can I extract the log marginal likelihood from SMC?
Based on https://blackjax-devs.github.io/blackjax/examples/TemperedSMC.html,
It seems from that I would need to use something like

tempered = blackjax.adaptive_tempered_smc(...)
state = random()
state = tempered.init(state)
for t in range(T):
   step, info = tempered.step(key, state)
   ll[t] = info.log_likelihood_increment
logZ = jnp.sum(ll)

The functionality seems to be buried in this code, but I'm not entirely sure what's going on.

It would be very helpful to make a little demo comparing this to AIS.

[rlouf]

Indeed the functionality is a little buried and is not presented in any example as you noticed. The code that computes these log_likelihood_increments from the weights is right below the line you linked to:

def _normalize(log_weights):
    """Normalize log-weights into weights and return resulting weights and log-likelihood increment."""
    n = log_weights.shape[0]
    max_logw = jnp.max(log_weights)
    w = jnp.exp(log_weights - max_logw)
    w_mean = w.mean()

    log_likelihood_increment = jnp.log(w_mean) + max_logw

    w = w / (n * w_mean)
    return w, log_likelihood_increment

Do you have a particular example we could add in mind?

As you might have seen in #279 the SMC part of Blackjax is going to be substantially refactored, hopefully to make it more general and flexible. It would be great to have your input given your experience and current projects.

[gerdm]

Hi @rlouf!
To follow up a bit more on this, we’re looking to compute the marginal log-likelihood of various latent variable models defined using flax.

A simple example is the comparison of the marginal log-likelihood between two different VAE architectures using SMC/AIS.

[rlouf]

Hey-did you figure it out in the end?

[murphyk]

On Thu, Nov 24, 2022 at 3:43 AM Rémi Louf ***@***.***> wrote: Hey-did you figure it out in the end?

No sorry i didn’t get time to try this It would still be useful though

— Reply to this email directly, view it on GitHub <#394 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABDK6EGSYN5HIIAVWGX53E3WJ5IHRANCNFSM6AAAAAARR5KIBU> . You are receiving this because you authored the thread.Message ID: ***@***.***>

-- Sent from Gmail Mobile

[rlouf]

To get the normalizing constant you need to sum the loglikelihood_increments, as is done in this test. From the look of it the code snippet you provided should work (provided ll is an array of length the number of SMC steps).

For anyone else reading this (I assume the OP knows this already), an explanation can be found in "An Introduction to Sequential Monte Carlo", in sections 16.4.1 (Eq. 16.9) and section 17.1