Automated construction of diffeomorphism/bijector trace translators

[bgroenks96]

One feature of Turing (and similar PPLs) that I am sorely missing in Gen is the ability to automatically reparameterize a probabilistic program with one or more variables (or I guess, "choices" in the language of Gen) that have bounded support under a transformed space with unbounded support. This means automatically applying transforms such as log and logit for the cases of single and double bounded support respectively. As far as I can tell, it's possible to do this manually for any particular program with @transform, but there does not appear to be any kind of automatic transformation based on the support of the prior distributions.

My question is, therefore, how hard would it be to implement a function (or macro?) that would automatically generate a trace translator alongside the generative function which applies such a reparameterization? Being somewhat new to Gen and its internals, it's a bit hard for me to judge this. I can see where it might be tricky, especially for dynamic generative functions.

I would be curious to hear the opinions of the primary maintainers on what it would take to implement this. I would imagine that this has been discussed before.

[ztangent]

Hi @bgroenks96, I discussed this briefly with @alex-lew, and he suggested the following alternative solution that makes use of the @dist DSL to define custom distributions:

1. Defining a @dist on the transformed space, e.g. @dist loggamma(…) = log(gamma(…))
2. Defining a generative function @gen gamma(…) = exp({:loggamma} ~ loggamma(…))

Then use the generative function to define your model. The trace’s variables are in an unconstrained space, but your program looks to you like you are working in the constrained space, and given a trace, you can still extract the gamma value as tr[addr] (it’s the return value of the gamma call).

This way, it should be possible to use unconstrained proposal distributions that target the value of the unconstrained random variable, without having to use trace translators as well. I believe gradient-based moves like MALA and HMC should also work, (though it's possible that gradient-support in the @dist DSL isn't fully debugged).

Of course, the use of a wrapper @dist and @gen function adds a bit of overhead, but it should also be possible to write a custom generative function that does the same thing with less overhead.

Hope this helps!

[bgroenks96]

Ok, thanks. That sounds like a workable solution for many cases.

However, I think it is a bit unfortunate to offload the burden onto the user of writing their generative function already reparameterized in this way. In my opinion, a (really good) PPL should let you define the model as close as possible to its formal specification and then automatically apply such reparameterizations, which are often only really necessary for algorithmic purposes, with minimal effort from the user.

Do you think it would be at all possible to design such a scheme? Or is the necessary information lost when the generative function is constructed?