基于概率神经网络的异方差不确定性估计

摘要

捕获任意不确定性是许多机器学习系统的关键部分。在深度学习中，达到这一目的的一种常用方法是训练神经网络，通过最大化观测数据下似然函数的对数来估计异方差高斯分布的参数。在这项工作中，我们研究了这种方法，并确定了与使用对数似然与基于梯度的优化器相关的潜在危险。首先，我们给出一个综合的例子，说明这种方法如何导致非常差但稳定的参数估计。其次，我们确定罪魁祸首是对数概率损失，以及某些加剧问题的条件。第三，我们提出了一种替代公式，称为β- NLL，其中每个数据点对损失的贡献由β-指数方差估计加权。我们表明，在我们的说明性示例中，使用适当的β在很大程度上缓解了这个问题。第四，我们在一系列领域和任务上评估了这种方法，并表明它在预测RMSE和对数似然标准方面取得了相当大的改进，并且在超参数方面表现得更加稳健。

作者：Maximilian Seitzer, Arash Tavakoli, Dimitrije Antic, Georg Martius
题目：On the Pitfalls of Heteroscedastic Uncertainty Estimation with Probabilistic Neural Networks
代码：https://www.catalyzex.com/paper/on-the-pitfalls-of-heteroscedastic
ICLR2022 poster: https://openreview.net/forum?id=aPOpXlnV1T

关注微信公众号，获取更多资讯内容：
在这里插入图片描述

1 介绍

赋予模型捕捉不确定性的能力在机器学习中是至关重要的。不确定性可以分为两种主要类型:认知（epistemic）不确定性和任意（aleatoric）不确定性(Kiureghian & Ditlevsen, 2009)。认知不确定性解释了模型中的主观不确定性，这种不确定性在给定足够数据的情况下是可简化的。相比之下，任意不确定性捕获了观测中固有的随机性，并且本身可以细分为同均方差不确定性和异方差不确定性。同方差不确定性对应于在整个输入空间中恒定的噪声，而异方差不确定性对应于随输入空间变化的噪声。

对每种不确定性建模都有公认的好处。例如，捕获认知不确定性可以在主动学习中有效地收集预算数据(Gal等人，2017)，允许在强化学习中进行有效的探索(Osband等人，2016)，并且在成本敏感的决策中不可或缺(Amodei等人，2016)。另一方面，量化任意不确定性可以学习随机过程的动态模型(例如，用于基于模型或离线强化学习)(Chua等人，2018;Yu等人，2020)，提高了语义分割、深度回归和目标检测的性能(Kendall & Gal, 2017;Harakeh & Waslander, 2021)，并允许风险敏感决策(Dabney等人，2018;Vlastelica et al.， 2021)。

我们研究了神经网络回归中量化任意不确定性的一种常用方法。通过假设回归目标遵循特定的分布，我们可以使用神经网络来预测该分布的参数，通常是假设异方差高斯分布时的输入相关均值和方差。然后，可以使用最大似然估计(MLE)来学习网络的参数，即使用随机梯度下降最小化负对数似然(NLL)准则。这个简单的程序，这是事实上的标准(Nix & Weigend, 1994;Lakshminarayanan等人，2017;Kendall & Gal, 2017;Chua等人，2018)，已知会受到过度自信的方差估计的影响。鉴于已经提出了缓解这一特定问题的策略(Detlefsen等人，2019;Stirn & Knowles, 2020)，我们认为，同样重要的问题是，这一过程还可能导致均值拟合欠佳。在这项工作中，我们分析并提出了一个简单的修改来缓解这个问题。
在这里插入图片描述

图1:训练一个概率神经网络来拟合一个简单的正弦信号是失败的。左: $10^7$ 次更新后的学习预测(橙线)，阴影区域显示预测的标准差。目标函数为y(x)= 0.4sin(2πx)+ ξ，其中ξ为标准差为0.01的高斯噪声。右:训练的均方根误差(RMSE)，超过10个随机种子的均值和标准差。为了进行比较，我们绘制了训练曲线，使用均方误差作为训练目标- 在 $10^5$ 次更新中实现最佳平均拟合(虚线)。这种行为在不同的优化器、超参数和体系结构中是稳定的(参见章节B.2)。

