Aleatoric uncertainty losses

[jimfleming]

This paper discusses two loss functions for regression and classification losses which capture aleatoric uncertainty: https://arxiv.org/abs/1703.04977

def aleatoric_regression_loss(dist, labels, weights):
    losses = (labels - dist.loc)**2 / (2 * dist.scale**2)
    losses += (1 / 2) * log(dist.scale**2)
    loss = compute_weighted_loss(losses, weights)
    return loss

def aleatoric_classification_loss(dist, labels, weights):
    # requires sampling
    return loss

There are lots of ways of formulating these and maybe even something more direct using distributions. Open to suggestions.

[wenkesj]

I really like this. It seems like a nice addition on top of the tf.losses API. Should we adjust it to take the tfp.distributions.Distribution parameters directly instead? This way there isn't any assumption about the object containing the attributes and pass the loc and scale (or samples in other cases). It may not be necessary in any case, since this assumes it is a Gaussian, therefore should come from the already existing tfp API like your example.

To illustrate what I was suggesting:

def aleatoric_regression_loss(predictions_loc, predictions_scale, labels, weights):
    losses = (labels -predictions_loc)**2 / (2 * predictions_scale **2)
    losses += (1 / 2) * log(predictions_scale**2)
    loss = compute_weighted_loss(losses, weights)
    return loss
dist = tfp.distributions.MultiVariateNormalDiag(...)
loss = aleatoric_regression_loss(dist.loc, dist.scale, labels, weights)