Integrating out variables via Laplace approximation (or other methods)

[ElOceanografo]

This is an idea/feature request that I've been thinking about. Relates to this conversation on Discourse, and #1340 . See also this paper. How hard would it be to implement an interface for integrating out variables in a Turing model, using for example the Laplace approximation? This is used in TMB to efficiently fit models with mixed effects or other latent variables/fields, and I think it would be a great addition.

I'm picturing being able to do something like this:

@model function latent_model(x, z)
    a ~ SomeDistribution(1, 2) # parameter
    mu = x * a
    y ~ AnotherDistribution(mu) # latent variable
    z ~ YetAnotherDistribution(y) # likelihood
end

mod = latent_model(xdata,  zdata)
mod_marginal = marginalize(mod, LaplaceApprox(:y))
# ...
chn = sample(mod, sampler, nsamples)
opt = optimize(mod_marginal, MLE()) # fit `a` while integrating out `y`

Does this seem in principle desirable/doable? If so I'd be willing to do some work on it, given a couple pointers on where to get started...

[devmotion]

Also related to https://github.com/TuringLang/Turing.jl/issues/976

[devmotion]

And the proposal in TuringLang/DynamicPPL.jl#94

[ElOceanografo]

Or, alternatively, the interface could be something like this, more along the lines of how Gibbs works for samplers:

scheme = MarginalOptimizer(LBFGS(:a), LaplaceApprox(:y))
opt = optimize(mod, MLE(), scheme)

[yebai]

This looks like quite hard to do well for general cases, so it might make sense to implement it in a separate library that depends on DynamicPPL/Turing.