我们展示了优化神经网络回归的NLL损失的一个缺陷，它阻碍了准确均值预测器的训练(见图1的说明性示例)。其主要原因是梯度对预测方差的高度依赖。虽然这种依赖性通常被认为是导致均值和方差估计器联合优化不稳定的原因(Takahashi等人，2018;Stirn & Knowles, 2020)，我们从一个新的角度来看待这种依赖如何进一步成为问题。也就是说，我们假设这个问题是由于NLL损失降低了预测差的数据点相对于预测好的数据点的梯度，导致对预测差的数据点进行了有效的欠采样。

然后，我们引入了一种替代损失公式，称为β- NLL，通过其β指数方差估计加权每个数据点对总体损失的贡献来抵消这一点，其中β控制梯度对预测方差的依赖程度。该公式包含了β = 0时的标准NLL损失，并允许减少梯度对0 < β≤1的方差估计的依赖。有趣的是，使用β = 1完全消除了训练均值估计器的这种依赖性，产生标准均方误差(MSE)损失-但具有不确定性估计的额外能力。最后，我们的经验表明，我们改进的损失公式在很大程度上减轻了拟合不良的问题，在一系列领域和任务上取得了相当大的改进，同时对超参数配置表现出更强的鲁棒性。

2 预热

设X, Y为描述输入和目标的两个随机变量，服从联合分布P(X, Y)。我们假设在给定X的情况下，Y是条件独立的，并且它遵循某个概率分布P(Y | X)。在下文中，我们使用一个常见的假设，即给定X, Y是正态分布的，即 $P(Y | X)= N(µ(X)，σ^2(X))$ ，其中 $µ,σ^2$ 分别是真实的依赖于输入的均值和方差函数。同样,我们可以写 $\epsilon(X), \sim \epsilon(X)∼N(0,σ^2 (X))$ 。Y由X通过µ(X)加上方差为 $σ^2(X)$ 的零均值高斯噪声生成。这种输入依赖的方差量化了异方差不确定性或输入依赖的任意不确定性。

为了学习真实均值和方差函数的估计μ’ (X)， σ’2(X)，通常使用由θ参数化的神经网络f_θ。在这里，µ(X)和σ^2(X)可以是最后一层的输出(Nix &Weigend, 1994)，或者使用两个完全独立的网络(Detlefsen et al.， 2019)。因此，使用合适的激活函数(例如softplus)将方差输出约束到正区域。然后可以使用最大似然估计(MLE)，通过最小化分布P(X, Y)下的负对数似然(NLL)准则LNLL来找到最优参数 $θ^*_{NLL}$ :
在这里插入图片描述
相比之下，标准回归最小化均方误差(MSE):

在实践中，Eq. 1和Eq. 2使用随机梯度下降(SGD)对从P(X, Y)中抽取的小批样本进行优化。LNLL w.r.t.(相对于) $\hat{µ}(X)， \hat{σ}^2(X)$ 的梯度由式给出

3 分析

现在我们回到第1节中尝试拟合正弦函数的例子。回顾图1，使用高斯NLL作为目标导致次优拟合。以MSE为目标，模型在合理的时间内收敛到最优均值拟合。现在我们来分析一下这一惊人结果背后的原因。

从Eq. 3中，我们可以看到真正的平均值µ(X)是NLL损失的最小值。因此，很明显，a)图1中的解不是最优解，b) NLL目标原则上应该将μ (X)驱动到最优解µ(X)。那么，为什么模型不收敛到最优解呢?我们确定了高斯NLL目标的这种行为的两个主要罪魁祸首:

特征空间的初始平坦性可以产生一个不复杂但局部稳定的均值拟合。这种拟合是由局部对称性造成的，需要某种形式的对称性破缺才能逃脱。
相对于预测良好的点，NLL损失将预测不好的点的梯度缩小，有效地对这些点进行欠采样。这种影响随着训练的进行而恶化。

这些罪魁祸首及其对训练的影响如图2(左)所示。如果网络由于其特征空间(由最后一个隐藏层跨越)过于粗糙而无法拟合某个区域，那么它将感知到一个高效的数据方差。这导致来自这些地区的数据权重降低，从而加剧了权重日益失衡的自我放大的恶性循环。下面，我们将更详细地分析这些影响及其原因。
在这里插入图片描述

图2:与我们的解决方案相比，使用NLL(负对数似然)进行训练时的陷阱说明。初始的非均匀特征空间粒度(见第3.1节)在早期会导致不同的拟合质量。NLL中平方误差的隐式加权可以看作是p(x)∝σ2(x)的有偏数据抽样(见Eq. 6)。在训练过程中，拟合不良的部分越来越被忽略。在右边，1显示了我们的解决方案(Eq. 7)对数据点的相对重要性的影响。
在这里插入图片描述

