Probability: feature request: KL divergence for Gaussian Mixture Model

[BelhalK]

Hi all,

Would it be possible to add KL between two Mixture of gaussians distirbutions or even between one multivariate gaussian and a mixture of gaussian?
Example:

A =tfd.Normal( loc=[1., -1],scale=[1, 2.])
B =tfd.Normal( loc=[1., -1],scale=[1, 2.])
C = tfd.MixtureSameFamily(
    mixture_distribution=tfd.Categorical(
        probs=[0.3, 0.7]),
    components_distribution=tfd.Normal(
      loc=[-1., 1],       # One for each component.
      scale=[0.1, 0.5]))  # And same here.
D = tfd.MixtureSameFamily(
    mixture_distribution=tfd.Categorical(
        probs=[0.1, 0.9]),
    components_distribution=tfd.Normal(
      loc=[-1., 1],       # One for each component.
      scale=[0.1, 0.5]))  # And same here.

Then

tf.distributions.kl_divergence(A, B)

works fine and returns <tf.Tensor 'KullbackLeibler/kl_normal_normal/add:0' shape=(2,) dtype=float32>
Yet, tf.distributions.kl_divergence(A, D)and tf.distributions.kl_divergence(C, D) returns

No KL(distribution_a || distribution_b) registered for distribution_a type Normal and distribution_b type MixtureSameFamily

This would be of great value since more and more KL of this type are being calulated (or at least approximated) while training Neural Nets.

Thanks
Belhal

[srvasude]

+1. For Mixtures, unfortunately the KL divergence is not analytically tractable. However, I do agree having approximations / bounds to this would be generally useful (possibly living in the MixtureSameFamily class).

Out of curiousity, does computing the sample KL work well for you?

[BelhalK]

Indeed it's not tractable but we can use a Monte Carlo estimator to approximate a bound on it, that would be the idea on how to build a loss function when the BNN presents mixture of distributions as priors or posteriors on the weights.

What do you mean by sample KL? Is there a function that calculates KL between those distributions pointwise? If yes, I'd be happy to try.
Let me know

[srvasude]

