贝叶斯深度学习中的子网络推断

Bayesian Deep Learning via Subnetwork Inference

贝叶斯深度学习中的子网络推断

https://proceedings.mlr.press/v139/daxberger21a/daxberger21a.pdf

摘要
贝叶斯范式有望解决深度神经网络的核心问题，如校准性能差和数据利用效率低。然而，将贝叶斯推断扩展至大规模参数空间通常需引入强约束性近似。本文指出：仅对模型权重的一小部分子集进行推断，即可获得准确的预测后验分布；其余权重则保持为点估计。该子网络推断（subnetwork inference）框架使我们得以在子集上使用表达能力更强、原本难以处理的后验近似方法。具体而言，我们实现了一种简洁、可扩展的贝叶斯深度学习方法——子网络线性化拉普拉斯近似（subnetwork linearized Laplace）：首先获得全网络权重的最大后验（MAP）估计，随后基于线性化拉普拉斯近似，在选定子网络上推断一个全协方差高斯后验分布。我们提出一种子网络选择策略，旨在最大程度保留模型的预测不确定性。实验表明，该方法在性能上优于集成方法（ensembles）及对全网络采用表达能力较弱后验近似的其他方法。

1. 引言
   深度神经网络（NNs）的一个关键缺陷是：其预测往往校准不良且过度自信——尤其当训练与测试数据分布存在偏移时（Nguyen et al., 2015; Guo et al., 2017）。为支持可靠决策，神经网络需稳健地量化其预测不确定性（Bhatt et al., 2020），这对医疗、自动驾驶等安全攸关应用尤为重要（Amodei et al., 2016）。

贝叶斯建模（Bishop, 2006; Ghahramani, 2015）通过模型参数的后验分布，为不确定性量化提供了原则性途径。不幸的是，神经网络中精确后验推断不可行。尽管贝叶斯深度学习领域近年取得进展（Osawa et al., 2019; Maddox et al., 2019; Dusenberry et al., 2020），现有方法为适配大规模网络，仍不得不采用不切实际的假设，严重限制了后验分布的表达能力，进而损害不确定性估计质量（Ovadia et al., 2019; Fort et al., 2019; Foong et al., 2019a）。

或许，这些不切实际的推断近似可被避免。鉴于神经网络高度过参数化，其精度可由一个小型子网络良好保持（Cheng et al., 2017）；且在低维权重子空间中进行推断，即可实现准确的不确定性量化（Izmailov et al., 2019）。这引出如下问题：一个完整神经网络的模型不确定性能否被小型子网络充分保留？ 本文证明：全网络的后验预测分布可由子网络的后验预测分布良好近似。具体贡献如下：

1. 提出 子网络推断 ——一种可扩展的贝叶斯深度学习通用框架：仅对神经网络权重的一小部分子集进行推断，其余权重保持为确定性点估计；由此允许使用原本在大规模网络中难以处理的高表达力后验近似方法。我们给出该框架的一种具体实现：先拟合全网络的最大后验（MAP）估计，再以线性化拉普拉斯近似在子网络上推断全协方差高斯后验（见图1）。

2. 提出一种基于 全网络近似后验 与 子网络近似后验 之间Wasserstein距离的子网络选择策略。为提升可扩展性，子网络选择阶段采用对角近似；选定小规模子网络后，即可推断权重间的协方差。实验发现： 在子网络选择阶段做近似，对后验预测的影响远小于在推断阶段做近似 。

3. 我们在一系列不确定性校准与分布偏移鲁棒性基准上评估该方法。实验表明：高表达力的子网络推断方法，其性能优于对全网络进行低表达力推断的主流贝叶斯深度学习方法，也优于深度集成（deep ensembles）。

1. 子网络后验近似

设 ∈ ℝᴰ 为所有神经网络权重的 D 维向量（即所有层权重矩阵的拼接与展平）。贝叶斯神经网络（BNNs）旨在捕捉模型不确定性，即由于训练数据 = {, } 存在多种合理解释而产生的关于权重 选择的不确定性。其中， ∈ ℝᴼ 为输出变量（例如分类标签）， ∈ ℝᴺˣᴵ 为特征矩阵。首先，需在 BNN 的权重 上指定一个先验分布 ()。随后，我们希望推断其完整的后验分布。