在这里插入图片描述

图3:不同训练阶段NLL损失的模型拟合(橙色)，不确定性带为±σ。黑点表示训练数据。函数的拟合从左边开始，显然很慢。

3.1对称性和特征非线性

从图3可以看出模型在训练过程中是如何演变的。网络首先学习本质上最好的线性拟合，同时调整方差以匹配残差。局势是局部稳定的。也就是说，由于误差低于和高于均值拟合的对称性，没有动机去改变这种情况。进一步的研究需要对称性破缺。对称性破坏的一种形式来自于SGD中小批量抽样的固有随机性，或者数据集中包含的自然不对称性，例如，异常值。此外，我们假设特征空间的局部非线性在创建必要的非线性拟合中起着重要作用。

让我们考虑特征空间的非线性。这个数量不容易捕捉。为了对数据集D进行近似，我们计算特征f(x) w.r.t.的雅可比矩阵Jf在半径为r的l2球中围绕点x的变化，表示为雅可比方差：
在这里插入图片描述
图4显示了输入空间上的雅可比方差作为训练进度的函数。虽然最初是相对平坦的，但它在输入空间的某些部分变得更细粒度，这些部分后来很好地配合。具有低雅可比方差的区域仍然停留在这种配置中(见图1)。这证明了特征空间的非线性对该数据集学习的成败至关重要。但是，为什么梯度下降法没有摆脱这种情况呢?
在这里插入图片描述

图4:雅可比方差随训练时间的变化，使用矩阵V(x)的均值(见Eq. 5)。图5:在训练时间内，在输入x处采样数据点的概率。

3.2逆方差加权有效欠样本

答案在于输入空间中数据点的权重不平衡。回想一下，NLL w.r.t.的梯度 $_µL_{NLL}$ 的平均值缩放误差µ(X) ?当对称性被打破，真实函数开始在局部拟合时，这些区域的方差迅速缩小，以匹配减少的MSE。如果方差校准得很好，梯度变为 $\hat{μ} (X)-Y by \frac{1}{\hat{\sigma^2(X)}}$ 具有低误差的数据点在批梯度中的贡献将相对于高误差的数据点进行放大——“越富越富”的自我放大。因此，NLL与专注于高误差样本的MSE相反。如果良好拟合区域上的真实方差σ^2远远小于不良拟合区域上的误差，或者良好拟合点比不良拟合点多，则在不良拟合区域上的学习进度完全受阻。

另一种看待这一点的方法是将点的不同权重解释为将训练分布P(X, Y)更改为修改后的分布P’(X, Y)，其中误差高的点被采样的概率较低。这可以通过定义 $P'(X, Y)= Z^{-1}P(X,Y)/ σ^2(X)$ 其中 $\int P(x,y)/σ^2(x) dxdy$ 是一个归一化常数，并且认识到在修改后的数据分布P’(x,y)下，NLL的梯度与Eq. 2中MSE损失的梯度成正比:
在这里插入图片描述
在图5中，我们绘制了正弦示例在训练时间上的P’(X, Y)。可以看出，从高误差区域采样点的虚拟概率随着时间的推移而下降，直到极不可能从该区域采样点( $10^{-5}$ 比 $10^{-3}$ 为均匀取样)。我们在第B.3节中展示了这种行为也延续到现实世界的数据集。

有时，“反方差加权”被视为高斯NLL的一个特征(Kendall & Gal, 2017)，它通过允许网络“忽略”具有高误差的离群点来引入自正则化特性。如果预测的方差对应于数据固有的不可预测性(噪声)，这可能是可取的，但如果它导致过早收敛和忽略难以拟合的区域，则是不可取的，如上所示。在我们的方法中，我们可以控制自正则化的程度。

4 方法

我们想要解决的问题是NLL训练对高度次优均值拟合的过早收敛。在第3节中，我们确定了NLL损失中不良拟合数据点的相对降权以及其自放大特性是罪魁祸首。实际上，NLL用1/σ^2对每个数据点的均方误差进行加权，这可以解释为 $P(x)∝1/σ^2$ 的抽样数据点。因此，我们建议通过引入参数β来修改这种分布，允许在NLL和完全统一的数据点重要性之间进行插值。得到的抽样分布由 $P(x)∝σ^{2β}/σ^2$ 给出，如图2(右)所示。