By sample KL, I meant a Monte carlo estimator for the KL divergence between two Gaussian Mixtures (which currently doesn't exist as a function).
There are analytic upper / lower bounds for Gaussian mixtures that could possibly be library functions.

[BelhalK]

Yes I can build a simple function that returns the value of a Monte carlo estimator for the KL divergence between two Gaussian Mixtures evaluated at one point.
But for integration in a tf neural net I would need more than that, I would need a loss function that takes as inputs trainable variables, i.e. a <tf.Tensor 'KullbackLeibler/kl_normal_normal/add:0' shape=.. dtype=float32> object

There are analytic upper / lower bounds for Gaussian mixtures that could possibly be library functions

You mean external library, or even TF, could be used to calculate these bounds?

[dave-fernandes]

This problem is more general than just for Gaussian Mixtures. It would also be nice to use Horseshoe or StudentT distributions for priors in the BNN layers. Currently you get the same error when attempting to use those as priors.

[dave-fernandes]

Hmm. I just tried again using tfp.distributions instead of tf.distributions, and I can use Gaussian Mixture and Horseshoe distributions as priors. So looks like this issue can be closed @BelhalK ?

[BelhalK]

It works perfectly, Thanks guys!

[whusym]

It seems that it still hasn't been implemented?

I tried the following code:

tfd = tfp.distributions
mix = 0.3
bimix_gauss_1 = tfd.Mixture(
  cat=tfd.Categorical(probs=[mix, 1.-mix]),
  components=[
    tfd.Normal(loc=-1., scale=0.1),
    tfd.Normal(loc=+1., scale=0.5),
])

bimix_gauss_2 = tfd.Mixture(
  cat=tfd.Categorical(probs=[mix, 1.-mix]),
  components=[
    tfd.Normal(loc=-1., scale=0.1),
    tfd.Normal(loc=+1., scale=0.5),
])

bimix_gauss_1.kl_divergence(bimix_gauss_2)

But the errror is still

NotImplementedError: No KL(distribution_a || distribution_b) registered for distribution_a type Mixture and distribution_b type Mixture

[dave-fernandes]

Yes, I just tried again and it doesn't work. I'm not sure if this is a regression or if I was hallucinating that it worked a few days ago. Maybe it should be reopened as a bug.

[csuter]

Is it only in tfp-nightly? We haven't made a stable (tensorflow-probability) release in a while. We hope to have one out in the next few weeks (pending TF 1.14).

[dave-fernandes]

I've been using Colab, I haven't kept track of what build is installed there.

[csuter]

Ah, colab is on most recent stable by default. You can install nightly by placing a cell with !pip install tf-nightly[-gpu] tfp-nightly prior to your imports.

[dave-fernandes]

Doesn't seem to work in the nightly.

[csuter]

hm, can you share an example failure?

[csuter]

(Sorry I'm not immediately familiar with the original resolution of this issue, but happy to help dig in)

[whusym]

So my version of tfp on colab is '0.7.0-dev20190430'

The error message is still the same. I used my same code above.

[dave-fernandes]

@whusym has a nice short example posted above. Mine is a bit more complex. I want to use a Horseshoe distribution as a kernel prior in a DenseFlipout layer.

def kernel_prior(dtype, shape, name, trainable, add_variable_fn):
  return tfd.Horseshoe(scale=5.0)

...

logits = tfp.layers.DenseFlipout(5, name='Output', kernel_prior_fn=kernel_prior)(previous_layer)

[csuter]

I think the issue with @whusym's example is it's using Mixture, not MixtureSameFamily

[dave-fernandes]

The error I get is:

NotImplementedError: No KL(distribution_a || distribution_b) registered for distribution_a type Independent and distribution_b type Horseshoe

[whusym]

Thanks @csuter . I tried @BelhalK 's code above as follows:

import tensorflow_probability as tfp
tfd = tfp.distributions
A =tfd.Normal( loc=[1., -1],scale=[1, 2.])
B =tfd.Normal( loc=[1., -1],scale=[1, 2.])
C = tfd.MixtureSameFamily(
    mixture_distribution=tfd.Categorical(
        probs=[0.3, 0.7]),
    components_distribution=tfd.Normal(
      loc=[-1., 1],       # One for each component.
      scale=[0.1, 0.5]))  # And same here.
D = tfd.MixtureSameFamily(
    mixture_distribution=tfd.Categorical(
        probs=[0.1, 0.9]),
    components_distribution=tfd.Normal(
      loc=[-1., 1],       # One for each component.
      scale=[0.1, 0.5]))  # And same here.

print (tfp.distributions.kl_divergence(C, D))

It seems that it's still not implemented. The error is:

NotImplementedError: No KL(distribution_a || distribution_b) registered for distribution_a type MixtureSameFamily and distribution_b type MixtureSameFamily

[csuter]

ok actually reading more closely i don't think this was (or could be) implemented. sounds like the workaround was to use monte carlo estimates of the KL.

[whusym]

Thanks again @csuter . I agree. The closest form I've found is:

https://www.tensorflow.org/probability/api_docs/python/tfp/monte_carlo/expectation

Here is the way to do MC estimates using sklearn:
https://stackoverflow.com/questions/26079881/kl-divergence-of-two-gmms?noredirect=1&lq=1

[dave-fernandes]

Thanks @csuter . Given the kernel_prior_fn argument for TFP layers, would it be reasonable for the framework to automatically use a MC estimate if it doesn't have an analytic formula for KL? Are there any priors besides the default one which would not require this approach?

[csuter]

I can imagine that being a reasonable feature. I'd worry a bit (ok, maybe a lot) about noise in the stochastic gradients of the MC-estimated KL, but something is infinitely better than nothing :)

We have lots of analytical KL pairs defined (grep for "RegisterKL" in the distributions package). The requirement for using these with variational layers is that 1. the KL is defined, and 2. the layer is configured to use the given prior and (variational) posterior distributions. I don't think we have a lot of analytical KL's for distributions from 2 distinct families, though.

[csuter]

here's a demo of registering MC KL in TFP:

https://colab.research.google.com/drive/1RHx23ViJL-1SIgVYbUjB8QjCeQ9b2fVS

[csuter]

(updated link with a new version that i can share outside of our corp system)

[dave-fernandes]

Thank you!

Hopefully the registered pairs will make it into the docs for the distributions eventually.

[csuter]

Ah, yeah, that's a good idea! We'll have to give it some thought, since the registrations don't produce any actual documented, exported symbols. But it would certainly be nicer than having to grep around!

[whusym]

Thanks @csuter ! I change tfd.MixtureSameFamily to tfd.Mixture, and it also works!

[csuter]

Right. Note that if your mixture is over distributions of the same family, MSF is going to be more efficient because operations can be vectorized.

[brianwa84]

https://github.com/tensorflow/probability/blob/d24e333cb2742dee0fad6a4aabf1484766aea17f/tensorflow_probability/g3doc/api_docs/python/tfp/distributions/kl_divergence.md has an updated doc listing built-in KLs

[dave-fernandes]

Thanks, @brianwa84 !
What does the + mean in the KL registration list?

[brianwa84]

The +'s are subclasses for which some superclass has a registration. So
something w/ a + has a KL registered with everything on the other side up
to the preceding || and down to just before the next ||.

Brian Patton | Software Engineer | [email protected]

*From: *Dave Fernandes [e-mail]
*Date: *Mon, May 6, 2019 at 11:11 AM
*To: *tensorflow/probability
*Cc: *Brian Patton, Mention

Thanks, @brianwa84 https://github.com/brianwa84 !

What does the + mean in the KL registration list?

—
You are receiving this because you were mentioned.
Reply to this email directly, view it on GitHub
https://github.com/tensorflow/probability/issues/199#issuecomment-489656874,
or mute the thread
https://github.com/notifications/unsubscribe-auth/AFJFSIZ75ANIDTJHWGA5G73PUBDADANCNFSM4F76A2MA
.

[brianwa84]

I also added a listing of built-in KL registrations to each distribution's
kl/xent doc
e.g.
https://github.com/tensorflow/probability/blob/master/tensorflow_probability/g3doc/api_docs/python/tfp/distributions/Normal.md#kl_divergence
I think my work here is done.

*From: *Brian Patton 🚀 [e-mail]
*Date: *Mon, May 6, 2019 at 12:22 PM
*To: *tensorflow/probability

The +'s are subclasses for which some superclass has a registration. So

something w/ a + has a KL registered with everything on the other side up
to the preceding || and down to just before the next ||.

Brian Patton | Software Engineer | [email protected]

*From: *Dave Fernandes [e-mail]
*Date: *Mon, May 6, 2019 at 11:11 AM
*To: *tensorflow/probability
*Cc: *Brian Patton, Mention

Thanks, @brianwa84 https://github.com/brianwa84 !

What does the + mean in the KL registration list?

—
You are receiving this because you were mentioned.
Reply to this email directly, view it on GitHub
https://github.com/tensorflow/probability/issues/199#issuecomment-489656874,
or mute the thread
https://github.com/notifications/unsubscribe-auth/AFJFSIZ75ANIDTJHWGA5G73PUBDADANCNFSM4F76A2MA
.

[dave-fernandes]

here's a demo of registering MC KL in TFP:

https://colab.research.google.com/drive/1RHx23ViJL-1SIgVYbUjB8QjCeQ9b2fVS

I'm having the following problem when I try to use this MC method to train a variational layer...

The prior is created by a method with signature kernel_prior_fn(dtype, shape, name, trainable, add_variable_fn), and I return a distribution using shape as the batch_shape. In my case, batch_shape = [5, 1, 32]; event_shape = []

However the registered KL method is called with the default posterior distribution (Independent) as the first argument. And it has batch_shape = []; event_shape = [5, 1, 32]. So the shapes don't match.

Should my prior have event_shape = [5, 1, 32]? And if so, how do I set the event shape? The docs only describe how to set the batch shape.

[csuter]

You can use tfd.Independent to "reinterpret" batch dims as event dims. Think of it as building the joint distribution over a batch of independent (non-identical) distributions in the same family. In your case you'll want tfd.Independent(your_prior, reinterpretted_batch_ndims=3).

[dave-fernandes]

Ah, okay. Got it. Thank you!