Transforming variables

Motivation

In a probabilistic programming language (PPL) such as DynamicPPL.jl, one crucial functionality for enabling a large number of inference algorithms to be implemented, in particular gradient-based ones, is the ability to work with "unconstrained" variables.

For example, consider the following model:

@model function demo()
    s ~ InverseGamma(2, 3)
    return m ~ Normal(0, √s)
end

Here we have two variables s and m, where s is constrained to be positive, while m can be any real number.

For certain inference methods, it's necessary / much more convenient to work with an equivalent model to demo but where all the variables can take any real values (they're "unconstrained").

For a large family of constraints encountered in practice, it is indeed possible to transform a (partially) constrained model to a completely unconstrained one in such a way that sampling in the unconstrained space is equivalent to sampling in the constrained space.

In DynamicPPL.jl, this is often referred to as linking (a term originating in the statistics literature) and is done using transformations from Bijectors.jl.

For example, the above model could be transformed into (the following pseudo-code; it's not working code):

@model function demo()
    log_s ~ log(InverseGamma(2, 3))
    s = exp(log_s)
    return m ~ Normal(0, √s)
end

Here log_s is an unconstrained variable, and s is a constrained variable that is a deterministic function of log_s.

But to ensure that we stay consistent with what the user expects, DynamicPPL.jl does not actually transform the model as above, but instead makes use of transformed variables internally to achieve the same effect, when desired.

In the end, we'll end up with something that looks like this:

Below we'll see how this is done.

What do we need?

There are two aspects to transforming from the internal representation of a variable in a varinfo to the representation wanted in the model:

1. Different implementations of AbstractVarInfo represent realizations of a model in different ways internally, so we need to transform from this internal representation to the desired representation in the model. For example,

   • VarInfo represents a realization of a model as a "flattened" / vector representation, regardless of the form of the variable in the model.
   • SimpleVarInfo represents a realization of a model exactly as in the model (unless it has been transformed; we'll get to that later).

2. We need the ability to transform from "constrained space" to "unconstrained space", as we saw in the previous section.

Working example

A good and non-trivial example to keep in mind throughout is the following model:

using DynamicPPL, Distributions
@model demo_lkj() = x ~ LKJCholesky(2, 1.0)

demo_lkj (generic function with 2 methods)

LKJCholesky is a LKJ(2, 1.0) distribution, a distribution over correlation matrices (covariance matrices but with unit diagonal), but working directly with the Cholesky factorization of the correlation matrix rather than the correlation matrix itself (this is more numerically stable and computationally efficient).

As mentioned, some implementations of AbstractVarInfo, e.g. VarInfo, works with a "flattened" / vector representation of a variable, and so in this case we need two transformations:

1. From the Cholesky object to a vector representation.
2. From the Cholesky object to an "unconstrained" / linked vector representation.

And similarly, we'll need the inverses of these transformations.

From internal representation to model representation

To go from the internal variable representation of an AbstractVarInfo to the variable representation wanted in the model, e.g. from a Vector{Float64} to Cholesky in the case of VarInfo in demo_lkj, we have the following methods:

to_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])

from_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])

These methods allow us to extract the internal-to-model transformation function depending on the varinfo, the variable, and the distribution of the variable:

• varinfo + vn defines the internal representation of the variable.
• dist defines the representation expected within the model scope.

Continuing from the example above, we can inspect the internal representation of x in demo_lkj with VarInfo using DynamicPPL.getindex_internal:

model = demo_lkj()
varinfo = VarInfo(model)
x_internal = DynamicPPL.getindex_internal(varinfo, @varname(x))

4-element Vector{Float64}:
 1.0
 0.5681958064382993
 0.0
 0.8228933865002992

f_from_internal = DynamicPPL.from_internal_transform(
    varinfo, @varname(x), LKJCholesky(2, 1.0)
)
f_from_internal(x_internal)

LinearAlgebra.Cholesky{Float64, Matrix{Float64}}
L factor:
2×2 LinearAlgebra.LowerTriangular{Float64, Matrix{Float64}}:
 1.0        ⋅ 
 0.568196  0.822893

Let's confirm that this is the same as varinfo[@varname(x)]:

x_model = varinfo[@varname(x)]

LinearAlgebra.Cholesky{Float64, Matrix{Float64}}
L factor:
2×2 LinearAlgebra.LowerTriangular{Float64, Matrix{Float64}}:
 1.0        ⋅ 
 0.568196  0.822893