如何实现这种权重?我们简单地将方差加权项 $σ^{2β}$ 引入LNLL损失使其成为梯度上的一个因素。我们将由此产生的损失记为:
在这里插入图片描述
|_ _|表示停止梯度操作。通过停止梯度，方差加权项作为自适应的、依赖于输入的学习率。这样， $L_{β-NLL}$ 的梯度为:

自然地，当β =0时，我们恢复原来的NLL损失。当β =1时，Eq. 8中的梯度w.r.t. µ与MSE的梯度相等。但是，对于方差，Eq. 9中的梯度是一个新的量，分母是 $2σ^2$ 。对于0 <β < 1的值，我们得到不同的损失插值。特别有趣的是β =0.5的情况，其中数据点用1/σ(逆标准差而不是逆方差)加权。在我们的实验(第5节)中，我们发现β =0.5通常可以在准确性和对数似然之间实现最佳权衡。第D.5节提供了损失函数的Pytorch实现。

注意新的损失 $L_{β-NLL}$ 不是用于性能评估，而是用于产生有意义的梯度。由于权重项的存在，损失值不能反映模型的质量。训练期间的模型性能应该用原始的负对数似然目标来监控，也可以选择使用RMSE来测试均值拟合的质量。

在这里插入图片描述

图6:ObjectSlide数据集依赖于损失函数的残差预测误差分布(参见5.2节)。虚线表示预测RMSE。(a) NLL损失产生多模态残差。困难的数据点会被忽略，而容易的数据点则适合较高的精度。(b) MSE损失(LMSE)导致残差呈对数正态分布。©我们的β?NLL损失(Lβ?NLL)在简单的数据点上产生高度精确的拟合，而不会忽略困难的数据点。

5 实验

在我们的实验中，我们提出了以下问题，并得出了以下结论:

第5.1节:β-NLL修复陷阱与NLL的收敛?
第5.2节:β-NLL在实际环境中优于NLL ?β - NLL对超参数有多敏感?我们研究了不同的现实世界领域:UCI数据集的回归，动态模型学习，MNIST和Fashion-MNIST的生成建模，以及自然图像的深度图预测。是的,β-NLL通常比NLL性能更好，而且更容易调优。

第5.3节:β-NLL与分布回归的其他损失函数比较?我们比较了一系列方法:学习匹配高斯矩(称为“矩匹配”(MM));参见A节)，使用Student 's t分布而不是aGaussian (Detlefsen等人，2019)，或者在方差上放置不同的先验并使用变分推理(xVAMP, xVAMP*， VBEM, VBEM*) (Stirn & Knowles, 2020)。视情况而定。不同的损失产生不同的权衡，我们将在第5.3节中讨论。

5.1合成数据集

一个带有小附加噪声的正弦曲线：我们首先对图1中不具说明意义的例子进行了扩展研究:y =0.4sin(2πx)+ ξ，其中ξ为标准差σ =0.01的高斯噪声。人们会期望一个有足够容量的网络可以很容易地学会适应这个函数。图7在一系列体系结构和学习率上检查了这一点。

我们发现，对于标准NLL损失(β =0)，网络不会收敛到合理的均值拟合。有一种趋势是，更大的网络和学习率显示出更好的结果，但当将其与β?当β =0.5时，我们看到网络确实能够没有任何问题地拟合函数。正如预期的那样，均方误差损失(MSE)和β =1的NLL。与NLL的拟合质量如图S3所示。

具有异方差噪声的正弦信号，我们对β?NLL仍然对Detlefsen等人(2019)的示例提供了良好的不确定性估计——一个振幅和噪声增加的正弦曲线:y = x sin(x)+ xξ1 + ξ2，其中ξ1和ξ2是标准差σ =0.3的高斯噪声。图8 (a-e)显示了最佳模型的预测结果(w.r.t. NLL验证损失)，(f)比较了10个独立试验的预测不确定性。在所有损失的情况下实现均值的拟合。在训练场上，β?β> 0的NLL学习几乎相同的不确定性
在这里插入图片描述

图7:正弦回归问题的收敛性分析。20万次epoch后的RMSE，在3次独立试验中平均，用颜色代码(越亮越好)作为学习率和模型架构的函数显示(见D.1节)。原始NLL (β =0)对大多数超参数设置不能获得良好的RMSE拟合。图S3显示了NLL度量的结果。

