Implement pathwise derivative for dirichlet distribution about numpyro

[martinjankowiak]

i guess to do that you'd to override sample and define a custom jvp or jvp rule

[fehiepsi]

More context about this issue:

Dirichlet sampler is available in #81. However, the derivative w.r.t. to concentration parameter is computed based on pathwise derivative of the gamma sampler which generates dirichlet samples. It would be better to have a pathwise derivative implementation for Dirichlet distribution on its own. An algorithm for implementing can be found in this paper.

[fehiepsi]

The current Dirichlet sampler has a reparameterized gradient but its variance would be higher than the version in PyTorch. Given that we mostly rely on the upstream jax samplers now, it is better to use the current high variance version and raise the request to jax team (if needed) instead.

[adamgayoso]

Hello,

I'm trying to sample from a Dirichlet with 4 dimensions, but over a batch of thousands of elements. The gradients are very slow compared to PyTorch. Has there been any progress on speeding this up ?

[fehiepsi]

Could you try TFP Dirichlet distribution? IIRC it has faster dirichlet sampler.

[adamgayoso]

I'm getting nan's there related to dirichlet sampling where I'm not with jax.random.dirichlet/numpyro!...

[adamgayoso]

I can confirm the sampling if faster in TFP, but I'll get NaN's during training, which I haven't narrowed down to a gradient issue or a sampling issue.

[fehiepsi]

You might want to create a wrapper to clip Dirichlet samples. Something like

class RobustDirichlet(...):
    def sample(self, ...):
        samples = super().sample(...)
        return clip(samples)

[martinjankowiak]

@adamgayoso it might help to make sure your parameter is bounded away from zero e.g. concentration=0.01 + positive_param

[adamgayoso]

Thank you both for the suggestions, although I'm still getting nan's. I'd need more time to inspect if it's the gradients, but do you know if I can stitch the jax random dirichlet gradient fn with tfp gamma sampler?

[martinjankowiak]

from tensorflow_probability.substrates.jax import distributions as tfd

def model():
    numpyro.sample("x", tfd.Normal(0, 1))

[adamgayoso]

Thank you. Where is that happening?

I see this wrapper: https://num.pyro.ai/en/stable/distributions.html#numpyro.contrib.tfp.distributions.TFPDistribution

But not where a custom gradient fn is being set.

[adamgayoso]

@fehiepsi would you be able to check out google/jax#12943 ?

The sampling wasn't the bottleneck, it was actually the lax map for the grads on cpu. I also fixed the lax cond in the sampling loop