这种后验预测分布将权重中的不确定性转化为预测中的不确定性。遗憾的是，由于神经网络（NNs）的非线性特性，推断精确的后验分布 (|) 是不可行的；又因权重 的高维度，即使要忠实地近似后验分布也面临巨大的计算挑战。因此，通常采用粗略的后验近似方法，例如完全因子化近似，即 (|) ≈ ∏ᴰ_{d=1} (_d)，其中 _d 是权重向量 中的第 d 个权重（Hernández-Lobato & Adams, 2015; Blundell et al., 2015; Khan et al., 2018; Osawa et al., 2019）。然而，已有研究表明，此类近似存在严重缺陷（Foong et al., 2019a,b）。

在本工作中，我们质疑广泛存在的隐含假设——即一个表达能力强的后验近似必须包含全部 D 个模型权重。相反，我们尝试仅对权重的一个小规模子集 ≪ 进行推断。以下论证支持这一方法：

1. 过参数化：Maddox 等人（2020）表明，在局部最优解附近，存在许多方向不会改变神经网络的预测结果。此外，神经网络可被大量剪枝而不牺牲测试集精度（Frankle & Carbin, 2019）。这表明，神经网络的大部分预测能力可集中于一个小规模子网络中。
2. 子模型上的推断：先前研究¹ 已提供证据表明，即使推断未在完整参数空间上进行，仍可有效。例如，Izmailov 等人（2019）和 Snoek 等人（2015）分别在权重的低维投影空间和神经网络的最后一层上执行推断。

因此，我们将上述两个想法结合起来，对公式 (1) 中的后验分布做出如下两步近似：

与权重剪枝方法的关系。注意，（4）中的后验近似可以被视为将权重的方差修剪为零。这与权重剪枝方法（Cheng et al., 2017）形成对比，后者将权重本身设置为零。即，权重剪枝方法可以被视为移除权重以保留预测均值（即保持与完整模型接近的准确性）。相比之下，子网络推断可以被视为仅移除某些权重的方差——同时保持它们的均值——以保留预测不确定性（例如，保持与完整模型接近的校准）。因此，它们是互补的方法。重要的是，通过不剪枝权重，子网络推断保留了完整神经网络的全部预测能力以保持其预测准确性。

1. 背景：线性化拉普拉斯近似

在本工作中，我们通过使用线性化拉普拉斯近似（MacKay, 1992）对权重上的后验分布进行近似，从而满足公式（4）。这是一种可处理的推断技术，近期已被证明表现优异（Foong 等，2019b；Immer 等，2020），并可事后应用于预训练模型。下面我们将在一般设定下对其进行描述。

我们将神经网络函数记为 : ℝᴵ → ℝᴼ。首先，我们定义一个关于神经网络权重的先验分布，我们选择其为完全因子化的高斯分布 () = (; , )。接着，我们寻找后验分布的一个局部最优解，也称为权重的最大后验（MAP）估计：

随后，利用在 MAP 估计点处的二阶泰勒展开对后验分布进行近似：

有趣的是，当采用高斯似然时，以广义高斯-牛顿（GGN）精度矩阵定义的高斯分布，恰好对应于将神经网络在处进行一阶泰勒展开线性化后的 真实后验分布 （Khan et al., 2019；Immer et al., 2020）。该局部线性化函数为：

这些闭式表达式颇具吸引力，因其所得预测均值与分类决策边界 与 MAP 估计所得神经网络完全一致 。
然而，存储现代神经网络（即参数维度 D 极大）权重空间上完整的 D × D
协方差矩阵在计算上是不可行的。尽管已有研究致力于开发更廉价的近似方案（例如仅存储对角元（Denker & LeCun, 1990）或块对角元（Ritter et al., 2018; Immer et al., 2020）），但这些近似均以降低预测性能为代价。