Similarly, we can go from the model representation to the internal representation:

f_to_internal = DynamicPPL.to_internal_transform(varinfo, @varname(x), LKJCholesky(2, 1.0))

f_to_internal(x_model)

4-element reshape(::LinearAlgebra.LowerTriangular{Float64, Matrix{Float64}}, 4) with eltype Float64:
 1.0
 0.5681958064382993
 0.0
 0.8228933865002992

It's also useful to see how this is done in SimpleVarInfo:

simple_varinfo = SimpleVarInfo(varinfo)
DynamicPPL.getindex_internal(simple_varinfo, @varname(x))

LinearAlgebra.Cholesky{Float64, Matrix{Float64}}
L factor:
2×2 LinearAlgebra.LowerTriangular{Float64, Matrix{Float64}}:
 1.0        ⋅ 
 0.568196  0.822893

Here see that the internal representation is exactly the same as the model representation, and so we'd expect from_internal_transform to be the identity function:

DynamicPPL.from_internal_transform(simple_varinfo, @varname(x), LKJCholesky(2, 1.0))

identity (generic function with 1 method)

Great!

From unconstrained internal representation to model representation

In addition to going from internal representation to model representation of a variable, we also need to be able to go from the unconstrained internal representation to the model representation.

For this, we have the following methods:

to_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])

from_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])

These are very similar to DynamicPPL.to_internal_transform and DynamicPPL.from_internal_transform, but here the internal representation is also linked / "unconstrained".

Continuing from the example above:

f_to_linked_internal = DynamicPPL.to_linked_internal_transform(
    varinfo, @varname(x), LKJCholesky(2, 1.0)
)

x_linked_internal = f_to_linked_internal(x_model)

1-element Vector{Float64}:
 0.6448544199102265

f_from_linked_internal = DynamicPPL.from_linked_internal_transform(
    varinfo, @varname(x), LKJCholesky(2, 1.0)
)

f_from_linked_internal(x_linked_internal)

LinearAlgebra.Cholesky{Float64, Matrix{Float64}}
L factor:
2×2 LinearAlgebra.LowerTriangular{Float64, Matrix{Float64}}:
 1.0        ⋅ 
 0.568196  0.822893

Here we see a significant difference between the linked representation and the non-linked representation: the linked representation is only of length 1, whereas the non-linked representation is of length 4. This is because we actually only need a single element to represent a 2x2 correlation matrix, as the diagonal elements are always 1 and it's symmetric.

We can also inspect the transforms themselves:

f_from_internal

DynamicPPL.ToChol('L') ∘ DynamicPPL.FromVec{Tuple{Int64, Int64}}((2, 2))

vs.

f_from_linked_internal

Bijectors.Inverse{Bijectors.VecCholeskyBijector}(Bijectors.VecCholeskyBijector(:L))

Here we see that f_from_linked_internal is a single function taking us directly from the linked representation to the model representation, whereas f_from_internal is a composition of a few functions: one reshaping the underlying length 4 array into 2x2 matrix, and the other converting this matrix into a Cholesky, as required to be compatible with LKJCholesky(2, 1.0).

Why do we need both to_internal_transform and to_linked_internal_transform?

One might wonder why we need both to_internal_transform and to_linked_internal_transform instead of just a single to_internal_transform which returns the "standard" internal representation if the variable is not linked / "unconstrained" and the linked / "unconstrained" internal representation if it is.

That is, why can't we just do

Unfortunately, this is not possible in general. Consider for example the following model:

@model function demo_dynamic_constraint()
    m ~ Normal()
    x ~ truncated(Normal(); lower=m)

    return (m=m, x=x)
end

demo_dynamic_constraint (generic function with 2 methods)

Here the variable x is constrained to be in the domain (m, Inf), where m is sampled according to a Normal.

model = demo_dynamic_constraint()
varinfo = VarInfo(model)
varinfo[@varname(m)], varinfo[@varname(x)]

(0.35945980388141124, 0.6054721554645925)

We see that the realization of x is indeed greater than m, as expected.

But what if we link this varinfo so that we end up working on an "unconstrained" space, i.e. both m and x can take on any values in (-Inf, Inf):

varinfo_linked = link(varinfo, model)
varinfo_linked[@varname(m)], varinfo_linked[@varname(x)]

