Turing.jl: Help with intializing the `getq` function from the tutorials

[ClaudMor]

Hello,

I'm trying to optimize the inference of a DifferentialEquations.jl complex model using a Turing model. I tried to implement the suggestions of these three issues:

1. Type inference in Turing model when using vectorized ~
2. :forwarddiff backend with type specification goes OOM
3. Inference using :reversediff, :tracker and :zygote performs worse than :forwarddiff

And now it calibrates properly using Turing.NUTS(), anyway, I also wanted to try ADVI. The standard interface works fine, I wondered if I could get more ( like parameter values initialization or simply performance) from the more advanced VI interfaces :

vi(model, alg::VariationalInference, getq::Function, θ::AbstractArray)

The problem is that I'm not experienced enough to write the correct getq function for the Turing model:

# with for loops
@model function turing_model_with_for(data::Array{Float64,2}, prob1, priors, ::Type{T} = Float64) where {T} #::Array{Distribution{Univariate,Continuous},1}
    σ ~ InverseGamma(2, 3)


    variable_parameters_extractions ~ arraydist(priors)

    predicted  = solve(prob1, Tsit5(), p = variable_parameters_extractions, saveat=0.1).u  #

    if length(predicted) == size(data, 2)
        for i = 1:length(predicted)
            data[:,i] ~ MvNormal(predicted[i], σ) 
        end
    else
        Turing.@addlogprob! -Inf
        return
    end

end

I tried to use the getq function from the tutorials

function getq(θ)
   # print(typeof(θ))
    d = length(θ) ÷ 2
    A = @inbounds θ[1:d]
    b = @inbounds θ[d + 1: 2 * d]
    b = to_constrained ∘ Shift(b; dim = Val(1)) ∘ Scale(exp.(A); dim = Val(1))
    return transformed(base_dist, b)
end

where

• d is the number of parameters in turing model
• base_dist = Turing.DistributionsAD.TuringDiagMvNormal(zeros(d), ones(d))
• to_constrained = inv(bijector(model));

, together with the initialization

θ = vcat([repeat([mean(InverseGamma(2,3))],d)...,mean(InverseGamma(2,3)), variable_parameters_values ]...)

where variable_parameters_values are the parameter values found with a previous, successful calibration with Optim ( BFGS ).

Is the initialization θ written in the correct form?

The reasoning behind it was: I have to initialize A and b, where b is the Shift ( so I translate using the mean), and A the Scale of the MvNormal, so I use the same std for all.

Thank you very much

[pitmonticone]

Do you have any update? I would be very interested to know.