1. 线性化拉普拉斯子网络推断

我们概述以下程序，用于在子网络推断框架内将线性化拉普拉斯近似扩展至大规模神经网络模型。

1. 子网络选择

理想情况下，我们希望所选子网络诱导出的预测后验分布尽可能接近对全网络进行推断所得的预测后验分布（式11）。这种随机过程之间的差异通常通过函数空间的 KL 散度（functional Kullback–Leibler divergence）来量化（Sun 等，2019；Burt 等，2020）：

在权重空间中，我们的目标是最小化全网络精确后验分布（式1）与子网络近似后验分布（式4）之间的差异。这带来了两个挑战：首先，计算精确后验分布仍是不可行的；其次，常见的差异度量（如 KL 散度或 Hellinger 距离）对于式（4）中出现的狄拉克δ分布并未良好定义。

为解决第一个问题，我们再次借助第3节中引入的局部线性化方法。线性化模型的真实后验分布是高斯分布或近似高斯分布²：

我们通过选用平方 2-Wasserstein 距离来解决第二个问题，该度量对于支撑集不相交的分布仍有良好定义。对于全协方差高斯分布（式21）与一个全协方差高斯分布和若干狄拉克δ函数的乘积（式16）的情形，该度量具有如下形式：

表面上看，我们似乎又回到了最初试图避免的性能较差的对角假设（Ovadia 等，2019；Foong 等，2019a；Ashukha 等，2020）。然而，这里存在一个关键区别：我们是在 子网络选择阶段 做出对角假设，而非在 推断阶段 ；我们在子网络 上执行的是 全协方差推断 。在第6节中，我们将提供证据表明，在子网络选择阶段采用对角假设是合理的，原因如下：1）相较于在推断阶段做相同假设，它对预测性能的损害要小得多；2）它优于随机子网络选择。

1. 实验

我们通过实验评估子网络推断的有效性，并将其与以下方法进行比较：（1）对全网络采用表达能力较弱的推断方法；（2）深度学习中当前最先进的不确定性量化方法。我们考虑三类基准设置：
1）小规模玩具回归任务；
2）中等规模表格数据回归任务；
3）基于 ResNet-18 的图像分类任务。
更多实验结果与设置细节分别见附录 A 与附录 D。

6.1 子网络推断如何保留后验预测不确定性？

我们首先定性评估：在选定子网络上采用全协方差高斯后验所得预测分布，与以下方法所得预测分布的对比情况：
1）全网络上的全协方差高斯后验（Full Cov）；
2）全网络上的因子分解高斯后验（Diag）；
3）仅在网络最后一层上采用全协方差高斯后验（Final layer）（Snoek 等，2015）；
4）点估计（MAP）。

对于子网络推断，我们同时考虑两种子网络选择策略：第5节所述的Wasserstein策略（Wass）与均匀随机选择策略（Rand），以构建仅包含模型参数总量50%、3%和1%的子网络。在此玩具实验中，精确计算后验边际方差以指导子网络选择尚属可行。

我们的神经网络包含2个ReLU隐藏层，每层50个隐藏单元。采用同方差高斯似然函数，其噪声方差通过最大似然估计优化。我们在网络权重（不含偏置）上采用GGN拉普拉斯推断，并结合式（18）中的线性化预测分布。因此，所考察的所有方法共享相同的预测均值，便于更公平地比较其不确定性估计。

我们将全网络先验精度设为  = 3 （经验上表现良好），子网络先验精度设为 = ⋅ S/D。

我们采用 Antorán 等人（2020）提出的合成一维回归任务——输入数据形成两个分离的簇，从而可检验模型对“簇间区域”的不确定性响应（Foong 等，2019b）。结果如图2所示：

• 子网络推断在 推断更少权重 的同时，比对角高斯或仅最后一层推断更能保留全网络推断的不确定性；

• 通过捕捉权重间的相关性，子网络推断可在数据簇之间维持较高不确定性；

