Integrate SSMProblems & GeneralisedFilters with Turing.jl

[mhauru]

SSMProblems.jl establishes interfaces to define and represent state-space models. By limiting the classes of models that can be represented, SSMProblems.jl provides more specific interface and has the potential to be more computationally efficient.

State-space models are probabilistic models, and major classes of them can be seen as distributions that can be both sampled and scored. Through this lens, we can integrate SSMProblems models through the same interface DynamicPPL.jl currently uses to support "submodels".

to_samplable is implemented in TuringLang/DynamicPPL.jl#696, while to_distribution is tracked at TuringLang/DynamicPPL.jl#802. (Please refer to the comment below for more idea of using these interfaces).

[yebai]

We can likely use to_distribution(m::SSMProblem) to support linear Gaussian state-space models. For more general SSMs, we can look into to_sampleable(m::SSMProblem), which samples the latent states using a sampler in GeneralisedFilters—which can be viewed as a "submodel" but has its inference outsourced to GenerasedFilters instead of relying on MCMC on the tracing data structure VarInfo.

[yebai]

This is already possible:

using GeneralisedFilters, SSMProblems
using Random, StatsFuns, Distributions
using LinearAlgebra
using Turing


function ssm_model(θ::T) where {T<:Real}
    μ0 = T[1.0, 0.0]
    Σ0 = diagm(ones(T, 2))

    A = T[0.8 θ/2; -0.1 0.8]
    Q = Diagonal(T[0.2, 1.0])
    b = zeros(T, 2)

    H = Matrix{T}(I, 1, 2)
    R = Diagonal(T[0.2])
    c = zeros(T, 1)

    return create_homogeneous_linear_gaussian_model(μ0, Σ0, A, b, Q, H, c, R)
end

# data generation process
rng = MersenneTwister(1234)
true_ssm_model = ssm_model(1.0)
_, _, ys = sample(rng, true_ssm_model, 1000)


@model function demo(data::AbstractVector)
    θ ~ Normal()
    state, logZ = GeneralisedFilters.filter(ssm_model(θ), KF(), data)
    Turing.@addlogprob! logZ
    return state
end


chain = sample(demo(ys[1:100]), NUTS(), MCMCSerial(), 3000, 1)

[mhauru]

I just had a call with @charlesknipp where we talked about this. I came to open an issue asking whether SSMProblems could be used as submodels, and found this. (In case someone is confused: The opening post is nominally by me, but @sunxd3 wrote all of it.)

Would it be possible to take something like what @charlesknipp wrote and @yebai quoted above, and turn it into something more like this?

using GeneralisedFilters, SSMProblems
using Turing

function ssm_model(θ::T) where {T<:Real}
    μ0 = T[1.0, 0.0]
    Σ0 = diagm(ones(T, 2))

    A = T[0.8 θ/2; -0.1 0.8]
    Q = Diagonal(T[0.2, 1.0])
    b = zeros(T, 2)

    H = Matrix{T}(I, 1, 2)
    R = Diagonal(T[0.2])
    c = zeros(T, 1)

    return create_homogeneous_linear_gaussian_model(μ0, Σ0, A, b, Q, H, c, R)
end


@model function demo(data::AbstractVector)
    θ ~ Normal()
    data, state ~ to_distribution(ssm_model(θ))
    return state
end

Not sure what the syntax around to_distribution should exactly be. Really conditioning-wise I want data ~ to_distribution(ssm_model(θ)) but I also somehow want to extract the hidden states. I think this is the same problem as discussed in #2485. The implementation of this would be another wrapper type like ReturnedModelWrapper in a DynamicPPL SSMProblems extension.

The sampler would need to be something like Gibbs, where you could say e.g. Gibbs(:θ => NUTS(), :state => KF()), to have the submodel-SSMProblem use KF as the filter. I haven't thought about how to implement that, if we could really just use our current Gibbs or had to come up with something new.

I bet @yebai and maybe @sunxd3 have thought about this way more than I have. Thus the above make sense and would it be easyish to implement?

[yebai]

My view is that we could turn this into a documentation issue.

@model function demo(data::AbstractVector)
    θ ~ Normal()
    state, logZ = GeneralisedFilters.filter(ssm_model(θ), KF(), data)
    Turing.@addlogprob! logZ
    return state
end

The above could be a sensible solution to integration. Some benefits of this approach:

• This would save us the time of writing an extra wrapper by utilising the existing feature Turing.addlogprob!
• It is also consistent with how Turing integrates with HiddenMarkovModels.
• For non-Gaussian state space models, the estimate of logZ is stochastic, so it is helpful to keep Turing.@addlogprob! logZ explicit to inform users what is happening.