Poor sampling performance with some complex posteriors compared to HMC

[sefffal]

Hi all, I have isolated a relatively lightweight example where I find that HMC is significantly outperforming PT. I tested most variations supported in Pigeons: Slice Sampler & AutoMALA, fixed, variational and stabilized-variational).

A corner & trace plot are attached below with HMC in blue and PT in gold.

A script to produce this plot is available here: https://github.com/sefffal/OrbitPosteriorDB/blob/main/models/astrom-GL229A.jl
Use the latest #main commit of Octofitter (e.g. ] add Octofitter#main).

I would be curious to understand better why PT is struggling so much on this target, and if there is a way to improve performance to be at least within the same ball-park.

The HMC series in this plot is not exactly converged since there's pretty high correlation between samples, but it nonetheless successfully explores the posterior while PT is stuck in a much smaller region.

[alexandrebouchard]

Thanks William, I'm excited to see what happens here and hopefully make improvements based on what we observe! So far the focus has been on problems where HMC performs poorly, but I agree that having reasonable performance on problems where HMC work well is very useful too.

I would have expected automala to do well here since there are some funnel-like shapes in the pair plot. It could be the round-based initial step size adaptation running wild. That part of automala is relatively less understood compared to the per step auto step.

Another possibility is that long trajectories are really needed for this one. With @tgliu0406 we prototyped recently an auto HMC, it would be interesting to see if it performs better than automala on that problem.

Since multiple chains is not needed here, it might be worth focussing initially on running pigeons with 1 chain.

I'll start with trying the reproducibility script. Thanks again!

[alexandrebouchard]

For reference, here are the stats from the HMC run:

  iterations:                                   5100
  ratio_divergent_transitions:                  0.0
  ratio_divergent_transitions_during_adaption:  0.03
  n_steps:                                      4095
  is_accept:                                    true
  acceptance_rate:                              0.9686561630493195
  log_density:                                  19.1778839507093
  hamiltonian_energy:                           -15.535406645231566
  hamiltonian_energy_error:                     0.04783412097392059
  max_hamiltonian_energy_error:                 0.088189176494609
  tree_depth:                                   12
  numerical_error:                              false
  step_size:                                    0.0004985427615975861
  nom_step_size:                                0.0004985427615975861
  is_adapt:                                     false
  mass_matrix:                                  DenseEuclideanMetric(diag=[0.16992803377878063, 7.203 ...])

One other guess is that the preconditioner might not be off...

[alexandrebouchard]

Some stats based off

chain_pt,pt = octofit_pigeons(
    model,
    # 12 rounds is about the same wall-clock time as the above HMC
    n_rounds=15,
    n_chains=1,
    n_chains_variational=0,
    variational=nothing,#GaussianReference(),
    # explorer=Compose(SliceSampler(),AutoMALA()),
    explorer=AutoMALA(),
)

julia> pt.shared.explorer.step_size
0.004920727982147332

julia> pt.shared.explorer.estimated_target_std_deviations
12-element Vector{Float64}:
 0.009694861110282131
 0.011218920887184406
 9.84596575830136e-5
 0.008446406116308973
 0.007492746456207485
 0.006611899705535479
 0.02192204882563213
 0.006264372034808068
 0.006598652035268915
 0.014764926897502917
 0.006588233373333677
 0.0019844912584754994

[miguelbiron]

For reference, here are the stats from the HMC run:

  iterations:                                   5100
  ratio_divergent_transitions:                  0.0
  ratio_divergent_transitions_during_adaption:  0.03
  n_steps:                                      4095
  is_accept:                                    true
  acceptance_rate:                              0.9686561630493195
  log_density:                                  19.1778839507093
  hamiltonian_energy:                           -15.535406645231566
  hamiltonian_energy_error:                     0.04783412097392059
  max_hamiltonian_energy_error:                 0.088189176494609
  tree_depth:                                   12
  numerical_error:                              false
  step_size:                                    0.0004985427615975861
  nom_step_size:                                0.0004985427615975861
  is_adapt:                                     false
  mass_matrix:                                  DenseEuclideanMetric(diag=[0.16992803377878063, 7.203 ...])

One other guess is that the preconditioner might not be off...

Some questions/comments

• Is n_steps the number of leapfrog steps on an average iteration? Because wow. That is 4 times the default (2^10) in stan if im not mistaken
• The DenseEuclideanMetric might also be doing a lot of lifting. We haven't implemented that yet in pigeons
• Finally, the nasty concentration on lower dimensional submanifolds that are not axis-aligned definitely should kill the slicer, so that is not too surprising.

