Handling failed samples in custom likelihood function

I am working with a model that goes through four steps to get a likelihood function:

  1. Draw samples of parameters
  2. Compute fixed points of a non-linear system
  3. Compute first-order Taylor approximation of the non-linear system around the fixed points via perturbation
  4. Compute likelihood using first-order approximation

It’s all implemented in Aesara, so I have gradients through the whole procedure and I can run NUTS on it for low-dimensional parameter spaces. When I try to apply it to higher dimensional parameter spaces, I can’t guarantee that step 3 will be successful, but I still want to estimate the model via MCMC.

If I were rolling my own Metropolis sampler, I would add a check that would automatically reject any samples that caused the perturbation algorithm to fail. Is there any way to get this kind of behavior from the NUTS sampler? Would something like setting the log-likelihood to a specific value (-inf? 0? nan?) if a sample fails induce the desired behavior? Or do I need to modify the step function itself to add these checks myself (i hope not)?

Alternative/supplementary question:

Since something like what I proposed above (throwing away bad samples) might break detailed balance anyway, I figured I’d also try using ADVI or SMC to get posterior estimates. Here I am still having a problems, I suspect related to using ifelse. Here’s the relevant chunk of code. B_hat_solved is the system’s perturbation solution, which will be all nan if no solution exists.

    _, _, B_hat_final, _ = cycle_result
    B_hat_solved = B_hat_final[-1]
    pert_failure_flag = at.or_(at.any(at.isnan(B_hat_solved)), at.any(at.isinf(B_hat_solved)))
    
    log_likelihood = ifelse(pert_failure_flag,
                            at.constant(0.0, dtype='float64'),
                            kalman_likelihood(A, B, B_hat_solved, C, D, observed_data))

My understanding from reading the aesara docs is that if I use ifelse with linker=cvm, it will only evaluate one or the other branch of the conditional. This is important because I need to invert B_hat_solved inside of the likelihood function; if the perturbation failed that whole branch will throw errors.

When I compile this function with aesara.function it works as expected, I can give insane parameter values and I get a zero likelihood without any complaining. As soon as I hand the model to ADVI, however, it immediately fails with FloatingPointError: NaN occurred in optimization. Since I’m also getting Linalg warnings about singular matrices, my suspicion is that both branches of the ifelse are being evaluated, even though I changed the linker in my .aesararc.txt.

Any way to attack this problem via ADVI?

In general you should be able to switch your likelihood to -inf so that it is never accepted by NUTS. That is pretty common in PyMC. For the ifelse question I am not sure. In general ifelse is a bit more finicky so if you are not too concerned with loss in perfomance because of evaluation of all the branches you can use switch instead.

1 Like