[HMC] [IMM] gradient issue

[emilemathieu]

Naively run gradient on this infinite mixture model fails cf JuliaDiff/ForwardDiff.jl#280 (comment).

• Make gradient works on such a model.

• Check that the obtain gradient matches the gradient of the analytical joint of a mixture model with a NRM or PK prior.

[ae-foster]

Ok @emilemathieu I think I have (at least partially) understood the problem. The main problem is the transformation we apply to weights:
https://github.com/yebai/Turing.jl/blob/project-bnp/src/samplers/support/transform.jl#L36

link{T<:Real}(d::TransformDistribution, x::Union{T,Vector{T}}) = begin
  a, b = minimum(d), maximum(d)
  lowerbounded, upperbounded = isfinite(a), isfinite(b)
  if lowerbounded && upperbounded
    logit((x - a) ./ (b - a))
  elseif lowerbounded
    log(x - a)
  elseif upperbounded
    log(b - x)
  else
    x
  end
end

In our case, maximum(d) = T_surplus depends on the value of other variables. This already seems like a bad idea, because a change in the value of the first weight implicitly changes the values of all other weights.

And then, to cause the actual error we observed (making the total mass negative) the problem is that the value of T_surplus is not propagated to the link function.

At https://github.com/yebai/Turing.jl/blob/project-bnp/src/samplers/hmc.jl#L167, I added the following

println("dist of vi[vn] ", getdist(vi, vn))
println("dist from assume (current runmodel) ", dist)

and saw

dist of vi[vn] Turing.SBS_DP(alpha=10, T_surplus=Dual(0.22720660976717053,..))
dist from assume (current runmodel) Turing.SBS_DP(alpha=10, T_surplus=Dual(0.009301975557913407,..))

and then, by transforming the HMC proposed value using a T_surplus of 0.22720, but trying to compute a log-likelihood assuming a T_surplus of 0.009301975. Well... it's bad.

If we change the transform to simply log, is there anything that would allow us to reject proposed weights that are out-of-bounds?

[ae-foster]

Another thing to do would be to never sample T itself. Instead sample W_1, ..., W_n, T^* all constrained just to be positive.

[emilemathieu]

There are 2 things here:

• Generically, vi.dists[vn] is never updated in HMC which means that for distributions which need to be transformed, the transformation can only depends on the distribution type and NOT on parameters of that distribution. Following the 2sd point that's seems like a good safeguard. @xukai92 is that intended ?

• This is the case for the lazy size-biased sampling case since J[k] ~ SBS_DP(param, T_surplus) has support on [0, T_surplus].. This is an issue since updating all J[1], ..., J[K] at once in an HMC step will yield new J[1], ..., J[K]with potentially sum(J[1:K])>T implying T_surplus < 0.
  Consequently we might:

  • Sample consecutively the J[i] conditioned on the other J[1],...,J[i-1],J[i+1],...,J[k], i.e. applying a kernel for each J[i] (such as MH or HMC). Still not sure how build this update without violating constraints.
  • Sample them all at once but with a simplex constraint transformation ?

[emilemathieu]

@ae-foster Is HMC working on the atoms' location :x at least ?
The weights could be updated with CSMC if we don't find a MH or HMC way.

[mlomeli1]

@ae-foster , note that each size biased weight has different support in the current parameterisation: namely, say we first draw the total mass T, then, J_1 \in (0,T), J_2 \in (0,T-J1), ... J_k \in (0,T-J1-...T_{k-1}), so we have to be careful about this.
An alternative is to use the parameterisation from Favaro and Walker (2012)"Slice sampling sigma-stable Poisson-Kingman mixture models" JCGS. They reparameterise in such a way that the size biased weights are between (0,1). This could be the simplex transformation that @emilemathieu had in mind? You can see the joint distribution, conditional of the total mass in my thesis Section 3.6 equation 3.8. This joint is potentially be useful for a global update with an HMC sampler.

[ae-foster]

@emilemathieu I think it is working (would have to check though). Also I don't think it's time to give up yet! What about the idea of avoiding T altogether and focusing on T^*? I'm getting that idea from https://arxiv.org/pdf/1401.1137.pdf page 26.

In general, though, HMC won't work when the state space is highly dynamic. I wonder precisely which probabilistic programs it wouldn't work for... 🤔

[emilemathieu]

@ae-foster @mlomeli1 We could do the inference on the size-biased stick length V_is otherwise (~ Beta(1, alpha) for the DP). See https://github.com/yebai/Turing.jl/blob/project-bnp/test/rpm.jl/imm_hmc_sticks.jl#L27.

[mlomeli1]

We could do @emilemathieu , if we handle the DP case separately from the Q-class (which is actually good to do since it is a limit case of such class). For the Q-class have a look at the Favaro and Walker (2013) paper I cited. It has a reparameterisation of the size-biased weights in terms of (0,1) random variables. We know that, for all priors apart from DP and PY, these are no longer independent nor Beta distributed. The expression in equation 3.8 of my thesis (which reviews the paper) has the joint conditional distribution of the (0,1) r.v.s. given the total mass so if we multiply it by the prior of the total mass we get the joint distribution of interest.
Another comment is, just keep in mind that the Caron et al. graphs paper uses NRMs. Not all NRMs belong to the Q-class and not all priors in the Q-class are NRMs, for instance, the PY is not a NRM but it is a member of the Q-class. So before we could use the parameterisation suggested by @ae-foster we need to check if it is also valid for the Q-class.

[emilemathieu]

@mlomeli1 Indeed the stick length are no longer Beta distributed and independent for other classes than the DP and PYP. Yet the dependancy issue with weights Js is specifically that their support depends on previous weights. Otherwise the fact that their distribution parameters is a function of the previous weight is not an issue, such as for this simple example.

In Favaro et al (2014) you propose a way to sample the stick length V_k (conditioned on the total mass and the previous V_i). Such as for the NIGP:

Distributions.rand(d::SBS_NIGP) = begin
  X = sqrt(rand(Gamma(3/4, 1)))
  Y = sqrt(rand(InverseGamma(1/4, 1/(64 * d.T_surplus^2))))
  return X / (X + Y)
end

I believe we should directly do inference on these stick lengths by rewriting the model with these variables: https://github.com/yebai/Turing.jl/blob/project-bnp/test/rpm.jl/imm_hmc_sticks.jl#L27.

[mlomeli1]

Indeed, @emilemathieu, that identity in distribution is useful for sampling a new size biased weight for any (sigma=0.5)-Stable PK model. However, I thought the issue in the thread is about how to update the existing table weights together with all the other continuous r.v. such as table parameters. Which is why I mention the joint distribution in Favaro and Walker (2012) in which the support doesn't depend on previous weights. It is just a reparameterisation of the size-biased weights.

Alternatively, we could update each weight conditioned on the previous weights. I used slice sampler within Gibbs but I am not sure it's the best choice.

[yebai]

This might be fixed in #515.