• 该特性在随机与Wasserstein子网络选择下均成立，但后者在子网络更小时能保留更多不确定性；

• 相较于对角拉普拉斯，其显著优势表明： 在子网络选择阶段采用对角假设，但在推断阶段转而使用全协方差高斯后验（即本文做法），显著优于直接对推断后验采用对角假设 （参见第5节）。

综上，结果表明：在精心选择的子网络上进行高表达力推断，相较对全网络采用粗糙近似，能更好地保留预测不确定性。

6.2 大型模型中的子网络推断 vs 小型模型中的全网络推断

我们首先获得每个神经网络权重的最大后验（MAP）估计及其同方差似然函数的噪声方差。随后，对每个网络执行全网络 GGN-Laplace 推断。我们还使用所提出的 Wasserstein 规则修剪每个网络的权重方差，使剩余方差数量匹配每一个较小网络的规模。我们采用对角拉普拉斯近似来廉价地估算用于子网络选择的后验边际方差。我们利用式（12）和（18）中的线性化方法计算预测分布。因此，具有相同权重数量的神经网络会产生相同的预测均值；增加所考虑的权重方差数量只会提升预测不确定性。

我们选用三个规模递增的表格数据集（输入维度、样本点数）：wine（11维，1439点）、kin8nm（8维，7373点）和 protein（9维，41157点）。我们采用其标准训练-测试划分（Hernández-Lobato & Adams, 2015），以及专为测试分布外不确定性的变体划分（Foong 等, 2019b）。具体细节见附录 D.4。对于每个划分，我们将训练数据的15%留作验证集，用于在寻找MAP估计及选择权重先验精度时进行早停。所有模型和数据集保持其他超参数固定。结果如图3所示。

我们呈现平均测试对数似然（LL）值，因其同时考虑了准确率与不确定性。当结合全网络推断时，规模更大的模型通常表现最佳，尽管 Wine-gap 和 Protein-gap 是例外。有趣的是，即使我们仅在与小型模型同等规模的子网络上进行推断，这些大型模型的表现依然最优。我们推测，这源于权重后验神经网络模型中存在大量退化方向（即冗余权重）（Maddox 等，2020）。小型模型的全网络推断会同时捕获有用与无用权重的信息；而在大型模型中，我们的子网络选择策略使我们能将更多计算资源用于建模信息丰富的权重方差与协方差。在6个数据集中有3个，我们发现：随着推断所涉及权重数量的增加，LL 值出现骤升，随后进入平台期。这种平台现象可能是因为大部分信息丰富的权重方差已被纳入模型。考虑到计算 GGN 的成本远高于神经网络训练成本，这些结果表明： 在相同计算量下，对大型模型执行子网络推断比对小型模型执行全网络推断更优 。

6.3 分布偏移下的图像分类

我们现在评估采用子网络推断的大型卷积神经网络在图像分类任务中对分布偏移的鲁棒性，并与以下基线方法进行比较：

• 点估计网络（MAP）；

• 对全网络采用表达能力较弱推断的贝叶斯深度学习方法：MC Dropout（Gal & Ghahramani, 2016）、对角拉普拉斯、VOGN（Osawa 等，2019）——三者均假设权重后验完全因子化；以及 SWAG（Maddox 等，2019）——假设后验为“对角+低秩”结构；

• 深度集成（deep ensembles）（Lakshminarayanan 等，2017）——目前被公认为深度学习不确定性量化的最先进方法（Ovadia 等，2019；Ashukha 等，2020）。

我们采用5个网络构成的集成（据 Ovadia 等建议），并对 MC Dropout、对角拉普拉斯与 SWAG 均采样16次。Dropout 概率设为 0.1；对角拉普拉斯的先验精度通过网格搜索确定为 = 4 × 10⁴。所有方法均应用于 ResNet-18（He 等，2016）：包含1个输入卷积块、8个残差块和1个线性层，共计 11,168,000 个参数。

