Training a Neural Network on standard deviation

Question

Right now I have a training dataset that kinda looks like this (Image, Float) array, where the Image is independent variable and the float is the dependent variable:

(Dog1, 3.8)
(Dog2, 0.2)
(Dog3, 2.3)
(Dog4, 4.1)
(Dog5, 1.8)
(Dog6, 5.2)
(Cat1, -1.01)
(Cat2, -1.12)
(Cat3, -0.93)
(Cat4, -0.98)
(Cat5, -0.95)

In this case, the dogs average 2 with stddev 3. And the cats avg -1 with stddev 0.05.

I'm trying to get my Neural Network to train on this data. Where, once it's done training, it can take an Input Image, and Output a float tuple, representing the mean and stddev of what the dependent variable is likely to be. i.e., The Neural Network, when it sees a dog, should output (2, 3). And when it sees a Cat, should output (-1, 0.05).

In essence, the Neural Network needs to receive an image, and figure out both what the float is likely to be on average, and how far the float might be from that average (By outputting stddev). The float has a 0% chance of being the average itself.

~

I know that training the output variable "mean" on the 2nd value in the tuple, the NN will train to figure out how to predict the average value. But, my main problem is, I'm just not sure what to train the 2nd output of the NN on, so that the 2nd output learns the stddev (i.e., in some sense, the 2nd output is also a kind of "confidence value" for the 1st output).

Answer 1

I've found an answer here: https://medium.com/deeplearningmadeeasy/how-to-add-uncertainty-to-your-neural-network-afb5f855e66a

Notably, let's say for a given (Image, Val) your NN outputs (Mean, Stddev). You want your loss function to be:

distr = normal.distribution(mean=Mean, stddev=Stddev);
loss = -distr.log_prob(Val);

Meaning, figure out the probability P that a number Val or worse occurs in the distribution of (Mean, Stddev). The loss is -log(P). If P is very low, -log(P) is high. If P is very high, -log(P) is low. It will learn to make the probabilities accurate.

It's not clear if this harms the NN's ability to learn the mean, because it's not using MSE for learning the mean anymore.