Choose parametrization in sampling statement

[aseyboldt]

Non-centered parametrizations are often a better choice than the centered parametrizations that are more natural to code in pymc3:

val = pm.Normal('val', mu=mu, sd=sd)

vs

val_raw = pm.Normal('val_raw', mu=0, sd=1)
val = pm.Deterministic('val', val_raw * sd + mu)

Couldn't we add a parameter parametrization or similar to distributions, such that we can automate that? This would also make it easier to for new users to try a couple of different parametrizations if something does not work well. I think we could reuse the framework we have in place for transformations. All we need for the non-centered parametrization is an affine transformation.

Some useful candidates:

# non-centered Normal
pm.Normal('val', mu=mu, sd=sd, parametrization='non-centered')
# as shortcut for
val_raw = pm.Normal('val_raw', mu=0, sd=1)
val = pm.Deterministic('val', val_raw * sd + mu)

# non-centered MvNormal
pm.MvNormal('val'" mu=mu, cov=cov, parametrization='non-centered')
# ->
val_raw = pm.Normal('val_raw', mu=0, sd=1)
pm.Deterministic('val', tt.slinalg.cholesky(cov).dot(val_raw))

# non-centered Lognormal on log space
pm.Lognormal('val', mu=mu, sd=sd, parametrization='log-non-centered')
# ->
val_raw = pm.Normal('val_raw', mu=0, sd=1)
val = tt.exp(val_raw * sd + mu)

# logit-sf for eg cauchy (not sure about this one, but it would work
# for pretty much any distribution if we can compute the survival function)
pm.Cauchy('val', alpha=alpha, beta=beta, parametrization='logit-sf')
# ->
val_raw = pm.Uniform('val_raw', lower=0, upper=1, transform='logit')
val = stats.cauchy(alpha=alpha, beta=beta).sf(val_raw)

Do you think this would be useful / would work? Other ideas for common reparametrizations (or better names ;)?

[twiecki]

That's a really neat idea. I strongly support it. Could you mock up an example Distribution to see what it would like in code?

An alternative would be to add new distributions e.g. NormalNonCentered, not sure which one is better.

[junpenglao]

I support the idea of adding a new distribution. Additional parameter is not always visible (especially for new users).

[aseyboldt]

As to parameter vs new class:
I can see advantages for both, but I still prefer the parameter. Mainly because I think it is a bit weird to have different classes that represent the exact same distribution. The class feels more like "which mathematical object does this represent" and the parameter more like "how exactly do I do this internally". Having different objects also makes it harder to see at a glance what parametrizations I could use. If it is a parameter, I can see it in the docstring.

The implementation could look something like this:

class AffineTransform(pm.distributions.transforms.ElemwiseTransform):
    name = 'affine'
    
    def __init__(self, loc, scale):
        self.scale = scale
        self.loc = loc

    def forward(self, x):
        return (x - self.loc) / self.scale

    def backward(self, y):
        return self.scale * y + self.loc
    
    def jacobian_det(self, y):
        return tt.log(self.scale)

class Normal_(pm.Normal):
    def __init__(self, mu, sd, parametrization='centered', *args, **kwargs):
        if parametrization == 'centered':
            transform = None
        elif parametrization == 'non-centered':
            transform = AffineTransform(mu, sd)
        else:
            raise ValueError('Unknown parametrization: %s' % parametrization)
        
        super().__init__(mu, sd, *args, **kwargs, transform=transform)

We could also take advantage of the fact that is easier to compute the logp in the untransformed space in this case (using inheritance instead of changing the original implementation for clarity):

class TransformedDistribution(pm.distributions.transforms.TransformedDistribution):
    def __init__(self, dist, transform, untransformed_dist, *args, **kwargs):
        super().__init__(dist, transform, *args, **kwargs)
        self.untransformed_dist = untransformed_dist

    def logp(self, x):
        if self.untransformed_dist is not None:
            return self.untransformed_dist.logp(x)
        else:
            return super().logp(x)
        
        
class UnitNormal(pm.Continuous):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.mean = self.mode = self.median = 0
        self.variance = 1

    def logp(self, x):
        return - 0.5 * (x ** 2 + tt.log(2 * np.pi))

    
class AffineTransform(pm.distributions.transforms.ElemwiseTransform):
    name = 'affine'
    
    def __init__(self, loc, scale, untransformed_dist=None):
        self.scale = scale
        self.loc = loc
        self.untransformed_dist = untransformed_dist

    def forward(self, x):
        return (x - self.loc) / self.scale

    def backward(self, y):
        return self.scale * y + self.loc
    
    def jacobian_det(self, y):
        return tt.log(self.scale)
    
    def apply(self, dist):
        return TransformedDistribution.dist(dist, self, self.untransformed_dist)


class Normal_(pm.Normal):
    def __init__(self, mu, sd, parametrization='centered', *args, **kwargs):
        if parametrization == 'centered':
            transform = None
        elif parametrization == 'non-centered':
            transform = AffineTransform(mu, sd, UnitNormal.dist())
        else:
            raise ValueError('Unknown parametrization: %s' % parametrization)
        
        super().__init__(mu, sd, *args, **kwargs, transform=transform)

[fonnesbeck]

I would vote for a subclass of Normal, rather than an argument. It would make things clearer, I think. NonCenteredNormal or similar would be great.

[AustinRochford]

I am with @fonnesbeck on this one and think it's a great idea.

[bwengals]

A lack of a non-centered MvNormal is a bit of a blocker for me on the gp stuff at the moment. I wouldn't mind taking this on. Has anyone been working on this?

[aseyboldt]

Not really. I've been working a bit on a refactoring of the distribution base class and the transformations, and I'm hoping that this will pretty much write itself with that. But I don't know if this refactoring will ever see the light of day.

[junpenglao]

I suggest the following:

Create a new class in pymc3/distributions, and called it pymc3/distributions/reparameterized_dist. Inside it would be a collection of distribution such as noncentered Normal (also reparameterization for eg horseshoe and Cauchy). This class of distribution will not take observed kwarg.

I still firmly believe they should be a new class of distribution instead of just a kwarg inside of the Distribution we already have. As different parameterization is actually a different model and should be treated differently.

[twiecki]

This is still a cool idea.

[bridgeland]

Super-useful. But don't forget NonCenteredTruncatedNormal.

Even @aseyboldt's original proposal is useful, as it allowed me to check my parameterization of a non-centered lognormal.

[twiecki]

@bridgeland Want to do a PR with @aseyboldt's code? Would be super useful.

[bridgeland]

Indeed. But I am not the right guy, at least not yet. I am still a pymc3 noob.

[twiecki]

@bridgeland Fair enough, but note that this is the best way to learn and we'd help getting it through.

[canyon289]

Per meeting discussion symbolic pymc has this as its use case but there still might be value in having this natively in PyMC

[twiecki]

What could we do here now with aeppl? CC @ricardoV94 @brandonwillard

[brandonwillard]

We can do reparameterizations like that with only built-in Aesara features now.

[twiecki]

Is there a code example I could look at? Ideal would be centered <> non-centered.

[brandonwillard]

No, we need developers for these kinds of things.

[ricardoV94]

This was always an API issue, not a real backend limitation. As @aseyboldt showed, this could be easily achieved with an Affine transform, which is easy enough to implement.