对于子网络推断，我们采用式（19）中的线性化预测分布；并使用 Wasserstein 子网络选择策略，仅保留 0.38% 的权重，得到一个仅含 42,438 个权重的子网络——这是当前计算条件下可处理全协方差矩阵的最大规模（其大小为）。我们采用对角 SWAG（Maddox 等，2019）估算子网络选择所需的边际权重方差。我们曾尝试对角拉普拉斯，但发现其所选权重对应于在训练点上神经网络雅可比恒为零的位置（即“死亡 ReLU”）；此类权重的后验方差虽大（近似先验），但对网络输出几乎无影响。SWAG 不受此问题困扰，因其忽略了训练梯度为零的权重。子网络推断的先验精度经网格搜索设为 = 500。

为评估原则性子网络选择的重要性，我们另设一基线：均匀随机选择子网络（记为 Ours (Rand)）。我们开展以下两个实验，结果见图4：

旋转 MNIST：参照（Ovadia 等，2019；Antorán 等，2020），所有方法在 MNIST 上训练，并在逐步增大的数字旋转角度下评估其预测分布。尽管所有方法在原始 MNIST 测试集上表现良好，但当旋转角度超过 30 度时，准确率迅速下降。就对数似然（LL）而言，集成在基线方法中表现最优；而子网络推断的 LL 显著高于几乎所有基线（包括集成），唯一例外是 VOGN（表现略优）。值得注意的是，Ovadia 等（2019）亦观察到：平均场变分推断（VOGN 属于此类）在 MNIST 上表现极强，但在更大规模数据集上性能显著下降。子网络推断在分布内能做出准确预测，同时对分布外样本赋予比基线更高的不确定性。

损坏 CIFAR：同样参照（Ovadia 等，2019；Antorán 等，2020），所有方法在 CIFAR10 上训练，并在经 16 类不同损坏（每类5个强度等级）的数据上评估（Hendrycks & Dietterich, 2019）。由于局部线性化使预测均值与 MAP 一致，子网络推断在预测误差上与 MAP 网络相当；集成与 SWAG 准确率最高。然而，子网络推断的独特优势在于过自信程度最低——在所有损坏强度等级下，其对数似然均优于所有基线方法。此时 VOGN 表现较差；但这与其在 MNIST 上的优异表现看似矛盾——实则再次印证了 Ovadia 等（2019）的发现：平均场变分推断在 MNIST 上表现良好，但在更大数据集上性能下降。

此外，在两项基准测试中，随机选择子网络的表现显著劣于我们提出的 Wasserstein 选择策略，凸显了子网络选择方式的重要性。

综上，这些结果表明：子网络推断在不确定性校准与分布偏移鲁棒性方面，优于其他主流不确定性量化方法。

1. 适用范围与局限性

   多输出模型中的雅可比矩阵计算仍具挑战性。在当前主流深度学习框架中，由于采用反向模式自动微分，其计算所需反向传播次数等于模型输出数量。这使得线性化拉普拉斯方法难以应用于语义分割（Liu 等，2019）或类别数极多的分类任务（Deng 等，2009）。需注意，该问题仅限于线性化拉普拉斯方法本身；其他无此限制的推断方法仍可纳入本框架使用。

   先验精度 的选择在很大程度上决定了拉普拉斯近似的性能。我们提出的子网络先验精度更新方案依赖于对全网络已有合理参数设定。然而，由于全网络推断常不可行，目前选择 的最佳方式是直接在子网络近似上进行交叉验证。

   海森矩阵的存储需求限制了子网络权重的最大规模。例如，存储 4 万个权重对应的海森矩阵约需 6.4 GB 内存。对于现代 Transformer 等超大规模模型，可计算的子网络仅占总权重极小比例。尽管我们已证明优异性能未必依赖大型子网络（见图5），但探索更优的子网络选择策略仍是未来研究的关键方向。

   
   

1. 相关工作

