VI interface is too low-level

[penelopeysm]

Turing in general has been moving away from making people declare vectorised inputs and getting vectorised outputs. While recognising that this can sometimes be useful, it's also unfortunately quite unhelpful because vectorised formats lose meaning: they rely on internal DynamicPPL details of how variables are organised, they don't convey where values are linked or unlinked, etc.

It is pretty much always more useful to always store samples as VarNamedTuples of raw values, and leave the vectorised samples for real experts who really need to know every detail.

In particular, MCMC sampling and optimisation now refuse to take vectors as initial parameters. Likewise, optimisation does not return a vectorised optimum but rather a VarNamedTuple; if people want to get a vector they have to call a separate function vector_names_and_params(::ModeResult).

Since the MCMC and optimisation interfaces have been updated, the remaining one that still needs some coaxing is VI. In particular, we have that

q, info, state = vi(...)

and q is right now a Bijectors.TransformedDistribution, from which you sample a vector of values. In light of the above changes, I think this is no longer appropriate. q should be a wrapper around the transformed distribution plus a LogDensityFunction, from which you can sample a vector of values, and then feed it into the LDF to get back a VarNamedTuple of values. Calling rand(q) should give you a VarNamedTuple, or possibly a ParamsWithStats, I'm not entirely sure right now.

(Note that the requisite LDF is actually stored as state.prob, so all the info we need is already there.)

This would solve the issue in the docs where "extracting VarNames is a bit verbose at the moment" (https://turinglang.org/docs/tutorials/variational-inference/#obtaining-summary-statistics) (in fact, it's not only verbose -- it is also straight up wrong, for example, it will run into the same LKJCholesky bug that optimisation used to face #2734).

It would also solve a current thorny bit where VI expects that the bijector from constrained to unconstrained space is a constant (and can be obtained via bijector(model)). This works for many models, but is not really correct for models such as

@model function f()
    x ~ Normal()
    y ~ truncated(Normal(); lower=x)
end

because the bijector for y depends on the value of x.

By always transforming to unconstrained space, running VI on the unconstrained LDF and then including the wrapper to convert back to constrained space, we ensure that it behaves correctly even on cases like this.

cc. @Red-Portal

[penelopeysm]

The good news is that this is actually not too much work, because all the machinery to convert between LDF (vectorised samples) and VarNamedTuples (meaningful samples) is already present, and has been well-exercised and well-tested in both MCMC and optimisation now. So we just need some nice data structures plus documentation. I think even Claude would be capable of doing this!

[Red-Portal]

This sounds all very sensible to me. Let me try to work on it. Though my bandwidth these days is really limited due to a change in my status (I have a day job now and not a full-time academic anymore), can't promise fast headway

[penelopeysm]

I didn't mean to ask you to do it, I was going to work on it -- just wanted to let you know. And congratulations! 😄

[Red-Portal]

Oh, I haven't graduated yet. I am taking a leave for personal reasons. But anyway, okay, let me know anytime when you need a review. I am also working on the RJMCMC library in the meantime. I'll let you know when it approaches fruition.

[penelopeysm]

@Red-Portal On #2800 (needs DynamicPPL#1348 too) the interface now looks like this:

using Turing, AbstractMCMC

@model function f()
    x ~ Normal()
    y ~ MvNormal(zeros(4), I)
    1.0 ~ Normal(x + sum(y))
end

result = vi(f(), q_fullrank_gaussian, 1000; show_progress=false)
#= ^ returns a struct:
VIResult
  ├ q    : MvLocationScale
  ├ info : 1000-element Vector{@NamedTuple{elbo::Float64, iteration::Int64}}
  │        final iteration:
  │         ├ elbo = -0.9382047808926819
  │         └ iteration = 1000
  └ (2 more fields: state, ldf)
=#

rand(result)
#= ^ returns something that can be indexed into
VarNamedTuple
├─ x => -0.20039579379888028
└─ y => [0.5965836451905732, -1.4645270348220594, 0.43428159927122983, 1.5212093708261305]
=#

spls = rand(result, 100); # array of VarNamedTuple
AbstractMCMC.from_samples(MCMCChains.Chains, hcat(spls)) # convert to MCMCChains
#= 
Chains MCMC chain (100×5×1 Array{Float64, 3}):

Iterations        = 1:1:100
Number of chains  = 1
Samples per chain = 100
parameters        = x, y[1], y[2], y[3], y[4]
internals         =

Use `describe(chains)` for summary statistics and quantiles.
=#

You can also get it to use static transformations (i.e. pre-construct the equivalent of what used to be bijector(model), and then keep using that, by passing fix_transforms=true to vi. But in general I've found that that's not always helpful: it's often quite trivial to just rederive the transformations.

[Red-Portal]

@penelopeysm Thanks for letting me know. This looks great!

[penelopeysm]

Closed by #2800