# Likelihood model in Bayesian NODE

**URL:** <https://discourse.julialang.org/t/likelihood-model-in-bayesian-node/138664>\
**Category:** Modelling & Simulations\
**Created:** [August 7, 2026, 4:22pm UTC](https://discourse.julialang.org/t/likelihood-model-in-bayesian-node/138664 "2026-08-07T16:22:36Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![slwu89](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/slwu89/32/217323_2.png) [@slwu89](https://discourse.julialang.org/u/slwu89)\
**Post date:** [August 7, 2026, 4:22pm UTC](https://discourse.julialang.org/t/likelihood-model-in-bayesian-node/138664/1 "2026-08-07T16:22:36Z")

</div>

Hi all, I’m reading [Uncertainty Quantified Deep Bayesian Model Discovery · Overview of Julia's SciML](https://docs.sciml.ai/Overview/stable/showcase/bayesian_neural_ode/#Step-3:-Define-the-loss-function-for-the-Neural-ODE). 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)

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 8, 2026, 2:16pm UTC](https://discourse.julialang.org/t/likelihood-model-in-bayesian-node/138664/2 "2026-08-08T14:16:42Z")

</div>

> [@slwu89](#):
>
> My question is this one, say I want to use `Normal` from [Distributions.jl](https://juliaregistries.github.io/General/packages/redirect_to_repo/Distributions) 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)

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

---

<div class="post-metadata">

**Author:** ![slwu89](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/slwu89/32/217323_2.png) [@slwu89](https://discourse.julialang.org/u/slwu89)\
**Post date:** [August 10, 2026, 10:39pm UTC](https://discourse.julialang.org/t/likelihood-model-in-bayesian-node/138664/3 "2026-08-10T22:39:21Z")

</div>

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](https://docs.sciml.ai/Overview/dev/showcase/bayesian_neural_ode/) 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.

```julia-auto
# 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))

```