贝叶斯深度学习：针对神经网络权重后验分布 p ( w ∣ D ) 的刻画已有大量研究。迄今为止，哈密顿蒙特卡洛（Hamiltonian Monte Carlo, HMC；Neal, 1995）仍是贝叶斯神经网络（BNNs）中近似推断的黄金标准。尽管其在渐近意义上无偏，但基于采样的方法难以扩展至大规模数据集（Betancourt, 2015）。因此，近年来更流行的做法是在某一近似分布族（通常为高斯分布）中寻找最优代理后验。其中最早的是 MacKay（1992）提出的拉普拉斯近似，他也同时建议使用线性化模型的后验来近似预测后验（Khan 等，2019；Immer 等，2020）。随着更大规模神经网络的普及，能捕捉权重间相关性的代理分布因计算不可行而受限；因此，绝大多数现代方法转而采用平均场假设（Blundell 等，2015；Hernández-Lobato & Adams，2015；Gal & Ghahramani，2016；Mishkin 等，2018；Osawa 等，2019），但这牺牲了模型表达能力（Foong 等，2019a）并导致实证性能下降（Ovadia 等，2019；Antorán 等，2020）。Farquhar 等（2020）曾提出：在更深网络中，平均场假设或许并不构成限制；但我们的实证结果似乎与该观点相悖。我们发现，通过降低权重空间维度来扩展那些能考虑权重相关性的近似方法（如 MacKay，1992；Louizos & Welling，2016；Maddox 等，2019；Ritter 等，2018），其性能优于对角近似。由此我们认为，该方向仍需进一步深入研究。

神经线性方法（Neural Linear Methods）：此类方法可视为广义线性模型，其基函数由神经网络前 l − 1 层定义；即仅对神经网络最后一层进行推断，其余层保持固定（Snoek 等，2015；Riquelme 等，2018；Ovadia 等，2019；Ober & Rasmussen，2019；Pinsler 等，2019；Kristiadi 等，2020）。它们也可被视作子网络推断的特例——其中子网络被简单地定义为网络的最后一层。

子空间推断：神经网络剪枝这一子领域旨在通过识别实现准确预测所需的最小权重子集来提升计算效率（例如 Frankle & Carbin，2019；Wang 等，2020）。我们的工作与其不同：我们保留全部网络权重，但目标是找到一个用于概率推理的小型权重子集。与我们更密切相关的是 Izmailov 等（2019）的工作，他们提出在低维权重子空间（例如由 SGD 轨迹主成分构造的子空间）上进行推断。此外，若干近期方法在变分推断框架下采用低秩参数化来近似后验（Rossi 等，2019；Swiatkowski 等，2020；Dusenberry 等，2020），这亦可视为在权重空间的某种隐式子空间上进行推断。相比之下，我们提出了一种显式识别与预测不确定性相关权重子集的技术——即寻找坐标轴对齐的子空间（axis-aligned subspaces）。

1. 结论

本研究得出三项主要结论：
1）在神经网络中建模权重相关性对获得可靠的预测后验至关重要；
2）在考虑此类相关性的前提下，单峰后验近似即可与多峰近似（如深度集成）相媲美；
3）为获得可靠的预测后验，无需对全部权重进行推断。

基于上述洞见，我们构建了一种将贝叶斯推断扩展至大规模神经网络的框架：仅对权重子集进行后验近似，其余权重保持为确定性点估计。该框架将计算成本与总参数量解耦，从而可灵活权衡计算开销与近似质量，并得以采用更具表达力的后验近似（如全协方差高斯分布）。

线性化拉普拉斯子网络推断方法可事后应用于任意预训练模型，极具实用价值。实证分析表明，该方法：
1）相较于全网络采用粗糙近似的方法，表达能力更强，能保留更多不确定性；
2）允许我们使用容量更大、函数拟合能力更广的神经网络，而不牺牲不确定性估计质量；
3）性能与当前最先进的不确定性量化方法（如深度集成）相当。

我们期待未来进一步探索：将子网络推断与不同近似推断方法结合、开发更优的子网络选择策略，并深入研究子网络对预测分布特性的影响。

原文链接：https://proceedings.mlr.press/v139/daxberger21a/daxberger21a.pdf