Numpyro: Refactor MCMCKernel code and make them composable

[Joshuaalbert]

Currently, it's very difficult to work on code inside mcmc because all the logic is mixed up in various places. It makes it impossible, for example, to modularly compose sampling. I want to write nested sampling using the composition MetropolisHastings(NUTS). The NUTS proposes a new state that samples the likelihood weighted by the prior, and then the MetropolisHastings accepts only those where the likelihood at the proposed point is better than a certain value (not necessarily the previous). This sort of composition is desireable for cleaness of code, otherwise I need to start messing around with monolithic code.

Specifically, MCMCKernel is written so that you call a monolithic function hmc that accepts algo argument and then it chooses what to do inside this big function. I understand that HMC and NUTS share a lot of code in common but these code bits should be modular and go in their own functions. I would personally like to see dependency injection here, and composable MCMCKernels.

[fehiepsi]

Yeah, we like to make MCMCKernel as flexible as possible. We want to understand more about your refactor request. Let me illustrate what is going on in the pseudo-codes below:

Currently, MCMCKernel contains 3 methods

class MCMCKernel(ABC):
    def postprocess_fn(self, model_args, model_kwargs):
        return identity

    @abstractmethod
    def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs):
        raise NotImplementedError

    @abstractmethod
    def sample(self, state, model_args, model_kwargs):
        raise NotImplementedError

With composed kernel, I guess you mean something like?

class ComposeKernel(MCMCKernel):
    def __init__(self, kernels):
        self.kernels = kernels

    def postprocess_fn(self, model_args, model_kwargs):
        def fn(x):
            for kernel in self.kernels:
                 x = kernel.postprocess_fn(model_args, model_kwargs)(x)
            return x
        return x

    def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs):
        return self.kernel[0].init(rng_key, ...)

    def sample(self, state, model_args, model_kwargs):
        for kernel in self.kernels:
            state = kernel.sample(state, model_args, model_kwargs)
       return state

I am not sure if that is suitable for your usage. In addition, I don't quite understand the monolithic stuff that you mentioned. I guess you want to refactor hmc into smaller functions? That sounds reasonable to me (e.g. we can move hmc_next, nuts_next out of hmc) but does it relevant to the composability of MCMC kernels?

[neerajprad]

@Joshuaalbert - The internals of NUTS are complicated because of various performance related optimizations but this complexity should be internal. Is there anything that you would want to change at the level of MCMCKernel interface to allow for better composability?

If I understood your problem correctly, my suggestion would be to use composition within your kernel (instead of inheritance). Arbitrarily allowing for compositions of MCMC kernels as in @fehiepsi's example is also a possibility but I am unsure how useful it would be to generalize at this point (for instance what does the intermediate data structure look like?). Is there anything within the NUTS kernel or the sample HMCState data structure that you need to help you with this? For this, if you need to muck around with the internals of the NUTS class, then that is definitely something that we should fix so that you shouldn't have to.

class MetropolisNUTS:
    def __init__(self, ..):
        self.nuts_kernel = NUTS(..)

    def mh(self, new_sample, old_sample):
         ...

    def sample(self, state, model_args, model_kwargs):
        proposed_sample = self.nuts_kernel.sample(state, model_args, model_kwargs)
        return self.mh(proposed_sample, sample)

[Joshuaalbert]

Thanks for the replies. In terms of composing, there isn't a need for general composability. @neerajprad your code is actually just was TFP does for the MH kernel. Their MH kernel accepts an inner_kernel argument and then inside its one_step (equiv to your sample) does:

 [
          proposed_state,
          proposed_results,
      ] = self.inner_kernel.one_step(
          current_state,
          previous_kernel_results.accepted_results)
  log_acceptance_prob = proposed_results.target_log_prob \
    - previous_kernel_results.accepted_results.target_log_prob \
    + proposed_results.log_acceptance_correction
  #accept or reject based on log_acceptance_prob

log_acceptance_correction is to be returned as part of the results of inner_kernel or else assumed zero. This is the type of composability I need and see a use for. I'd need the MH kernel which takes an inner_kernel as well as an intermediate kernel that also takes an inner_kernel and computes the log_acceptance_correction required for nested sampling.
In the end it would look something like MH(Intermediate(NUTS or another)).

I feel like we don't need too much additional code actually, other than those two kernels which I can write.

[Joshuaalbert]

Just wondering what the reason for not using *args and **kwargs here,

    def postprocess_fn(self, model_args, model_kwargs):

    def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs):

    def sample(self, state, model_args, model_kwargs):

[Joshuaalbert]

@fehiepsi I need to compute the log-likelihood at the current proposal inside sample. What's the best way to do that? It looks I need the model and then call numpyro.infer.util.log_likelihood(model, proposed_state, *model_args, **model_kwargs). Is that correct? Then, do I need to manually pass in model to the sampler and construction time or is it possible to get the model some other way?

[fehiepsi]

Just wondering what the reason for not using args and *kwargs here,

It is just our preference to use model_args, model_kwargs for intermediate functions/methods (i.e. those not interacted by users). That way, we don't have to worry about name conflicts when we pass *args, **kwargs from user inputs to those intermediate utilities. (e.g. a conflict will happen when user's model contains state keyword).

compute the log-likelihood at the current proposal inside the sample.

Yes, it is basically like that. You will need to convert state to sample use kernel.postprocess_fn(args, kwargs)(state), then reshape the sample to have a batch dimension 1, then use log_likelihood utility to get, then reshape again to remove the leading 1 dimension. We expose model from HMC/NUTS here.

@neerajprad This is a usage case where batch_ndim in log_likelihood, Predictive is useful. Users won't need to do reshape if we support that keyword.

[fritzo]

Just wondering what the reason for not using args and *kwargs here,

In addition to the added safety in avoiding name conflicts, there is a slight performance improvement by avoiding unrolling and constructing the tuple and dict each call.

[neerajprad]

@neerajprad This is a usage case where batch_ndim in log_likelihood, Predictive is useful. Users won't need to do reshape if we support that keyword.

I usually find it easier to think of a single batch dim for these utilities in the spirit of jax.vmap, but I think you have a few examples to suggest that having batch_ndims will be useful, so let us include this in our utilities. I think the default value of 1 would lead to no change to existing code, and users should be able to adjust this to 0 like in this case, or 2 for multi chains.

[Joshuaalbert]

Thanks for the information. It turns out I'll actually need to touch some NUTS/HMC code to implement constrained HMC. Basically, after each spatial step in constrained HMC a constraint (infinite potential barrier) is checked and if violated then the particle bounces off the barrier. This requires the addition of another cond, which is not great for accelerators but necessary.

[1] Constrained HMC. https://arxiv.org/pdf/1005.0157.pdf

EDIT: this thread is somehow merging into the idea of nested sampling implementation

[fehiepsi]

A composable solution, suggested by @neerajprad above, has worked well with the recently added HMCGibbs class. Other composable kernels might follow the same pattern. Please feel free to reopen this issue if needed.