PyMC/PyTensor Implementation of Pathfinder VI

[aphc14]

Another version to draft PR #386 which uses more of PyTensor's symbolic variables and compiling functions.

Questions for Review

1. Which implementations should I continue for future improvements?
2. Are there additional PyTensor optimisations we could leverage?

[aphc14]

Suppose the preferred approach is to stick with symbolic variables in PyTensor than the other non-symbolic approach in #386. In that case, I'd be happy to refactor the Multipath Pathfinder implementation in #386 to use symbolic variables and pytensor.function.

[aphc14]

This version runs much faster than #386, but the codes are messier due to the numerous pytensor symbolic variables created for the compiled pytensor functions (see the lines of code between def compute_logp and def single_pathfinder). Any suggestions for a cleaner setup would be appreciated

[fonnesbeck]

Cleaner to use a data class? Don't know.

[aphc14]

yep, I agree. dataclass now added

[fonnesbeck]

This PR will provide that, no?

[aphc14]

the latest commit addresses this

[aphc14]

Need to make an important change related to how important sampling is done. Based on some tests, for trickier posteriors, psir (Pareto smoothed importance resampling) tends to cause many large peaks. In contrast to the reference posterior (what you’d get using NUTS), it doesn’t have such peaks.

Turning off resampling, you'd get psis instead, and the final posterior better resembles NUTS, so you don't get the weird peaks behaviour. But this would differ from the original paper, which uses psir.

Since the choice of importance sampling (IS) can have a big impact on the final posterior, and there are several IS methods, I plan to use a class variable that controls how IS is done based on the user inputs. I'm thinking of making psis (and not psir) the default IS behaviour as the safest and most generally reliable option.

Shouldn't take long to fix.

[aphc14]

Results are looking pretty good now! Have done a couple of tests with a few datasets and returns a posterior similar to NUTS sometimes. Speed is comparable to ADVI, and for larger models it can be about 2x or more times faster than ADVI. I've made some changes to the code to make it more efficient and more stable.

https://gist.github.com/aphc14/019f172e44c0c767a4f48e91a045e896

[fonnesbeck]

What happens when either is not finite? Clearly the the arrays are uninitialized, so is it handled somewhere further down the line?

[aphc14]

yea it is. if neither is not finite, the computed value and grad at x0 doesn't get added to history_manager, however, lbfgs continues until it stops. at the very end of lbfgs, it'll store suitable value and grad values. if the arrays are uninitialised, then there is this check in place on line 114: elif (result.status == 2) or (history_manager.count <= 1) which lets the user know a path has failed to initialised and it'll move onto the next path.

[fonnesbeck]

The importance_sampling argument isn't passed anywhere and is only relevant to the function when it is "none". Is this intended, or is something missing?

[aphc14]

importance_sampling is passed into multipath_pathfinder (the return for this function is return _importance_sampling(...)) and
convert_flat_trace_to_idata. when user chooses "none", it'll impact the sample size for idata, so convert_flat_trace_to_idata checks for "none".

for "psis", "psir", "identity", they are handled in pymc-experimental/pymc_experimental/inference/pathfinder/importance_sampling.py

[fonnesbeck]

Add type hints throughout (and docstrings, ideally).

[fonnesbeck]

Incomplete docstring and missing type hints.

[ricardoV94]

Add something to the docs?

[ricardoV94]

You may want to consider a value_and_grad function, as it can avoid many repeated operations?

[aphc14]

for my understanding, does "repeated operations" refer to value being computed twice: once in fn(x0) and again as part of grad_fn(x0)? or does "repeated operations" refer to something else?

I could change scipy.optimize.minimize(fun=logp_dlogp_fn, jac=True) where logp_dlogp_fn is:

first approach

logp_dlogp_fn = model.logp_dlogp_function(
    ravel_inputs=True,
    dtype="float64",
    mode=pytensor.compile.mode.Mode(linker="cvm_nogc"),
)
logp_dlogp_fn.set_extra_values({}) # without this, i'd get an error
logp_dlogp_fn._pytensor_function.trust_input = True

I would prefer this approach since I can toggle between jacobian=True/False if I need to:

second aproach

outputs, inputs = pm.pytensorf.join_nonshared_inputs(
        model.initial_point(),
        [model.logp(jacobian=jacobian), model.dlogp(jacobian=jacobian)],
        model.value_vars,
    )

logp_dlogp_fn = compile_pymc(
    [inputs], outputs, mode=pytensor.compile.mode.Mode(linker="cvm_nogc")
)
logp_dlogp_fn.trust_input = True

how does the second approach look?

[ricardoV94]

repeated operations means operations that are shared between the logp and dlogp functions. Second approach is fine, except the hardcoded mode, you should allow the user to pass compile_kwargs where they can define a custom Mode if they need to

[aphc14]

ahh I see. In that case, I might have considered the repeated operations here: #387 (comment)

I can make the changes if repeated operations weren't considered

[ricardoV94]

Why is this an Op? How is it constructed? Why can't you work directly with PyTensor graphs?

[aphc14]

LogLike is initialised using logp_func which is a compiled function that we already have. logp_func cannot take in a pytensor variable since its already compiled, so it has to be a numpy array with ndim=1. I needed to vectorise logp_func so that it can take in an array with ndim=3.

I wasn't sure how to make logp_func take in the symbolic input with batched dims, phi or psi, which was why I used Op.

how would you make an already compiled function receive symbolic pytensor variables as inputs?

[ricardoV94]

you can use pt.vectorize or pytensor.graph.replace.vectorize_graph

[ricardoV94]

These are quite big atols

[aphc14]

they are. I think its due to the Pathfinder algorithm. Have compared it to Stan's multipath Pathfinder (teal) and the mu appears about -2 away from the reference posterior (red) which is taken from posteriordb.

the image below are pymc's pathfinder and NUTS (blue) estimate. the default multipath pathfinder (orange) agrees with Stan's default multipath pathfinder (teal) above. getting closer to the reference or NUTS posterior requires a different setting for jitter and num_paths.

as for tau, the original reference value was 4.15. just had a look at what pymc NUTS returns, and its somewhere around 3.5, whereas estimates from pymc Pathfinder is around 3.0. I guess I can set it to np.testing.assert_allclose(idata.posterior["tau"].mean(), 3.5, atol=0.6)

[ricardoV94]

Can you test something more than the shapes?

[aphc14]

Relates to #387 (comment)

would the operations be shared between logp and dlogp here considering outputs comes from both model.logp and model.dlogp? Or it wouldn't since they are separated (compile fn for logp, compile separate fn for dlogp)?

This would help me determine if I need to change the existing codes to the one below:

outputs, inputs = pm.pytensorf.join_nonshared_inputs(
        model.initial_point(),
        [model.logp(jacobian=jacobian), model.dlogp(jacobian=jacobian)],
        model.value_vars,
    )

logp_dlogp_fn = compile_pymc(
    [inputs], outputs
)
logp_dlogp_fn.trust_input = True

[ricardoV94]

you have to compile both outputs (logp and dlogp) together, then the compiled function will avoid repeated operations.

If you compile different functions it won't. Does that answer your question?

[fonnesbeck]

It would be useful to have a progress bar. At the moment we only get number of parameters and maxcor, which is hard to translate into expected runtime.