[sefffal]

Hi @miguelbiron , thanks for taking a look.

Yes, I find that increasing the maximum number of steps is important for getting acceptable performance from HMC on these problems. Sometimes the data are not that constraining, in which case NUTS cuts the trajectories short and we average around 2^7 leapfrog steps.
In this case, almost every iteration uses 2^12 steps so an even higher limit might be warranted.

The DenseEuclideanMetric is critical, and HMC effectively does not work on these problems without it.

I'll note that while surely there are some reparameterizations that can help with this specific dataset, it is very hard to find parameterizations that solve the problem for all datasets.

[miguelbiron]

No problem and thanks for clarifying! It's good to know at least that NUTS works here because you're taking it to extreme tree sizes (and even then, I see it goes beyond the limit 10% of the time), and also using the Dense preconditioner.

From a Pigeons perspective, we instead would prefer to stick with a cheaper explorer and instead crank up the number of chains. But it's weird that that doesn't work. So I'll be doing some tests now to see if I can find a combination of dumb explorer with more chains that gives as good performance as NUTS in the same wallclock time. I'll keep you posted.

[sefffal]

Thanks @miguelbiron, I would say this matches my experiences with Pigeons too. The SliceSampler combined with ~40 chains and ~40 variational chains can usually power through just about anything. Unfortunately there are some challenging cases where even that leads to some weird results.

For example, this figure (from a different but related posterior) shows how sometimes PT leaves these gaps along the posterior that ought to be connected (see e.g. inclination vs eccentricity):

[miguelbiron]

The axis aligned gap kinda make sense for a slice sampler failure mode. That could fixed by substituting with a Hit and Run wrapper for the slicer that we implemented for the AutoStep paper with @tgliu0406. Or even fully switching to autoRWMH which showed promising results. We should be releasing a package with these additional explorers in the short term.

I'll go back to the other reproducible example now.

[alexandrebouchard]

Good point... I think this would be a good test bed for the auto step methods + Hit and Run...

[alexandrebouchard]

I noticed that the script does not give exactly the same result when I rerun it. Looking at the output of the initializer, I think it might be a cause of this. In my first run I got

┌ Info: Starting point drawn
└   initial_logpost = -850885.8759688964

vs for the second:

┌ Info: Starting point drawn
└   initial_logpost = -1.5588419366108126e6

[sefffal]

Hi @alexandrebouchard, my bad. If the starting point aren’t given explicitly they sample from the default RNG.

You should be able to:

• seed it with Random.seed!(0)
• Initialize with an explicit RNG argument Octofitter.default_initializer!(rng, model)
• Specify the starting points manually (make sure there’s at least one per chain):

model.starting_points = [
    model.link([1,2,3…]), # start values for each parameter 
    # repeat for each chain, or use fill(…)
]

[miguelbiron]

I also noticed the non-determinism, which is more annoying in that the HMC sample themselves are completely different from one run to the other. Like I've never seen on my runs the sinusoidal shapes you show above. So unless I f'd up something on my end, I wouldn't trust too much those samples either.

BTW, I'm doing a 20 round run with 60 chains = 4Lambda, autoRWMH with n_refresh=32. I did this following my heuristic advise from the NRST paper, and I actually get

• max energy autocorrelation <~ 90% (due to increased n_refresh)
• almost exact equirejection
• reversibility rate and acceptance probs consistent with autoRWMH in other test sets.
• With this I observe roundtrips from round 13 onwards. Which is right about the time when the normalizing constant begins stabilizing.

Main takeaway: you haven't reached adaptation convergence if you don't see

• Lambda ~ 15
• logZ ~ -16

[sefffal]

Okay, I can reproduce the non-determinism of the pathfinder initialization. I'll try to track that down.

[sefffal]

Wow, this is frustrating.

I think I found non-determinism issues with both Pathfinder.jl, and with AdvancedHMC.jl.
Clearly, I need to add some integration tests here.

For pathfinder, it is non-deterministic when using multiple threads (can be worked around by passing nruns=1 to Octofitter.default_initializer!).

For AdvancedHMC, I am passing an rng object explicitly, but it is ignoring the argument and using the default global RNG.

[sefffal]

I pushed an update to the script that works around those issues for now.