Turing model not optimising; unstable ODE solutions


I am getting an integration error with my Turing model. I have searched around but haven’t seen anything similar on discourse.

The ODE model is set up as follows:

using DifferentialEquations, LightGraphs, Random, Turing
using Turing: Variational


const N = 5
const P = 1.0

G = erdos_renyi(N, P)
L = laplacian_matrix(G)

function NetworkFKPP(u, p, t)
    κ, α = p 
    du = -κ * L * u .+ α .* u .* (1 .- u)

u0 = rand(N)
p = [2.0, 3.0]
t_span = (0.0, 2.0)

problem = ODEProblem(NetworkFKPP, u0, t_span, p)
sol = solve(problem, Tsit5(), saveat=0.05)

This seems to work fine.
The Turing model is:

@model function fit(data, func)
    σ ~ InverseGamma(2, 3)
    k ~ truncated(Normal(5,10.0),0.0,10)
    a ~ truncated(Normal(5,10.0),0.0,10)
    u ~ MvNormal(0.5 * ones(5), ones(5))
    p = [k, a] 

    prob = remake(problem, u0=u, p=p)
    predicted = solve(prob, AutoTsit5(Rosenbrock23()), saveat=0.05)

    for i ∈ 1:length(predicted)
        data[:,i] ~ MvNormal(predicted[i], σ)

When I run this using advi, I receive the following warning:

┌ Warning: dt <= dtmin. Aborting. There is either an error in your model specification or the true solution is unstable.
└ @ DiffEqBase ~/.julia/packages/DiffEqBase/PW6zI/src/integrator_interface.jl:343
┌ Info: [ADVI] Should only be seen once: optimizer created for θ
└   objectid(θ) = 0x9425e7b5a1e9ddba

advi continues to run, but the end results are not accurate estimates of true parameters.

This warning does not present if I remove the line

 u ~ MvNormal(0.5 * ones(5), ones(5))

and only optimise the model parameters without updating the initial conditions.

Does anybody know what’s going wrong/have any suggestions on fixing?

Any help appreciated!


I don’t think ADVI is stable for stiff equations.

Ok, thanks. Will this be true regardless of which optimiser I use for ADVI?

I will also try NUTS but I need to scale this problem up to 80-100 nodes, for which I thought ADVI would be better suited.