svilupp
commented
1 year ago I did a quick benchmark on a simple linear regression model and I don't see much difference in how quickly it adapts (perhaps it's too easy to see any difference)
This is the first 100 tuning draws and what values were sampled.
The hypothesis was that if Nutpie is better, it would hone in on the right values faster, but as you can see it takes only c. 10 draws to get it right for both algorithms -- ie, hard to compare.
# Example for Nuts-rs / Nutpie Adaptor
using AdvancedHMC, ForwardDiff
using LinearAlgebra
using Distributions
using Plots
const A = AdvancedHMC
using LogDensityProblems, TransformVariables, TransformedLogDensities, Parameters
using AbstractMCMC: LogDensityModel
# Example taken from https://www.tamaspapp.eu/DynamicHMCExamples.jl/latest/example_linear_regression/
"""
Linear regression model ``y ∼ Xβ + ϵ``, where ``ϵ ∼ N(0, σ²)`` IID.
Weakly informative prior for `β`, half-T for `σ`.
"""
struct LinearRegressionProblem{TY<:AbstractVector,TX<:AbstractMatrix,
Tν<:Real}
"Observations."
y::TY
"Covariates"
X::TX
"Degrees of freedom for prior."
ν::Tν
end
function (problem::LinearRegressionProblem)(θ)
@unpack y, X, ν = problem # extract the data
@unpack β, σ = θ # works on the named tuple too
ϵ_distribution = Normal(0, σ) # the error term
ℓ_error = mapreduce((y, x) -> logpdf(ϵ_distribution, y - dot(x, β)), +,
y, eachrow(X)) # likelihood for error
ℓ_σ = logpdf(TDist(ν), σ) # prior for σ
ℓ_β = loglikelihood(Normal(0, 10), β) # prior for β
ℓ_error + ℓ_σ + ℓ_β
end
# Random data
n_samples, n_adapts = 2_000, 1_000
N = 100
β = [-10.0, 2.0, -1.0, -0.3, 3, 0.1, 8]
σ = 2
X = hcat(ones(N), randn(N, length(β) - 1));
y = X * β .+ randn(N) .* σ;
p = LinearRegressionProblem(y, X, 1.0);
@info string("log density: ", p((β=β, σ=σ)))
# Transform it to unconstrained space
function problem_transformation(p::LinearRegressionProblem)
as((β=as(Array, size(p.X, 2)), σ=asℝ₊))
end
t = problem_transformation(p)
P = TransformedLogDensity(t, p)
initial_θ = ones(length(β) + 1)
metric = DiagEuclideanMetric(length(initial_θ))
# Sampling w Default
hamiltonian = Hamiltonian(metric, P, ForwardDiff)
initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
integrator = Leapfrog(initial_ϵ)
proposal = NUTS{MultinomialTS,GeneralisedNoUTurn}(integrator)
adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))
# sample(LogDensityModel(P), proposal, metric, adaptor, 100)
@time samples1, stats1 = sample(hamiltonian, proposal, initial_θ, n_samples, adaptor, n_adapts; drop_warmup=false, progress=true);
# NUTPIE
# https://github.com/pymc-devs/nuts-rs/blob/main/src/adapt_strategy.rs
hamiltonian = Hamiltonian(metric, P, ForwardDiff)
initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
integrator = Leapfrog(initial_ϵ)
proposal = NUTS{MultinomialTS,GeneralisedNoUTurn}(integrator)
pc = A.ExpWeightedWelfordVar(size(metric))
adaptor = A.NutpieHMCAdaptor(pc, StepSizeAdaptor(0.8, integrator))
@time samples2, stats2 = sample(hamiltonian, proposal, initial_θ, n_samples, adaptor, n_adapts; drop_warmup=false, progress=true);
# # Plots
# plot how many tuning draws it takes to get to the real values
compare_coefficient = let samples1 = samples1, samples2 = samples2, β = β, σ = σ
function plotter(idx)
pl = plot(getindex.(samples1, idx), xlim=(0, 100), label="Default", color=palette(:default)[1])
plot!(pl, getindex.(samples2, idx), xlim=(0, 100), label="Nutpie/Nuts-rs", color=palette(:default)[2])
if idx <= length(β)
hline!(pl, β[[idx]], color=:red, linestyle=:dash, label="True value")
else # if it's not beta, it must be sigma
hline!(pl, [1 / σ], color=:red, linestyle=:dash, label="True value (untrf σ")
end
return pl
end
end
k = 8
pl = plot([compare_coefficient(i) for i in 1:k]..., layout=(k, 1), size=(600, 400 + k * 100),
plot_title="Comparison of components", legend=:bottomright)
savefig(pl, "comparison.png")
aseyboldt
NOT READY!
Quick draft of the new mass matrix adaption begin tested/used by PyMC team for discussion. Attempt at #311
Notes:
It's a draft -- my VS Code made several formatting changes that I'd need to unwind + I need to take out changes to Project.toml (used for running examples)