AdvancedHMC.sample might not be respecting the passed rng object

[sefffal]

Hi all, thanks for the great package!

I noticed recently with some Octofitter tests that passing a seeded RNG object (e.g. rng = Xoshiro(0) into AdvancedHMC.sample(...;rng=rng) does not give reproducible results.

By contrast, calling Random.seed!(0) (with or without passing a seeded rng object) does lead to reproducible results.

Could it be that there is a rand() call somewhere in this package that is not respected the passed rng object?

Thanks very much, and let me know if there is any way I can help with tracking this down.

-William

[penelopeysm]

Hello @sefffal, thanks for getting in touch!

I dug into it a bit with the first example in the README. I don't know if this is similar to what you were doing (if not, it would be great if you could share an MWE).

using AdvancedHMC, ForwardDiff
using LogDensityProblems
using LinearAlgebra
using Random: Xoshiro
using Statistics: mean
struct LogTargetDensity
    dim::Int
end
LogDensityProblems.logdensity(p::LogTargetDensity, θ) = -sum(abs2, θ) / 2  # standard multivariate normal
LogDensityProblems.dimension(p::LogTargetDensity) = p.dim
LogDensityProblems.capabilities(::Type{LogTargetDensity}) = LogDensityProblems.LogDensityOrder{0}()
D = 10; initial_θ = rand(Xoshiro(0), D)
ℓπ = LogTargetDensity(D)
n_samples, n_adapts = 2_000, 1_000
metric = DiagEuclideanMetric(D)
hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff)
initial_ϵ = find_good_stepsize(Xoshiro(0), hamiltonian, initial_θ)
integrator = Leapfrog(initial_ϵ)
kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))

adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))
m = mean(mean(sample(Xoshiro(0), hamiltonian, kernel, initial_θ, n_samples, adaptor, n_adapts; progress=true)[1]))

I can see that (perhaps rather unfortunately) the call to sample() mutates adaptor; so if you just run sample() a second time it will result in different values. However, if you redefine a fresh adaptor before rerunning sample, you should get a reproducible result.

julia> adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))
StanHMCAdaptor(
    pc=WelfordVar,
    ssa=NesterovDualAveraging(γ=0.05, t_0=10.0, κ=0.75, δ=0.8, state.ϵ=1.6),
    init_buffer=75, term_buffer=50, window_size=25,
    state=window(0, 0), window_splits()
)

julia> m = mean(mean(sample(
           Xoshiro(0), hamiltonian, kernel, initial_θ,
               n_samples, adaptor, n_adapts; progress=false)[1]))
┌ Info: Finished 1000 adapation steps
│   adaptor =
│    StanHMCAdaptor(
│        pc=WelfordVar,
│        ssa=NesterovDualAveraging(γ=0.05, t_0=10.0, κ=0.75, δ=0.8, state.ϵ=0.6956947494412146),
│        init_buffer=75, term_buffer=50, window_size=25,
│        state=window(76, 950), window_splits(100, 150, 250, 450, 950)
│    )
│   κ.τ.integrator = Leapfrog(ϵ=0.696)
└   h.metric = DiagEuclideanMetric([1.064418895706038, 0.97647 ...])
┌ Info: Finished 2000 sampling steps for 1 chains in 0.053473125 (s)
│   h = Hamiltonian(metric=DiagEuclideanMetric([1.064418895706038, 0.97647 ...]), kinetic=GaussianKinetic())
│   κ = HMCKernel{AdvancedHMC.FullMomentumRefreshment, Trajectory{MultinomialTS, Leapfrog{Float64}, GeneralisedNoUTurn{Float64}}}(AdvancedHMC.FullMomentumRefreshment(), Trajectory{MultinomialTS}(integrator=Leapfrog(ϵ=0.696), tc=GeneralisedNoUTurn{Float64}(10, 1000.0)))
│   EBFMI_est = 1.0883690820849576
└   average_acceptance_rate = 0.8378527099573518
-0.004069571743129061

julia> adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))
StanHMCAdaptor(
    pc=WelfordVar,
    ssa=NesterovDualAveraging(γ=0.05, t_0=10.0, κ=0.75, δ=0.8, state.ϵ=1.6),
    init_buffer=75, term_buffer=50, window_size=25,
    state=window(0, 0), window_splits()
)

julia> m = mean(mean(sample(
           Xoshiro(0), hamiltonian, kernel, initial_θ,
               n_samples, adaptor, n_adapts; progress=false)[1]))
┌ Info: Finished 1000 adapation steps
│   adaptor =
│    StanHMCAdaptor(
│        pc=WelfordVar,
│        ssa=NesterovDualAveraging(γ=0.05, t_0=10.0, κ=0.75, δ=0.8, state.ϵ=0.6956947494412146),
│        init_buffer=75, term_buffer=50, window_size=25,
│        state=window(76, 950), window_splits(100, 150, 250, 450, 950)
│    )
│   κ.τ.integrator = Leapfrog(ϵ=0.696)
└   h.metric = DiagEuclideanMetric([1.064418895706038, 0.97647 ...])
┌ Info: Finished 2000 sampling steps for 1 chains in 0.0427485 (s)
│   h = Hamiltonian(metric=DiagEuclideanMetric([1.064418895706038, 0.97647 ...]), kinetic=GaussianKinetic())
│   κ = HMCKernel{AdvancedHMC.FullMomentumRefreshment, Trajectory{MultinomialTS, Leapfrog{Float64}, GeneralisedNoUTurn{Float64}}}(AdvancedHMC.FullMomentumRefreshment(), Trajectory{MultinomialTS}(integrator=Leapfrog(ϵ=0.696), tc=GeneralisedNoUTurn{Float64}(10, 1000.0)))
│   EBFMI_est = 1.0883690820849576
└   average_acceptance_rate = 0.8378527099573518
-0.004069571743129061

I'm running:

julia> versioninfo()
Julia Version 1.11.1
Commit 8f5b7ca12ad (2024-10-16 10:53 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin22.4.0)
  CPU: 10 × Apple M1 Pro
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, apple-m1)
Threads: 1 default, 0 interactive, 1 GC (on 8 virtual cores)

(ppl) pkg> st
Status `~/ppl/Project.toml`
  [0bf59076] AdvancedHMC v0.6.3
  [f6369f11] ForwardDiff v0.10.37
  [6fdf6af0] LogDensityProblems v2.1.2
  [10745b16] Statistics v1.11.1
  [37e2e46d] LinearAlgebra v1.11.0
  [9a3f8284] Random v1.11.0

[penelopeysm]

I recognise that the above doesn't quite account for your observation that Random.seed! makes it work reproducibly (unless, perhaps, you happened to place seed! at the top of a script and reran multiple lines including the definition of adaptor). If the above isn't relevant, do let us know.

[sefffal]

Thanks @penelopeysm for looking into this!
I'll see if I can narrow down the cause on my end, or reduce it to a smaller example.

One thing I might try is replacing the global RNG with one that logs when rand is called.

[torfjelde]

into AdvancedHMC.sample(...;rng=rng) does not give reproducible results.

Hmm, AFAIK, this is not the syntax for using the rng. The rng should be passed as the first argument, e.g.

sample(rng, ...)

Did you try this?