Turing and AdvancedHMC give different adaptors for NUTS

[ysfoo]

Minimal working example

using Turing, AdvancedHMC, AbstractMCMC, Random

@model function gdemo(x, y)
    m ~ Normal(0, sqrt(1))
    x ~ Normal(m, sqrt(1))
    y ~ Normal(m, sqrt(1))
end

turing_chain = sample(gdemo(1.5, 2), Turing.NUTS(0.8), 1000; save_state=true, n_adapts=500);
display(turing_chain.info.samplerstate.adaptor)

ahmc_transition, ahmc_state = AbstractMCMC.step(
    Random.default_rng(),
    AbstractMCMC._model(LogDensityFunction(gdemo(1.5, 2))),
    AdvancedHMC.NUTS(0.8); n_adapts = 500
);
display(ahmc_state.adaptor)

Description

Running NUTS with Turing and AdvancedHMC both use adaptors of StanHMCAdaptor, but the fields are different. In particular, the Turing version has a negative value for the endpoint of the mass matrix adaptation window. Consequently, Turing does not perform mass matrix adapatation.

Output of MWE:

 Info: Found initial step size
└   ϵ = 1.6
Sampling 100%|████████████████████████████████████████| Time: 0:00:07
StanHMCAdaptor(
    pc=WelfordVar{Float64} adaptor,
    ssa=NesterovDualAveraging(0.05, 10.0, 0.75, 0.8, 0.7804915138075489),
    init_buffer=75, term_buffer=50, window_size=25,
    state=window(76, -51), window_splits()
)
[ Info: Found initial step size 3.2
StanHMCAdaptor(
    pc=WelfordVar{Float64} adaptor,
    ssa=NesterovDualAveraging(0.05, 10.0, 0.75, 0.8, 7.472207330909228),
    init_buffer=75, term_buffer=50, window_size=25,
    state=window(76, 450), window_splits(100, 150, 250, 450)
)

Julia version info

Julia Version 1.11.6

Manifest

  [80f14c24] AbstractMCMC v5.10.0
  [0bf59076] AdvancedHMC v0.8.3
  [fce5fe82] Turing v0.41.1

[penelopeysm]

Thanks for reporting this @ysfoo! #2718 should fix this.

I'm not a statistician, so can I ask how you found this out?

This bug is really because of the default setting of nadapts=-1 in NUTS which is kind of questionable. I'd like to say that the HMC sampler constructors will be overhauled (they really should be), but I'm not sure I have the capacity to get to it in the short term :/ so will just have to patch things for now.

[ysfoo]

Thanks for addressing this quickly!

It's a weird story as to how I found this out: I had a model where NUTS mixed well when I was using Turing, but did not mix well when I was using AdvancedHMC. NUTS in AdvancedHMC by default prints out some sampler state info during sampling, and I could see the mass matrix being adapted at discrete time points, but remembered that Turing's sampler state was always a diagonal of ones (no adaptation). This then led to me checking whether the two used the same adaptor.

A tangential observation is the fact that unintentionally turning off mass matrix adaptation gives me better MCMC mixing for my particular model which had a weirdly shaped posterior. So after the fix, I would specify metricT=AHMC.UnitEuclideanMetric when constructing the NUTS sampler to turn off mass matrix adaptation intentionally.