Likelihood model in Bayesian NODE

Specific Domains Modelling & Simulations
1

Hi all, I’m reading Uncertainty Quantified Deep Bayesian Model Discovery · Overview of Julia's SciML. to play around with Bayesian NODEs. I noticed the loss function which actually gets passed to AdvancedHMC’s samplers is l(θ) = -sum(abs2, ode_data .- predict_neuralode(θ)) - sum(θ .* θ), where the “likelihood” part is proportional to a Gaussian likelihood model with a fixed standard deviation.

My question is this one, say I want to use Normal from Distributions.jl and simultaneously fit the standard deviation parameter during fitting/sampling; are there examples of such use? Are there any pitfalls I should be cautious of? (e.g. different scale of parameters, etc)

2

you can use that. Yeah the loglikelihoods from Distributions.jl can be helpful in these contexts

3

I hacked together a solution here. Not sure if it’s the recommended way to go but ComponentArrays seemed fine with it.

Modifies Uncertainty Quantified Deep Bayesian Model Discovery · Overview of Julia's SciML to simultaneously fit standard deviation (sigma) parameter of the Gaussian likelihood. Also modifies objective so that the priors for sigma and neural network weights use proper Distributions.jl objects, as well as the likelihood function. Minimizes the neg log posterior, so doing MAP estimation.

# original problem: https://docs.sciml.ai/Overview/dev/showcase/bayesian_neural_ode/
# this uses ADAM + LBGFS, also simultaneously optimizes the σ for Gaussian likelihood model

# SciML Libraries
import SciMLSensitivity as SMS
import OrdinaryDiffEq as ODE
import SciMLBase

# likelihood and prior
import Distributions

# ML Tools
import Lux
import Zygote

# Optimization
import Optimization
import OptimizationOptimisers
import OptimizationOptimJL

# External Tools
import Random
import Plots
import Printf
import ComponentArrays

using Plots

# 1: setup: get data from spiral ODE
u0 = [2.0; 0.0]
datasize = 40
tspan = (0.0, 1)
tsteps = range(tspan[1], tspan[2], length = datasize)
function trueODEfunc(du, u, p, t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u .^ 3)'true_A)'
end

prob_trueode = ODE.ODEProblem(trueODEfunc, u0, tspan)
ode_solve = ODE.solve(prob_trueode, ODE.Tsit5(), saveat = tsteps)
ode_data = Array(ode_solve)
ode_data .+= rand(Distributions.Normal(0, 0.5), size(ode_data))
plot(ode_solve, c=[1 2])
scatter!(tsteps, transpose(ode_data), legend=false, c=[1 2])

# 2: define neural ODE architecture
dudt2 = Lux.Chain(x -> x .^ 3,
    Lux.Dense(2, 35, tanh),
    Lux.Dense(35, 2))

rng = Random.default_rng()
p, st = Lux.setup(rng, dudt2)
const _st = st
function neuralodefunc(u, p, t)
    dudt2(u, p, _st)[1]
end
function prob_neuralode(u0, p)
    prob = ODE.ODEProblem(neuralodefunc, u0, tspan, p)
    ODE.solve(prob, ODE.Tsit5(), saveat = tsteps)
end
p = ComponentArrays.ComponentArray{Float64}(p)
const _p = p

# here's how to simultaneously optimize NN weights and likelihood model params, pack them into
# a nested ComponentArray.
# note that σ is optimized on log (unconstrained scale), take exp to go to constrained scale
# in the likelihood computation
θ0 = ComponentArrays.ComponentArray(nn = p, logσ = log(1))
const _θ0 = θ0

# these are just helpers to visualize the solution and see the neg log likelihood of the optimal parameters
function predict_neuralode(p)
    p = p isa ComponentArrays.ComponentArray ? p : convert(typeof(_p), p)
    Array(prob_neuralode(u0, p))
end
function nll_neuralode(θ)
    θ = θ isa ComponentArrays.ComponentArray ? θ : convert(typeof(_θ0), θ)
    pred = predict_neuralode(θ.nn)
    -sum(Distributions.logpdf.(Distributions.Normal.(pred, exp(θ.logσ)), ode_data))
end

# 3: objective fn for the optimizers
# Returns negative log posterior, by calculating the log prior and log likelihood seperately,
# then returning the sum (with negative sign)
# σ now gets its own prior term.
function objective(θ, _)
    σ = exp(θ.logσ)
    pred = predict_neuralode(θ.nn)
    lprior = Distributions.loglikelihood(Distributions.Normal(), θ.nn) + # nn weights prior
             Distributions.logpdf(Distributions.Cauchy(), θ.logσ) # σ prior
    llike = sum(Distributions.logpdf.(Distributions.Normal.(pred, σ), ode_data))
    return -(llike + lprior)
end

# functor to stop optimization when improvement in neg log posterior is below some rtol.
# also handle printing of info.
Base.@kwdef mutable struct ConvergeCallback
    rtol::Float64 = 1.0e-6      # relative change counted as no progress
    every::Int = 25             # print interval
    prev::Float64 = Inf         # previous objective value
end

function (cb::ConvergeCallback)(state, l)
    if state.iter == 1 || state.iter % cb.every == 0
        Printf.@printf "Iteration: %5d, Objective: %.6e\n" state.iter l
    end
    delta = abs(cb.prev - l) / max(abs(cb.prev), eps())
    cb.prev = l
    if delta < cb.rtol
        Printf.@printf "Convergence reached, exiting: Iteration: %5d, Objective: %.6e\n" state.iter l
        return true
    else
        return false
    end
end

optf = Optimization.OptimizationFunction(objective, Optimization.AutoZygote())
optprob = Optimization.OptimizationProblem(optf, θ0)

# Adam first: robust to a poor starting point. Each solve gets its own ConvergeCallback so the
# previous value starts clean.
res_adam = Optimization.solve(optprob, OptimizationOptimisers.Adam(0.05);
    callback = ConvergeCallback(every=10), maxiters = 300)

# then L-BFGS from Adam's answer for fast local convergence
optprob2 = SciMLBase.remake(optprob; u0 = res_adam.u)
res_lbfgs = Optimization.solve(optprob2, OptimizationOptimJL.LBFGS();
    callback = ConvergeCallback(every=10), maxiters = 200)

θ_map = res_lbfgs.u
prediction = predict_neuralode(θ_map.nn)
rmse = sqrt(sum(abs2, ode_data .- prediction) / length(prediction))
Printf.@printf "Objective: %.6e | NLL only: %.6e\n" objective(θ_map, nothing) nll_neuralode(θ_map)
Printf.@printf "σ: %.6f | RMSE: %.6f\n" exp(θ_map.logσ) rmse

# 4: plot the fit

pl = Plots.scatter(tsteps, ode_data[1, :], color = :red, label = "Data: Var1", xlabel = "t",
    title = "Spiral Neural ODE (MAP)")
Plots.scatter!(tsteps, ode_data[2, :], color = :blue, label = "Data: Var2")
Plots.plot!(tsteps, prediction[1, :], color = :red, w = 2, label = "MAP: Var1")
Plots.plot!(tsteps, prediction[2, :], color = :blue, w = 2, label = "MAP: Var2",
    ylims = (-2.5, 3.5))

# phase-space view
pl2 = Plots.scatter(ode_data[1, :], ode_data[2, :], color = :red, label = "Data",
    xlabel = "Var1", ylabel = "Var2", title = "Spiral Neural ODE (MAP)")
Plots.plot!(prediction[1, :], prediction[2, :], color = :black, w = 2,
    label = "MAP fit", ylims = (-2.5, 3))