Add MultivariateNormalDiag distribution

[samuela]

No description provided.

[vishwakftw]

cc: @fritzo @alicanb @fehiepsi

[alicanb]

I'm not sure we actually need this. This can be done with Normal only (+ Independent if you're strict about event_shape)

[fritzo]

@alicanb is correct, this can be accomplished via Independent(Normal(loc, scale_diag), 1). Maybe that should be an example in the Normal or Independent docs?

[alicanb]

@fritzo there is an example in Independent docs. (written by you actually 😄 )

[samuela]

I wasn't aware of the Independent(Normal(loc, scale_diag), 1) trick. I use diagonal normals quite a lot, so it would be nice if there was a direct way to construct them that was a little bit more explicit in intent than Independent(Normal(loc, scale_diag), 1). But I understand that there's a tradeoff here.

Are there any performance benefits to having a dedicated implementation?

[zou3519]

Have you tried benchmarking @samuela?

[samuela]

@zou3519 I have not. I assumed that this PR was sort of doomed based on the discussion.

[fritzo]

I think there will be negligible performance benefit, since Independent simply calls its base dist

    def rsample(self, sample_shape=torch.Size()):
        return self.base_dist.rsample(sample_shape)

    def log_prob(self, value):
        log_prob = self.base_dist.log_prob(value)
        return _sum_rightmost(log_prob, self.reinterpreted_batch_ndims)

Also IMO it is cleaner to provide a general Independent wrapper to create diagonal distributions of any base type, rather than implement DiagNormal, DiagGamma, DiagCauchy, DiagBeta etc. If you find Independent(Normal(...), ...) cumbersome, we could add a Normal.independent(...) method as we have in Pyro.

[samuela]

@fritzo I find myself using a lot of diagonal normal distributions in my work. For me it's just cleaner and communicates intent more clearly to write MultivariateNormalDiag instead of Independent(Normal(...), 1). Right now I've just written a function MultivariateNormalDiag that calls Independent(Normal(...), 1).

[kashif]

typo here?

[samuela]

ah, yes. will fix

[fritzo]

Can we replace this complex PR in favor of @samuela's simple wrapper?

def MultivariateNormalDiag(loc, scale_diag):
    if loc.dim() < 1:
        raise ValueError("loc must be at least one-dimensional.")
    return Independent(Normal(loc, scale_diag), 1)

EDIT sorry if I misunderstand the purpose. It would help to add a PR description.

[JakobHavtorn]

As far as I am aware, another reason for wanting a separate diagonal class is the torch.distributions.kl_divergence raises a NotImplementedError when called on a diagonal Normal defined with Independent(Normal(loc, scale_diag), 1). It works fine with MultivariateNormal but this can be memory intensive.

[fritzo]

torch.distributions.kl_divergence raises a NotImplementedError

Note this works with pyro.distributions. If you'd like this behavior with torch.distributions, feel free to move upstream the Pyro implementations of kl_divergence(Independent,Independent) and kl_divergence(Independent(Normal),MultivariateNormal).

@samuela What do you think of Pyro's abbreviated syntax .to_event(), used as

Normal(...).to_event(1) = Independent(Normal(...), 1)

If there is demand we can move that upstream to torch.distributions.Distribution. I think that would be a better solution than a custom new class that works for only one distribution.

[samuela]

@fritzo I've started using jax now actually, so I haven't been following this issue much. .to_event(...) looks interesting although I would worry about offering to ways to accomplish the same thing.

[JakobHavtorn]

@fritzo Thank you, it is completely straightforward to get that working.

[lqf96]

Is the current log_prob behavior something we want? Because the following implementation:

def log_prob(self, value):
        log_prob = self.base_dist.log_prob(value)
        return _sum_rightmost(log_prob, self.reinterpreted_batch_ndims)

is clearly different from the PDF described in the Tensorflow Distributions document:

pdf(x; loc, scale) = exp(-0.5 ||y||**2) / Z,
y = inv(scale) @ (x - loc),
Z = (2 pi)**(0.5 k) |det(scale)|,

I think many people expects the log of PDF when they are calling log_prob, but clearly PyTorch is not returning the log of PDF for MultivariateNormalDiag on Independent(Normal(loc, scale), 1).

[pytorchbot]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
Stale pull requests will automatically be closed 30 days after being marked Stale

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.