A model that learns its own uncertainty chooses what to fit

Asking a regression network for its own error bars sounds like pure gain. Instead of one output per prediction it produces two, a mean $\mu$ and a standard deviation $\sigma$, and it is trained to maximise the likelihood of the data under $\mathcal{N}(\mu, \sigma^2)$. Kendall and Gal (2017)¹ made this the standard way to model the noise inherent in the data, the aleatoric uncertainty, and it costs one extra output channel.

I had used this loss without thinking hard about it. Then I wrote down its gradient and noticed that the uncertainty doesn't only sit next to the prediction: it decides how much each data point is allowed to change the prediction. I wanted to see what that does in practice.

The gradient on paper

The negative log-likelihood of one observation $y$, dropping constants, is

Its gradients with respect to the two outputs are

The second one is the honest part. It vanishes when $\sigma^2$ equals the squared error, so at the optimum the predicted variance matches the actual one.

The first one is where the surprise is. It is the squared-error gradient divided by $\sigma^2$. A point the model believes is precise gets its error amplified; a point it believes is noisy gets its error discounted. If one region of the input has a noise level of 0.005 and another 0.6, and the model learns both, the gradient on the mean differs between them by a factor of $(0.6 / 0.005)^2 = 14{,}400$. The loss is not one regression problem with error bars attached. It is a weighted regression, and the model chooses its own weights.

Two predictions follow. The mean should fit quiet regions very well and noisy regions poorly. And since the weight $1/\sigma^2$ is unbounded as $\sigma$ shrinks, the gradients should be large and spiky, which suggests trouble at higher learning rates.

The experiment

I built a one-dimensional task that makes the weighting visible: $y = \sin(4x) + 0.5x$ plus Gaussian noise whose standard deviation is 0.005 on the left third of the interval, 0.1 in the middle and 0.6 on the right. The network is a small multilayer perceptron with two hidden layers of 128 units and two outputs. I trained it for 6,000 steps with Adam at learning rates $10^{-3}$ and $10^{-2}$, three seeds each, four ways:

1. squared error on the mean, no uncertainty;
2. the Gaussian negative log-likelihood with $\log\sigma$ as the raw output;
3. the same with a floor, $\sigma = 0.02 + \mathrm{softplus}(\cdot)$;
4. β-NLL with β = 0.5 (Seitzer et al., 2022²), described below.

I measured the error of the mean against the true function in each region, the largest gradient norm seen in training, and how often the observation fell within one predicted standard deviation, which should be about 68%.

Where the mean gets fitted

At learning rate $10^{-3}$:

Errors are mean absolute differences from the true function, averaged over three seeds.

The weighting shows up just as the gradient predicts. On the quiet region the likelihood loss fits the mean with a tenth of the error of squared error, 0.0026 against 0.026. On the noisy region it does almost twice as badly, 0.088 against 0.047. Squared error treats every point alike and spreads its effort evenly. The likelihood loss pours its effort into the points it trusts.

The gradient norms tell the second half of the story. The largest gradient under the likelihood loss was more than a thousand times the largest under squared error. Every time a quiet point is predicted slightly wrong, its error arrives multiplied by $1/\sigma^2$, and $\sigma$ there is around a hundredth.

Stable, calibrated and wrong

At learning rate $10^{-2}$ the plain likelihood loss did something I hadn't expected. It never produced a NaN, and its loss stayed finite in all three seeds. But the final model was useless: the error of its mean was about 7 in every region, on a function whose values stay between −1.5 and 1.5. Its predicted standard deviation was about 9. By its own standard it was well calibrated, with 82% to 99% of observations inside one predicted standard deviation.

The second gradient explains how it got there. After a large gradient spike knocked the mean off course, the cheapest way to lower the loss was no longer to fix the mean, which the spiky gradients made hard. It was to raise $\sigma$ until $\sigma^2$ matched the now enormous squared error, which makes $\partial\mathcal{L}/\partial\log\sigma$ zero. Once $\sigma$ is large, the gradient on the mean is divided by $\sigma^2 \approx 80$ and nearly stops. The model settles into a state where it is wrong and says so, and the loss rewards it for that. Seitzer et al. (2022) describe the same failure: likelihood training that leads to "very poor but stable parameter estimates".

Two ways to take back control

A floor on σ caps the weight $1/\sigma^2$ at $1/0.02^2 = 2{,}500$. At $10^{-3}$ that cut the largest gradient tenfold, and the noisy-region mean was the best of any loss, 0.036. The price is visible in the quiet region. The floor is four times the true noise there, so 97% of quiet observations fell inside one predicted standard deviation instead of 68%. A floor is a statement about the noise, and here it was the wrong statement for one region. At $10^{-2}$ the floor prevented the collapse, but the fit was poor everywhere.

β-NLL attacks the weighting directly. Seitzer et al. multiply each point's loss by $\sigma^{2\beta}$, treated as a constant so that no gradient flows through it. The gradient on the mean becomes

At β = 0 this is the likelihood loss; at β = 1 the mean is fitted by squared error while $\sigma$ is still learnt. With β = 0.5 the weight is $1/\sigma$, so the 14,400-fold spread between quiet and noisy points shrinks to 120. In my runs it was the only likelihood loss whose largest gradient stayed close to squared error's, about 10, and the only one that held up at the higher learning rate:

Its coverage at $10^{-3}$ was 77% in the quiet region and 68% to 69% elsewhere, close to the 68% it should be.

What I take from it

The variance head is not a passive report. Under the likelihood loss it sets the learning rate of every data point, and it does so with no upper bound. That has three consequences I would now design around:

• Plain likelihood training fits the points it trusts and neglects the rest. If the noisy region matters, squared error, a floor or β-NLL serves the mean better.
• Gradient norms under this loss can be orders of magnitude above those of squared error, so a learning rate tuned on squared error does not transfer.
• A finite, falling loss and good calibration are not evidence of a good model. The collapsed runs had both. The error of the mean against held-out data has to be watched separately.

The same loss, as one term among two in a shared network, is what produced the spike in taking a gradient spike apart by loss term and by module. A confidence head trained with cross-entropy has the opposite property, a well-defined target it converges to, in a confidence head learns the target you give it.

What I still don't know

My task is one-dimensional with sharp boundaries between noise levels, which makes the weighting easy to see and may exaggerate it. In a high-dimensional problem, quiet and noisy examples share far more parameters, and I don't know how much the reweighting of one leaks into the fit of the other.

I tried one floor value and one β. The floor was too high for the quiet region by design; a floor learnt from data, or one set per output, might keep the stability without the miscalibration. And I only used Adam. Its per-parameter normalisation probably softens some of the spikes, and I would like to know whether SGD makes the collapse more or less likely.

References

1. What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision? — Alex Kendall and Yarin Gal, NeurIPS 2017
2. On the Pitfalls of Heteroscedastic Uncertainty Estimation with Probabilistic Neural Networks — Maximilian Seitzer, Arash Tavakoli, Dimitrije Antic and Georg Martius, ICLR 2022