Nonparametrics

[rlouf]

Nonparametics

DRAFT

Gathering my thoughts on how to implement nonparametrics in mcx. This needs to be thought through before freezing the internal API as it is likely to impact it substantially. Implementation can wait, but API design must be ready before first release.

Here is a purely speculative design:

@mcx.model
def DirichletProcess(alpha, base_measure):
	
    sticks_cache = []
    atoms_cache = []

    @mcx.model
    def sticks(k):

        @mcx.model
        def new_stick():
            b <~ Beta(1, alpha)
            sticks_cache.append(b)
            return b

        return jax.lax.cond(
            k <~ len(sticks_cache),
            lambda _: sticks_cache[:k],
            lambda _: new_stick(),
            None
        )


	@mcx.model
    def atoms(k):
		
        @mcx.model
        def new_atom():
            a <~ base_measure
            atoms_cache.append(a)
            return a

        return jax.lax.cond(
            k <~ len(atoms_cache),
            lambda _: atoms_cache[k],
            lambda _: new_atom(),
            None
        )
		

    @mcx.model
    def distribution():
		k = 0
		while True:
			stick <~ sticks(k)
			do_pick <~ Bernoulli(stick)
			if do_pick:
				return atom(k)
            k += 1

    return distribution

This is nice, but to see how this would impact MCX's design we need to answer the following questions:

1. What does it look like to sample from distribution?
2. What does it look like to sample from DirichletProcess?
3. What does the logpdf of ditribution look like?
4. What does the logpdf of DirichletProcess look like?

References

General reference on probabilistic programming and what a language would need to be able to implement non-parametrics (look for "HOPPL"):

• https://arxiv.org/abs/1809.10756

General literature on nonparametrics:

• http://mlg.eng.cam.ac.uk/zoubin/talks/uai05tutorial-b.pdf

Stochastic memoization:

• http://danroy.org/papers/RoyManGooTen-ICMLNPB-2008.pdf
• https://dp.tdhopper.com/hdp-sampling/

Collection of samplers for infinite mixture models: https://github.com/tscholak/imm

Truncated Dirichlet Process in PyMC3:

• https://docs.pymc.io/notebooks/dp_mix.html