(0.35945980388141124, 0.6054721554645925)

Still get the same values, as expected, since internally varinfo transforms from the linked internal representation to the model representation.

But what if we change the value of m, to, say, a bit larger than x?

# Update realization for `m` in `varinfo_linked`.
varinfo_linked[@varname(m)] = varinfo_linked[@varname(x)] + 1
varinfo_linked[@varname(m)], varinfo_linked[@varname(x)]

(1.6054721554645925, 0.6054721554645925)

Now we see that the constraint m < x is no longer satisfied!

Hence one might expect that if we try to compute, say, the logjoint using varinfo_linked with this "invalid" realization, we'll get an error:

logjoint(model, varinfo_linked)

-3.327851247482377

But we don't! In fact, if we look at the actual value used within the model

first(DynamicPPL.evaluate!!(model, varinfo_linked, DefaultContext()))

(m = 1.6054721554645925, x = 1.851484507047774)

we see that we indeed satisfy the constraint m < x, as desired.

The reason for this is that internally in a model evaluation, we construct the transformation from the internal to the model representation based on the current realizations in the model! That is, we take the dist in a x ~ dist expression at model evaluation time and use that to construct the transformation, thus allowing it to change between model evaluations without invalidating the transformation.

But to be able to do this, we need to know whether the variable is linked / "unconstrained" or not, since the transformation is different in the two cases. Hence we need to be able to determine this at model evaluation time. Hence the internals end up looking something like this:

if istrans(varinfo, varname)
    from_linked_internal_transform(varinfo, varname, dist)
else
    from_internal_transform(varinfo, varname, dist)
end

That is, if the variable is linked / "unconstrained", we use the DynamicPPL.from_linked_internal_transform, otherwise we use DynamicPPL.from_internal_transform.

And so the earlier diagram becomes:

This is also the reason why we have two definitions of getindex:

• getindex(::AbstractVarInfo, ::VarName, ::Distribution): used internally in model evaluations with the dist in a x ~ dist expression.
• getindex(::AbstractVarInfo, ::VarName): used externally by the user to get the realization of a variable.

For getindex we have the following diagram:

While if dist is not provided, we have:

Notice that dist is not present here, but otherwise the diagrams are the same.

Other functionalities

There are also some additional methods for transforming between representations that are all automatically implemented from DynamicPPL.from_internal_transform, DynamicPPL.from_linked_internal_transform and their siblings, and thus don't need to be implemented manually.

Convenience methods for constructing transformations:

from_maybe_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])

to_maybe_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])

internal_to_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)

linked_internal_to_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])

Convenience methods for transforming between representations without having to explicitly construct the transformation:

to_maybe_linked_internal(vi::AbstractVarInfo, vn::VarName, dist, val)

from_maybe_linked_internal(vi::AbstractVarInfo, vn::VarName, dist, val)

Supporting a new distribution

To support a new distribution, one needs to implement for the desired AbstractVarInfo the following methods:

• DynamicPPL.from_internal_transform
• DynamicPPL.from_linked_internal_transform

At the time of writing, VarInfo is the one that is most commonly used, whose internal representation is always a Vector. In this scenario, one can just implement the following methods instead:

from_vec_transform(dist::Distribution)

from_linked_vec_transform(dist::Distribution)

These are used internally by VarInfo.

Optionally, if inverse of the above is expensive to compute, one can also implement:

• DynamicPPL.to_internal_transform
• DynamicPPL.to_linked_internal_transform

And similarly, there are corresponding to-methods for the from_*_vec_transform variants too

to_vec_transform(x)

to_linked_vec_transform(dist)

TL;DR

• DynamicPPL.jl has three representations of a variable: the model representation, the internal representation, and the linked internal representation.

  • The model representation is the representation of the variable as it appears in the model code / is expected by the dist on the right-hand-side of the ~ in the model code.
  • The internal representation is the representation of the variable as it appears in the varinfo, which varies between implementations of AbstractVarInfo, e.g. a Vector in VarInfo. This can be converted to the model representation by DynamicPPL.from_internal_transform.
  • The linked internal representation is the representation of the variable as it appears in the varinfo after linking. This can be converted to the model representation by DynamicPPL.from_linked_internal_transform.

• Having separation between internal and linked internal is necessary because transformations might be constructed at the time of model evaluation, and thus we need to know whether to construct the transformation from the internal representation or the linked internal representation.