Errors in Sampling When Parameter Bounds Depend on Parameters

[farr]

I am finding a lot of bugs in Turing samples when the bounds on distributions depend on parameters. It looks like the Jacobians of the mappings from the constrained and un-constrained space are not being incorporated properly in the sampling. For example, this simple model generates a ton of numerical errors:

using Turing

@model function truncated()
    x ~ Uniform(0,1)
    y ~ Uniform(0,x)
end

trace = sample(truncated(), NUTS(), 1000) # Generates a bunch of numerical errors, and `y` ends up "frozen."

and the sampling over y ends up "frozen" at a single value. For example:

More complicated models suffer from the same problem:

@model function truncated_normal()
    x ~ Normal(0,1)
    y ~ truncated(Normal(0,1), -Inf, x)
end

trace2 = sample(truncated_normal(), NUTS(), 1000)

I haven't tried to dig into the code to figure out how things are failing, but it has to be a problem with incorporating the Jacobian of the constrained<->unconstrained transform.

[devmotion]

You have to define a custom bijector since the support of the distribution of y is stochastic. The discussion in https://github.com/TuringLang/Turing.jl/issues/1270 shows how this can be done.

[farr]

Great---thanks! FWIW, I think it's definitely worth supporting this natively in the backend, even if it adds some complication. It's super useful. (I come from the Stan world, and all the time am building models where some parameter y is constrained to have a range that depends on some previous parameter x; it would be a huge pain to have to do it in a custom way each time.)

[torfjelde]

As mentioned in the issue mentioend by @devmotion , you can do:

struct Uniform2D <: ContinuousMultivariateDistribution end

import Random: AbstractRNG
function Distributions.rand(rng::AbstractRNG, ::Uniform2D)
    x = rand(rng, Uniform(0, 1))
    y = rand(rng, Uniform(0, x))
    return [x, y]
end

function Distributions.logpdf(::Uniform2D, x::AbstractVector{<:Real})
    return logpdf(Uniform(0, 1), x[1]) + logpdf(Uniform(0, x[1]), x[2])
end

function Bijectors.bijector(::Uniform2D)
    b1 = Bijectors.Stacked((Bijectors.TruncatedBijector{1}([0], [1]), Identity{1}()))
    m = Bijectors.PartitionMask(2, [2], [1])
    b2 = Bijectors.Coupling(x -> Bijectors.TruncatedBijector(0, x), m)

    return b1 ∘ b2
end

@model function demo()
    z ~ Uniform2D()
end

trace = sample(demo(), NUTS(), 5000) # Generates a bunch of numerical errors, and `y` ends up "frozen."
plot(trace)

which results in

Which seems more reasonable.

BUT what we definitively should and can improve is how to define such custom distributions with stochastic support. The Coupling construction is too ugly and the Stacked is also a bit annoying.