Support transformed RVs & track expressions

[trappmartin]

This is a summary issue of something floating around for quite some time.

We would like to be able to do this:

@model test(x) = begin
    m1 ~ Normal()
    @track m2 = f(m1) # some differentiable transformation
    x ~ Normal(m2, 1.0)
end

The resulting chain should contain states for m1, m2 respectively.

Comments by @mohamed82008 how to approach this:

• Define a Deterministic distribution.
• Make sure getspace doesn't return any variable symbol whose dists field isa Vector{<:Deterministic}. This is the trickiest part.
• Define logpdf for Deterministic to return 0, rand to return the wrapped value and bijector to return Identity.
• Define assume for Deterministic such that it writes down the value wrapped in the RHS to vi[vn] calling ForwardDiff.value, ReverseDiff.value or Tracker.data on the wrapped value before writing it to vi[vn]. We should probably define DistributionsAD.value to unify these.

[devmotion]

Even if that might be a way to solve the problem, I'm not sure that's how we should actually do it (sorry 😛). Adding a deterministic distribution and an identity bijector for a transformation seems like a bit of a hack to me (and the AD part seems to be focused on HMC?).

It seems the main purpose of @track would be to ensure that m2 would be part of the chain? Or is there anything else that we have to take care of?

If it's (mainly) about the interesting variables ending up in the chain, maybe

@model function test(x)
    m1 ~ Normal()
    m2 = f(m1) # some differentiable transformation
    x ~ Normal(m2, 1.0)
	return m1, m2
end

and respecting the return statement would be the more natural approach.

[trappmartin]

Cool idea to have the return statement determine what is stored "additionally" in the chain.

[trappmartin]

That said, I find the use of an @track macro more "user-friendly". But that's rather subjective. Also, @track could be used to inform a sampler that performs marginalisation automatically, that a certain RV should not be marginalised out as we want to track it.

[trappmartin]

What do you guys think about using the return statement (as described by @devmotion) in general to define the set of variables that should be tracked? This idea is kinda growing on me. :)

@model function gmm(x, K)
    m ~ filldist(Normal(), K) 
    v ~ filldist(InverseGamma(2, 3), K)
    s = sqrt.(v)
    w ~ Dirichlet(K, 1.0)
    z ~ filldist(Categorical(w), length(x))
    x .~ Normal.(m[z] s[z])
    return m, s, w
end

Then the resulting chain has only states for m, s and w but not for z or v and we can safely marginalise out z in the inference.

[torfjelde]

So just a heads up: I've already implemented a generated_quantities for "extracting" return-values from a chain by re-running the model on the values in the chain: https://github.com/cambridge-mlg/Covid19/blob/4d967d58eeada3f3399116fc815cb17fd4163c87/src/utils.jl#L79-L89.

In this case you could then just do:

@model function gmm(x, K)
    m ~ filldist(Normal(), K) 
    v ~ filldist(InverseGamma(2, 3), K)
    s = sqrt.(v)
    w ~ Dirichlet(K, 1.0)
    z ~ filldist(Categorical(w), length(x))
    x .~ Normal.(m[z] s[z])
    return m, s, w
end

m = gmm(x, K)
chain = sample(m, sampler, n)

return_values = generated_quantities(m, chain)

This is a fairly good temporary-fix for tackling this issue + it has additional use-cases:

1. Predicting. model_train is instantiated on "train-data" but now you want to generate the return-values for model_test which is instantiated on "test-data", then you can execute model_test using the chain (posterior) obtained from model_train.
2. If you for some reason wanted to compute the "generated values" for model1 using a chain obtained from model2 you could do that, assuming they have the same latent parameters.

Not suggesting this is as a solution to the problem, but it can in many cases be a good fix + it provides some additional functionality.

For completeness:

function generated_quantities(m::Turing.Model, c::MCMCChains.Chains)
    # if `c` is multiple chains we pool them into a single chain
    chain = length(chains(c)) == 1 ? c : MCMCChains.pool_chain(c)

    varinfo = Turing.DynamicPPL.VarInfo(m)

    return map(1:length(chain)) do i
        Turing.DynamicPPL._setval!(varinfo, chain[i])
        m(varinfo)
    end
end

[devmotion]

Nice!

IMO both approaches try to address a different problem (which we should both solve in some way). The approach with return (or some other similar model-based approach) tries to avoid saving auxiliary variables completely. This structure would allow samplers to marginalize out variables, which is otherwise not possible if always all variables have to be returned. Your implementation addresses the problem of rerunning the model with different (or the same) variables, which is interesting in itself and provides the tracked expressions if they are returned in the model definition (so you need the same/similar model definition).

Additionally, I think we should just have some convenience function for running models with NamedTuple instead of VarInfo, instead of letting users rewrite the same hacks over and over again. IMO it would be good to use NamedTuples in generated_quantities as well since otherwise reordering of the columns will break the function, which would not be clear to a user.

[torfjelde]

Additionally, I think we should just have some convenience function for running models with NamedTuple instead of VarInfo, instead of letting users rewrite the same hacks over and over again.

👍

IMO it would be good to use NamedTuples in generated_quantities as well since otherwise reordering of the columns will break the function, which would not be clear to a user.

Hooold up; what do you mean when you say "reordering of the columns"? Do you mean that generated_quantities will break if the ordering of the resulting vector in the model is wrong? If that's what you mean, _setval internally makes use of the symbols so everything should work just fine even then, right?
With that being said, I 100% support allowing use of NamedTuple instead of having to deal with VarInfo 👍

[devmotion]

You're absolutely right - I assumed _setval! just operates on vectors, I guess due to the new indexing behaviour in MCMCChains#dev (IIRC chains[i] just returns an array (with labelled axes)).

[ElOceanografo]

+1 for using the return statement to define which variables to track. I think it's clean and intuitive, and the ability to marginalize out variables without storing them in a chain first would be really great.

[cpfiffer]

Something I've always wondered about this is why we don't just track these generated quantities in VarInfo? It seems like it would be reasonably easy to build a VarName from the returned quantities and just stick them in some field somewhere.

[trappmartin]

I think using the return statement is indeed the best idea. I'm not sure we need to store transformed variables in VarInfo at all, tbh. I'd rather have it only in the resulting trace and keep VarInfo clean.