在这里插入图片描述

图8:适合Detlefsen等人(2019)的异方差正弦示例(a-e)。虚线分别表示真实均值和±2σ。(f)使用相同颜色代码的预测标准差(在10个独立试验中用阴影表示的标准差)。请注意，β =0.5和β =1的图形位于彼此的顶部。在训练制度中，所有β?NLL变体(a-d)产生校准良好的不确定性估计。矩匹配(e)显著低估了所有地方的方差。

5.2真实数据集

UCI Regression Datasets：作为预测不确定性估计的标准现实基准，我们考虑UCI数据集(Hernandez-Lobato & Adams, 2015)。表1概述了对不同损失变量的比较。有关所有12个数据集的完整结果，请参阅第B.4节。结果令人鼓舞:β?NLL实现了与NLL损失相当或更好的预测对数似然，同时在大多数数据集上明显提高了预测精度。
在这里插入图片描述

表1:UCI回归数据集的结果。我们报告了预测性对数似然和RMSE(±标准差)。Ties表示该方法在统计上无法与最佳方法区分的数据集的数量(12个数据集中)(参见第B.4节)。我们比较了Student-t(Detlefsen等人，2019)和xVAMP/VBEM (Stirn & Knowles, 2020)。B.4节列出了全部结果。

Dynamics models：由于不确定性估计的主要应用在于基于模型的强化学习(RL)，我们在两个不同难度的动态预测任务(ObjectSlide和Fetch-PickAndPlace)上测试了不同的损失函数。在这两个任务中，目标都是预测对象将如何从当前状态和智能体的动作中移动。ObjectSlide (Seitzer et al.， 2021)是一个简单的1d环境，Fetch-PickAndPlace (Plappert et al.， 2018)是一个复杂的3D机器人操作环境。这些模型是在RL代理收集的轨迹上训练的。

对于这两个数据集，我们在不同的超参数配置上执行网格搜索(参见第D.2节)，以便对超参数设置进行灵敏度分析，如图9所示。结果表明，NLL易受超参数选择的影响，而β?NLL在各种配置下都能取得良好的效果。然后在保留测试集上评估每种损失的最佳性能配置(表2)。可以看到，NLL损失导致较差的预测性能，并且在随机种子之间也表现出相当高的方差。我们的方法准确度高，对数似然拟合β值的范围，β =0.5通常达到最佳折衷。
在这里插入图片描述

图9:动态模型学习任务上损失函数对超参数的敏感性分析:ObjectSlide (a)和Fetch-PickAndPlace (b)。基于不同模型配置的网格搜索，验证RMSE和NLL上的分布显示为超参数的函数(参见第D.2节)。当评估RMSE时，NLL损失是高度敏感的，β?NLL损耗显示出更低的灵敏度，并产生良好的结果，而不管确切的配置。

生成建模和深度图预测: 对于生成建模，我们在MNIST和Fashion-MNIST上使用概率解码器训练变分自编码器(Kingma & Welling, 2014)。对于深度回归的任务，我们修改了最先进的方法(AdaBins;Bhat等人(2021))，并在NYUv2数据集(Silberman等人，2012)上使用我们的损失(图S6)进行测试。表3给出了这两个任务的选择结果，产生了与以前相似的趋势。我们参考章节B.5和章节B.6了解更多细节，包括定性结果。

6 结论

我们强调了在使用常见的NLL损失优化概率神经网络时经常出现的问题:训练陷入次优函数拟合。通过我们的分析，我们揭示了潜在的原因:最初的不良拟合区域在损失中得到的权重越来越小，从而导致过早收敛。我们提出了一个简单的解决方案，通过引入一组称为β-NLL的损失函数。有效地，原始NLL损失的梯度通过β指数的每样本方差进行缩放。这允许在NLL和MSE损失函数之间进行有意义的插值，同时提供良好的不确定性估计。超参数β为从业者提供了控制NLL自正则化强度的选择:高噪声区域或难以预测的数据点在拟合过程中的重要性。在大多数情况下，β =0.5将是一个很好的起点。我们认为本文讨论的问题主要是为什么在回归或生成建模任务中使用高斯分布的实践者经常选择常数或均方差(全局)方差，而不是更一般的异方差(数据依赖)方差。我们希望我们的简单解决方案能够通过提高深度神经网络建模数据不确定性的可用性和性能来改变这种情况。