Ordered priors

[luiarthur]

In Turing, do we currently have something like ordered vectors from STAN?

What I mean is this:

Say I have the following GMM with two components:

@model GMM(y)
  mu ~ filldist(Normal(), 2)
  sigma ~ filldist(LogNormal(), 2)
  eta ~ Dirichlet(2, 0.5)
  y .~ UnivariateGMM(mu, sigma, Categorical(eta))
end

In addition, I would like to impose that mu[1] < mu[2]. In STAN, I would accomplish this by defining mu as an ordered vector. What is the recommended way of doing this in Turing?

[luiarthur]

What I mean is this:

Say I have the following GMM with two components:

@model GMM(y)
  mu ~ filldist(Normal(), 2)
  sigma ~ filldist(LogNormal(), 2)
  eta ~ Dirichlet(2, 0.5)
  y .~ UnivariateGMM(mu, sigma, Categorical(eta))
end

In addition, I would like to impose that mu[1] < mu[2].

[torfjelde]

We don't, though we have specific instance(s) implemented I believe (e.g. OrderedLogistic

Turing.jl/src/stdlib/distributions.jl

Line 134 in 8886d35

struct OrderedLogistic{T1, T2<:AbstractVector} <: DiscreteUnivariateDistribution

).

But for the specific stuff you're mentioning here, you could implemented it as a Bijector and then transform the distribution to make the samples ordered using transform(distribution, ordered_bijector). Then you could make this your prior, e.g. your example would become:

@model function GMM(y)
    mu ~ transform(filldist(Normal(), 2), OrderedBijector())
    sigma ~ filldist(LogNormal(), 2)
    eta ~ Dirichlet(2, 0.5)
    y .~ UnivariateGMM(mu, sigma, Categorical(eta))
end

We could also just add in a convenience function, e.g.

ordered(d::Distribution) = Bijectors.transformed(d, OrderedBijector())

An implementation of OrderedBijector would probably be something like:

import Bijectors: Bijector, Inverse, logabsdetjac

struct OrderedBijector <: Bijector{1} end

function (b::OrderedBijector)(x::AbstractVector)
    y = similar(x)
    y[1] = x[1]
    @. y[2:end] = log(x[2:end] - x[1:end - 1])

    return y
end

function (ib::Inverse{<:OrderedBijector})(y::AbstractVector)
    x = similar(y)
    x[1] = y[1]
    x[2:end] = y[1] .+ cumsum(exp.(y[2:end]))

    return x
end

logabsdetjac(ib::Inverse{<:OrderedBijector}, y::AbstractVector) = sum(y[2:end])
logabsdetjac(b::OrderedBijector, x) = -logabsdetjac(inv(b), b(x))

The annoying bit is that we also need to implement custom adjoints for Zygote.jl and ReverseDiff.jl due to the fact that here we're using mutation.

[luiarthur]

Ah, thanks! This is very helpful.

The annoying bit is that we also need to implement custom adjoints for Zygote.jl and ReverseDiff.jl due to the fact that here we're using mutation.

This avoids mutation so should just work, right?

function (b::OrderedBijector)(x::AbstractVector)
    y = [(k == 1) ? x[k] : log(x[k] - x[k-1]) for k in eachindex(x)]
    return y
end

function (ib::Inverse{<:OrderedBijector})(y::AbstractVector)
    z = y[1] .+ cumsum(exp.(y[2:end]))
    x = [(k == 1) ? y[k] : z[k-1] for k in eachindex(y)]
    return x
end

[luiarthur]

Something still seems off to me. But I don't exactly know how to articulate it.

The inverse transform gets you an ordered vector (from the unconstrained space).

To do the forward transform, you need to start with an ordered vector. But the order constraint isn't imposed anywhere (besides the inverse transform). What am I missing? Does this constraint need to be imposed?