Loss Function for Probability Regression

Question

I'm trying to predict a probability with a neural network, but having trouble figuring out which loss function is best. Cross entropy was my first thought, but other resources always talk about it in the context of a binary classification problem where the labels are $\{0, 1\}$, but in my case I have an actual probability as the target. Is one of these options clearly best, or maybe are they all valid with just minor differences around the extreme 0/1 regions?

Assuming $x$ is the output of the final layer of my model.

Cross Entropy:
$target * -log(sigmoid(x)) + (1 - target) * -log(1 - sigmoid(x))$

Mean Squared Error with Sigmoid:
$(sigmoid(x) - target)^2$

Mean Squared Error with Clamp:
$(x - target)^2$

When I use the output I clamp the values between $[0, 1]$.

Answer 1

Given that you are trying to predict a scalar probability value, the cross-entropy formula you listed in the question is only valid if the target variable is discrete. So if your question is "predicts the odds of drawing a matching card from a deck", it will be fine.

The main difference between cross entropy and MSE will be how they penalize wrong predictions. Let's say given a target of 1 but prediction with 0. The cross-entropy is actually undefined in this case, but as the prediction gets closer to 0, the cross-entropy loss gets exponentially larger. On the other hand, your MSE is only 1. Which one is better, it depends on your application, if you want to avoid large margin of errors, it would seem the cross-entropy is more appropriate and vice versa.

The second and third approach only differs in how they make sure the prediction is within [0, 1], one uses a sigmoid function and another uses a clamp. Given you are using a neural network, you should avoid using the clamp function. The clamp function is the same as the identity function within the clamped range, but completely flat outside of it. So the gradient of that function is 1 within the clamp range but 0 outside of it. Because of this, you are more likely to run into the dead neural problem in the same way when people talk about "dead relu".

Answer 2

In the case when the domain of the output variable is discrete or continuous, and not binary, you could also consider the KL-divergence as the loss to compute the distance between the predicted probability distribution and the target one.

e.g. https://stats.stackexchange.com/questions/357963/what-is-the-difference-cross-entropy-and-kl-divergence

This is important also becasue in some implementations, for efficiency reasons, the cross entropy only get indexes as target variables and not distributions. This is the case with pytorch, e.g. see https://discuss.pytorch.org/t/cross-entropy-loss-implementation/43592