Bayesian Deep Learning via Subnetwork Inference

貝葉斯深度學(xué)習(xí)中的子網(wǎng)絡(luò)推斷

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

摘要
貝葉斯范式有望解決深度神經(jīng)網(wǎng)絡(luò)的核心問題，如校準(zhǔn)性能差和數(shù)據(jù)利用效率低。然而，將貝葉斯推斷擴(kuò)展至大規(guī)模參數(shù)空間通常需引入強(qiáng)約束性近似。本文指出：僅對模型權(quán)重的一小部分子集進(jìn)行推斷，即可獲得準(zhǔn)確的預(yù)測后驗分布；其余權(quán)重則保持為點估計。該子網(wǎng)絡(luò)推斷（subnetwork inference）框架使我們得以在子集上使用表達(dá)能力更強(qiáng)、原本難以處理的后驗近似方法。具體而言，我們實現(xiàn)了一種簡潔、可擴(kuò)展的貝葉斯深度學(xué)習(xí)方法——子網(wǎng)絡(luò)線性化拉普拉斯近似（subnetwork linearized Laplace）：首先獲得全網(wǎng)絡(luò)權(quán)重的最大后驗（MAP）估計，隨后基于線性化拉普拉斯近似，在選定子網(wǎng)絡(luò)上推斷一個全協(xié)方差高斯后驗分布。我們提出一種子網(wǎng)絡(luò)選擇策略，旨在最大程度保留模型的預(yù)測不確定性。實驗表明，該方法在性能上優(yōu)于集成方法（ensembles）及對全網(wǎng)絡(luò)采用表達(dá)能力較弱后驗近似的其他方法。

1. 引言
   深度神經(jīng)網(wǎng)絡(luò)（NNs）的一個關(guān)鍵缺陷是：其預(yù)測往往校準(zhǔn)不良且過度自信——尤其當(dāng)訓(xùn)練與測試數(shù)據(jù)分布存在偏移時（Nguyen et al., 2015; Guo et al., 2017）。為支持可靠決策，神經(jīng)網(wǎng)絡(luò)需穩(wěn)健地量化其預(yù)測不確定性（Bhatt et al., 2020），這對醫(yī)療、自動駕駛等安全攸關(guān)應(yīng)用尤為重要（Amodei et al., 2016）。

貝葉斯建模（Bishop, 2006; Ghahramani, 2015）通過模型參數(shù)的后驗分布，為不確定性量化提供了原則性途徑。不幸的是，神經(jīng)網(wǎng)絡(luò)中精確后驗推斷不可行。盡管貝葉斯深度學(xué)習(xí)領(lǐng)域近年取得進(jìn)展（Osawa et al., 2019; Maddox et al., 2019; Dusenberry et al., 2020），現(xiàn)有方法為適配大規(guī)模網(wǎng)絡(luò)，仍不得不采用不切實際的假設(shè)，嚴(yán)重限制了后驗分布的表達(dá)能力，進(jìn)而損害不確定性估計質(zhì)量（Ovadia et al., 2019; Fort et al., 2019; Foong et al., 2019a）。

或許，這些不切實際的推斷近似可被避免。鑒于神經(jīng)網(wǎng)絡(luò)高度過參數(shù)化，其精度可由一個小型子網(wǎng)絡(luò)良好保持（Cheng et al., 2017）；且在低維權(quán)重子空間中進(jìn)行推斷，即可實現(xiàn)準(zhǔn)確的不確定性量化（Izmailov et al., 2019）。這引出如下問題：一個完整神經(jīng)網(wǎng)絡(luò)的模型不確定性能否被小型子網(wǎng)絡(luò)充分保留？ 本文證明：全網(wǎng)絡(luò)的后驗預(yù)測分布可由子網(wǎng)絡(luò)的后驗預(yù)測分布良好近似。具體貢獻(xiàn)如下：

1. 提出 子網(wǎng)絡(luò)推斷 ——一種可擴(kuò)展的貝葉斯深度學(xué)習(xí)通用框架：僅對神經(jīng)網(wǎng)絡(luò)權(quán)重的一小部分子集進(jìn)行推斷，其余權(quán)重保持為確定性點估計；由此允許使用原本在大規(guī)模網(wǎng)絡(luò)中難以處理的高表達(dá)力后驗近似方法。我們給出該框架的一種具體實現(xiàn)：先擬合全網(wǎng)絡(luò)的最大后驗（MAP）估計，再以線性化拉普拉斯近似在子網(wǎng)絡(luò)上推斷全協(xié)方差高斯后驗（見圖1）。

2. 提出一種基于 全網(wǎng)絡(luò)近似后驗 與 子網(wǎng)絡(luò)近似后驗 之間Wasserstein距離的子網(wǎng)絡(luò)選擇策略。為提升可擴(kuò)展性，子網(wǎng)絡(luò)選擇階段采用對角近似；選定小規(guī)模子網(wǎng)絡(luò)后，即可推斷權(quán)重間的協(xié)方差。實驗發(fā)現(xiàn)： 在子網(wǎng)絡(luò)選擇階段做近似，對后驗預(yù)測的影響遠(yuǎn)小于在推斷階段做近似 。

3. 我們在一系列不確定性校準(zhǔn)與分布偏移魯棒性基準(zhǔn)上評估該方法。實驗表明：高表達(dá)力的子網(wǎng)絡(luò)推斷方法，其性能優(yōu)于對全網(wǎng)絡(luò)進(jìn)行低表達(dá)力推斷的主流貝葉斯深度學(xué)習(xí)方法，也優(yōu)于深度集成（deep ensembles）。

1. 子網(wǎng)絡(luò)后驗近似

設(shè) ∈ ?? 為所有神經(jīng)網(wǎng)絡(luò)權(quán)重的 D 維向量（即所有層權(quán)重矩陣的拼接與展平）。貝葉斯神經(jīng)網(wǎng)絡(luò)（BNNs）旨在捕捉模型不確定性，即由于訓(xùn)練數(shù)據(jù) = {, } 存在多種合理解釋而產(chǎn)生的關(guān)于權(quán)重 選擇的不確定性。其中， ∈ ?? 為輸出變量（例如分類標(biāo)簽）， ∈ ???? 為特征矩陣。首先，需在 BNN 的權(quán)重 上指定一個先驗分布 ()。隨后，我們希望推斷其完整的后驗分布。

這種后驗預(yù)測分布將權(quán)重中的不確定性轉(zhuǎn)化為預(yù)測中的不確定性。遺憾的是，由于神經(jīng)網(wǎng)絡(luò)（NNs）的非線性特性，推斷精確的后驗分布 (|) 是不可行的；又因權(quán)重 的高維度，即使要忠實地近似后驗分布也面臨巨大的計算挑戰(zhàn)。因此，通常采用粗略的后驗近似方法，例如完全因子化近似，即 (|) ≈ ∏?_{d=1} (_d)，其中 _d 是權(quán)重向量 中的第 d 個權(quán)重（Hernández-Lobato & Adams, 2015; Blundell et al., 2015; Khan et al., 2018; Osawa et al., 2019）。然而，已有研究表明，此類近似存在嚴(yán)重缺陷（Foong et al., 2019a,b）。

在本工作中，我們質(zhì)疑廣泛存在的隱含假設(shè)——即一個表達(dá)能力強(qiáng)的后驗近似必須包含全部 D 個模型權(quán)重。相反，我們嘗試僅對權(quán)重的一個小規(guī)模子集 ? 進(jìn)行推斷。以下論證支持這一方法：

1. 過參數(shù)化：Maddox 等人（2020）表明，在局部最優(yōu)解附近，存在許多方向不會改變神經(jīng)網(wǎng)絡(luò)的預(yù)測結(jié)果。此外，神經(jīng)網(wǎng)絡(luò)可被大量剪枝而不犧牲測試集精度（Frankle & Carbin, 2019）。這表明，神經(jīng)網(wǎng)絡(luò)的大部分預(yù)測能力可集中于一個小規(guī)模子網(wǎng)絡(luò)中。
2. 子模型上的推斷：先前研究1 已提供證據(jù)表明，即使推斷未在完整參數(shù)空間上進(jìn)行，仍可有效。例如，Izmailov 等人（2019）和 Snoek 等人（2015）分別在權(quán)重的低維投影空間和神經(jīng)網(wǎng)絡(luò)的最后一層上執(zhí)行推斷。

因此，我們將上述兩個想法結(jié)合起來，對公式 (1) 中的后驗分布做出如下兩步近似：

與權(quán)重剪枝方法的關(guān)系。注意，（4）中的后驗近似可以被視為將權(quán)重的方差修剪為零。這與權(quán)重剪枝方法（Cheng et al., 2017）形成對比，后者將權(quán)重本身設(shè)置為零。即，權(quán)重剪枝方法可以被視為移除權(quán)重以保留預(yù)測均值（即保持與完整模型接近的準(zhǔn)確性）。相比之下，子網(wǎng)絡(luò)推斷可以被視為僅移除某些權(quán)重的方差——同時保持它們的均值——以保留預(yù)測不確定性（例如，保持與完整模型接近的校準(zhǔn)）。因此，它們是互補(bǔ)的方法。重要的是，通過不剪枝權(quán)重，子網(wǎng)絡(luò)推斷保留了完整神經(jīng)網(wǎng)絡(luò)的全部預(yù)測能力以保持其預(yù)測準(zhǔn)確性。

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

在本工作中，我們通過使用線性化拉普拉斯近似（MacKay, 1992）對權(quán)重上的后驗分布進(jìn)行近似，從而滿足公式（4）。這是一種可處理的推斷技術(shù)，近期已被證明表現(xiàn)優(yōu)異（Foong 等，2019b；Immer 等，2020），并可事后應(yīng)用于預(yù)訓(xùn)練模型。下面我們將在一般設(shè)定下對其進(jìn)行描述。

我們將神經(jīng)網(wǎng)絡(luò)函數(shù)記為 : ?? → ??。首先，我們定義一個關(guān)于神經(jīng)網(wǎng)絡(luò)權(quán)重的先驗分布，我們選擇其為完全因子化的高斯分布 () = (; , )。接著，我們尋找后驗分布的一個局部最優(yōu)解，也稱為權(quán)重的最大后驗（MAP）估計：

隨后，利用在 MAP 估計點處的二階泰勒展開對后驗分布進(jìn)行近似：

有趣的是，當(dāng)采用高斯似然時，以廣義高斯-牛頓（GGN）精度矩陣定義的高斯分布，恰好對應(yīng)于將神經(jīng)網(wǎng)絡(luò)在處進(jìn)行一階泰勒展開線性化后的 真實后驗分布 （Khan et al., 2019；Immer et al., 2020）。該局部線性化函數(shù)為：

這些閉式表達(dá)式頗具吸引力，因其所得預(yù)測均值與分類決策邊界 與 MAP 估計所得神經(jīng)網(wǎng)絡(luò)完全一致 。
然而，存儲現(xiàn)代神經(jīng)網(wǎng)絡(luò)（即參數(shù)維度 D 極大）權(quán)重空間上完整的 D × D
協(xié)方差矩陣在計算上是不可行的。盡管已有研究致力于開發(fā)更廉價的近似方案（例如僅存儲對角元（Denker & LeCun, 1990）或塊對角元（Ritter et al., 2018; Immer et al., 2020）），但這些近似均以降低預(yù)測性能為代價。

1. 線性化拉普拉斯子網(wǎng)絡(luò)推斷

我們概述以下程序，用于在子網(wǎng)絡(luò)推斷框架內(nèi)將線性化拉普拉斯近似擴(kuò)展至大規(guī)模神經(jīng)網(wǎng)絡(luò)模型。

1. 子網(wǎng)絡(luò)選擇

理想情況下，我們希望所選子網(wǎng)絡(luò)誘導(dǎo)出的預(yù)測后驗分布盡可能接近對全網(wǎng)絡(luò)進(jìn)行推斷所得的預(yù)測后驗分布（式11）。這種隨機(jī)過程之間的差異通常通過函數(shù)空間的 KL 散度（functional Kullback–Leibler divergence）來量化（Sun 等，2019；Burt 等，2020）：

在權(quán)重空間中，我們的目標(biāo)是最小化全網(wǎng)絡(luò)精確后驗分布（式1）與子網(wǎng)絡(luò)近似后驗分布（式4）之間的差異。這帶來了兩個挑戰(zhàn)：首先，計算精確后驗分布仍是不可行的；其次，常見的差異度量（如 KL 散度或 Hellinger 距離）對于式（4）中出現(xiàn)的狄拉克δ分布并未良好定義。

為解決第一個問題，我們再次借助第3節(jié)中引入的局部線性化方法。線性化模型的真實后驗分布是高斯分布或近似高斯分布2：

我們通過選用平方 2-Wasserstein 距離來解決第二個問題，該度量對于支撐集不相交的分布仍有良好定義。對于全協(xié)方差高斯分布（式21）與一個全協(xié)方差高斯分布和若干狄拉克δ函數(shù)的乘積（式16）的情形，該度量具有如下形式：

表面上看，我們似乎又回到了最初試圖避免的性能較差的對角假設(shè)（Ovadia 等，2019；Foong 等，2019a；Ashukha 等，2020）。然而，這里存在一個關(guān)鍵區(qū)別：我們是在 子網(wǎng)絡(luò)選擇階段 做出對角假設(shè)，而非在 推斷階段 ；我們在子網(wǎng)絡(luò) 上執(zhí)行的是 全協(xié)方差推斷 。在第6節(jié)中，我們將提供證據(jù)表明，在子網(wǎng)絡(luò)選擇階段采用對角假設(shè)是合理的，原因如下：1）相較于在推斷階段做相同假設(shè)，它對預(yù)測性能的損害要小得多；2）它優(yōu)于隨機(jī)子網(wǎng)絡(luò)選擇。

1. 實驗

我們通過實驗評估子網(wǎng)絡(luò)推斷的有效性，并將其與以下方法進(jìn)行比較：（1）對全網(wǎng)絡(luò)采用表達(dá)能力較弱的推斷方法；（2）深度學(xué)習(xí)中當(dāng)前最先進(jìn)的不確定性量化方法。我們考慮三類基準(zhǔn)設(shè)置：
1）小規(guī)模玩具回歸任務(wù)；
2）中等規(guī)模表格數(shù)據(jù)回歸任務(wù)；
3）基于 ResNet-18 的圖像分類任務(wù)。
更多實驗結(jié)果與設(shè)置細(xì)節(jié)分別見附錄 A 與附錄 D。

6.1 子網(wǎng)絡(luò)推斷如何保留后驗預(yù)測不確定性？

我們首先定性評估：在選定子網(wǎng)絡(luò)上采用全協(xié)方差高斯后驗所得預(yù)測分布，與以下方法所得預(yù)測分布的對比情況：
1）全網(wǎng)絡(luò)上的全協(xié)方差高斯后驗（Full Cov）；
2）全網(wǎng)絡(luò)上的因子分解高斯后驗（Diag）；
3）僅在網(wǎng)絡(luò)最后一層上采用全協(xié)方差高斯后驗（Final layer）（Snoek 等，2015）；
4）點估計（MAP）。

對于子網(wǎng)絡(luò)推斷，我們同時考慮兩種子網(wǎng)絡(luò)選擇策略：第5節(jié)所述的Wasserstein策略（Wass）與均勻隨機(jī)選擇策略（Rand），以構(gòu)建僅包含模型參數(shù)總量50%、3%和1%的子網(wǎng)絡(luò)。在此玩具實驗中，精確計算后驗邊際方差以指導(dǎo)子網(wǎng)絡(luò)選擇尚屬可行。

我們的神經(jīng)網(wǎng)絡(luò)包含2個ReLU隱藏層，每層50個隱藏單元。采用同方差高斯似然函數(shù)，其噪聲方差通過最大似然估計優(yōu)化。我們在網(wǎng)絡(luò)權(quán)重（不含偏置）上采用GGN拉普拉斯推斷，并結(jié)合式（18）中的線性化預(yù)測分布。因此，所考察的所有方法共享相同的預(yù)測均值，便于更公平地比較其不確定性估計。

我們將全網(wǎng)絡(luò)先驗精度設(shè)為? = 3?（經(jīng)驗上表現(xiàn)良好），子網(wǎng)絡(luò)先驗精度設(shè)為?= ? S/D。

我們采用 Antorán 等人（2020）提出的合成一維回歸任務(wù)——輸入數(shù)據(jù)形成兩個分離的簇，從而可檢驗?zāi)Ｐ蛯Α按亻g區(qū)域”的不確定性響應(yīng)（Foong 等，2019b）。結(jié)果如圖2所示：

• 子網(wǎng)絡(luò)推斷在 推斷更少權(quán)重 的同時，比對角高斯或僅最后一層推斷更能保留全網(wǎng)絡(luò)推斷的不確定性；

• 通過捕捉權(quán)重間的相關(guān)性，子網(wǎng)絡(luò)推斷可在數(shù)據(jù)簇之間維持較高不確定性；

• 該特性在隨機(jī)與Wasserstein子網(wǎng)絡(luò)選擇下均成立，但后者在子網(wǎng)絡(luò)更小時能保留更多不確定性；

• 相較于對角拉普拉斯，其顯著優(yōu)勢表明： 在子網(wǎng)絡(luò)選擇階段采用對角假設(shè)，但在推斷階段轉(zhuǎn)而使用全協(xié)方差高斯后驗（即本文做法），顯著優(yōu)于直接對推斷后驗采用對角假設(shè) （參見第5節(jié)）。

綜上，結(jié)果表明：在精心選擇的子網(wǎng)絡(luò)上進(jìn)行高表達(dá)力推斷，相較對全網(wǎng)絡(luò)采用粗糙近似，能更好地保留預(yù)測不確定性。

6.2 大型模型中的子網(wǎng)絡(luò)推斷 vs 小型模型中的全網(wǎng)絡(luò)推斷

我們首先獲得每個神經(jīng)網(wǎng)絡(luò)權(quán)重的最大后驗（MAP）估計及其同方差似然函數(shù)的噪聲方差。隨后，對每個網(wǎng)絡(luò)執(zhí)行全網(wǎng)絡(luò) GGN-Laplace 推斷。我們還使用所提出的 Wasserstein 規(guī)則修剪每個網(wǎng)絡(luò)的權(quán)重方差，使剩余方差數(shù)量匹配每一個較小網(wǎng)絡(luò)的規(guī)模。我們采用對角拉普拉斯近似來廉價地估算用于子網(wǎng)絡(luò)選擇的后驗邊際方差。我們利用式（12）和（18）中的線性化方法計算預(yù)測分布。因此，具有相同權(quán)重數(shù)量的神經(jīng)網(wǎng)絡(luò)會產(chǎn)生相同的預(yù)測均值；增加所考慮的權(quán)重方差數(shù)量只會提升預(yù)測不確定性。

我們選用三個規(guī)模遞增的表格數(shù)據(jù)集（輸入維度、樣本點數(shù)）：wine（11維，1439點）、kin8nm（8維，7373點）和 protein（9維，41157點）。我們采用其標(biāo)準(zhǔn)訓(xùn)練-測試劃分（Hernández-Lobato & Adams, 2015），以及專為測試分布外不確定性的變體劃分（Foong 等, 2019b）。具體細(xì)節(jié)見附錄 D.4。對于每個劃分，我們將訓(xùn)練數(shù)據(jù)的15%留作驗證集，用于在尋找MAP估計及選擇權(quán)重先驗精度時進(jìn)行早停。所有模型和數(shù)據(jù)集保持其他超參數(shù)固定。結(jié)果如圖3所示。

我們呈現(xiàn)平均測試對數(shù)似然（LL）值，因其同時考慮了準(zhǔn)確率與不確定性。當(dāng)結(jié)合全網(wǎng)絡(luò)推斷時，規(guī)模更大的模型通常表現(xiàn)最佳，盡管 Wine-gap 和 Protein-gap 是例外。有趣的是，即使我們僅在與小型模型同等規(guī)模的子網(wǎng)絡(luò)上進(jìn)行推斷，這些大型模型的表現(xiàn)依然最優(yōu)。我們推測，這源于權(quán)重后驗神經(jīng)網(wǎng)絡(luò)模型中存在大量退化方向（即冗余權(quán)重）（Maddox 等，2020）。小型模型的全網(wǎng)絡(luò)推斷會同時捕獲有用與無用權(quán)重的信息；而在大型模型中，我們的子網(wǎng)絡(luò)選擇策略使我們能將更多計算資源用于建模信息豐富的權(quán)重方差與協(xié)方差。在6個數(shù)據(jù)集中有3個，我們發(fā)現(xiàn)：隨著推斷所涉及權(quán)重數(shù)量的增加，LL 值出現(xiàn)驟升，隨后進(jìn)入平臺期。這種平臺現(xiàn)象可能是因為大部分信息豐富的權(quán)重方差已被納入模型。考慮到計算 GGN 的成本遠(yuǎn)高于神經(jīng)網(wǎng)絡(luò)訓(xùn)練成本，這些結(jié)果表明： 在相同計算量下，對大型模型執(zhí)行子網(wǎng)絡(luò)推斷比對小型模型執(zhí)行全網(wǎng)絡(luò)推斷更優(yōu) 。

6.3 分布偏移下的圖像分類

我們現(xiàn)在評估采用子網(wǎng)絡(luò)推斷的大型卷積神經(jīng)網(wǎng)絡(luò)在圖像分類任務(wù)中對分布偏移的魯棒性，并與以下基線方法進(jìn)行比較：

• 點估計網(wǎng)絡(luò)（MAP）；

• 對全網(wǎng)絡(luò)采用表達(dá)能力較弱推斷的貝葉斯深度學(xué)習(xí)方法：MC Dropout（Gal & Ghahramani, 2016）、對角拉普拉斯、VOGN（Osawa 等，2019）——三者均假設(shè)權(quán)重后驗完全因子化；以及 SWAG（Maddox 等，2019）——假設(shè)后驗為“對角+低秩”結(jié)構(gòu)；

• 深度集成（deep ensembles）（Lakshminarayanan 等，2017）——目前被公認(rèn)為深度學(xué)習(xí)不確定性量化的最先進(jìn)方法（Ovadia 等，2019；Ashukha 等，2020）。

我們采用5個網(wǎng)絡(luò)構(gòu)成的集成（據(jù) Ovadia 等建議），并對 MC Dropout、對角拉普拉斯與 SWAG 均采樣16次。Dropout 概率設(shè)為 0.1；對角拉普拉斯的先驗精度通過網(wǎng)格搜索確定為 = 4 × 10?。所有方法均應(yīng)用于 ResNet-18（He 等，2016）：包含1個輸入卷積塊、8個殘差塊和1個線性層，共計 11,168,000 個參數(shù)。

對于子網(wǎng)絡(luò)推斷，我們采用式（19）中的線性化預(yù)測分布；并使用 Wasserstein 子網(wǎng)絡(luò)選擇策略，僅保留 0.38% 的權(quán)重，得到一個僅含 42,438 個權(quán)重的子網(wǎng)絡(luò)——這是當(dāng)前計算條件下可處理全協(xié)方差矩陣的最大規(guī)模（其大小為）。我們采用對角 SWAG（Maddox 等，2019）估算子網(wǎng)絡(luò)選擇所需的邊際權(quán)重方差。我們曾嘗試對角拉普拉斯，但發(fā)現(xiàn)其所選權(quán)重對應(yīng)于在訓(xùn)練點上神經(jīng)網(wǎng)絡(luò)雅可比恒為零的位置（即“死亡 ReLU”）；此類權(quán)重的后驗方差雖大（近似先驗），但對網(wǎng)絡(luò)輸出幾乎無影響。SWAG 不受此問題困擾，因其忽略了訓(xùn)練梯度為零的權(quán)重。子網(wǎng)絡(luò)推斷的先驗精度經(jīng)網(wǎng)格搜索設(shè)為 = 500。

為評估原則性子網(wǎng)絡(luò)選擇的重要性，我們另設(shè)一基線：均勻隨機(jī)選擇子網(wǎng)絡(luò)（記為 Ours (Rand)）。我們開展以下兩個實驗，結(jié)果見圖4：

旋轉(zhuǎn) MNIST：參照（Ovadia 等，2019；Antorán 等，2020），所有方法在 MNIST 上訓(xùn)練，并在逐步增大的數(shù)字旋轉(zhuǎn)角度下評估其預(yù)測分布。盡管所有方法在原始 MNIST 測試集上表現(xiàn)良好，但當(dāng)旋轉(zhuǎn)角度超過 30 度時，準(zhǔn)確率迅速下降。就對數(shù)似然（LL）而言，集成在基線方法中表現(xiàn)最優(yōu)；而子網(wǎng)絡(luò)推斷的 LL 顯著高于幾乎所有基線（包括集成），唯一例外是 VOGN（表現(xiàn)略優(yōu)）。值得注意的是，Ovadia 等（2019）亦觀察到：平均場變分推斷（VOGN 屬于此類）在 MNIST 上表現(xiàn)極強(qiáng)，但在更大規(guī)模數(shù)據(jù)集上性能顯著下降。子網(wǎng)絡(luò)推斷在分布內(nèi)能做出準(zhǔn)確預(yù)測，同時對分布外樣本賦予比基線更高的不確定性。

損壞 CIFAR：同樣參照（Ovadia 等，2019；Antorán 等，2020），所有方法在 CIFAR10 上訓(xùn)練，并在經(jīng) 16 類不同損壞（每類5個強(qiáng)度等級）的數(shù)據(jù)上評估（Hendrycks & Dietterich, 2019）。由于局部線性化使預(yù)測均值與 MAP 一致，子網(wǎng)絡(luò)推斷在預(yù)測誤差上與 MAP 網(wǎng)絡(luò)相當(dāng)；集成與 SWAG 準(zhǔn)確率最高。然而，子網(wǎng)絡(luò)推斷的獨特優(yōu)勢在于過自信程度最低——在所有損壞強(qiáng)度等級下，其對數(shù)似然均優(yōu)于所有基線方法。此時 VOGN 表現(xiàn)較差；但這與其在 MNIST 上的優(yōu)異表現(xiàn)看似矛盾——實則再次印證了 Ovadia 等（2019）的發(fā)現(xiàn)：平均場變分推斷在 MNIST 上表現(xiàn)良好，但在更大數(shù)據(jù)集上性能下降。

此外，在兩項基準(zhǔn)測試中，隨機(jī)選擇子網(wǎng)絡(luò)的表現(xiàn)顯著劣于我們提出的 Wasserstein 選擇策略，凸顯了子網(wǎng)絡(luò)選擇方式的重要性。

綜上，這些結(jié)果表明：子網(wǎng)絡(luò)推斷在不確定性校準(zhǔn)與分布偏移魯棒性方面，優(yōu)于其他主流不確定性量化方法。

1. 適用范圍與局限性

   多輸出模型中的雅可比矩陣計算仍具挑戰(zhàn)性。在當(dāng)前主流深度學(xué)習(xí)框架中，由于采用反向模式自動微分，其計算所需反向傳播次數(shù)等于模型輸出數(shù)量。這使得線性化拉普拉斯方法難以應(yīng)用于語義分割（Liu 等，2019）或類別數(shù)極多的分類任務(wù)（Deng 等，2009）。需注意，該問題僅限于線性化拉普拉斯方法本身；其他無此限制的推斷方法仍可納入本框架使用。

   先驗精度 的選擇在很大程度上決定了拉普拉斯近似的性能。我們提出的子網(wǎng)絡(luò)先驗精度更新方案依賴于對全網(wǎng)絡(luò)已有合理參數(shù)設(shè)定。然而，由于全網(wǎng)絡(luò)推斷常不可行，目前選擇 的最佳方式是直接在子網(wǎng)絡(luò)近似上進(jìn)行交叉驗證。

   海森矩陣的存儲需求限制了子網(wǎng)絡(luò)權(quán)重的最大規(guī)模。例如，存儲 4 萬個權(quán)重對應(yīng)的海森矩陣約需 6.4 GB 內(nèi)存。對于現(xiàn)代 Transformer 等超大規(guī)模模型，可計算的子網(wǎng)絡(luò)僅占總權(quán)重極小比例。盡管我們已證明優(yōu)異性能未必依賴大型子網(wǎng)絡(luò)（見圖5），但探索更優(yōu)的子網(wǎng)絡(luò)選擇策略仍是未來研究的關(guān)鍵方向。

   
   

1. 相關(guān)工作

貝葉斯深度學(xué)習(xí)：針對神經(jīng)網(wǎng)絡(luò)權(quán)重后驗分布 p ( w ∣ D ) 的刻畫已有大量研究。迄今為止，哈密頓蒙特卡洛（Hamiltonian Monte Carlo, HMC；Neal, 1995）仍是貝葉斯神經(jīng)網(wǎng)絡(luò)（BNNs）中近似推斷的黃金標(biāo)準(zhǔn)。盡管其在漸近意義上無偏，但基于采樣的方法難以擴(kuò)展至大規(guī)模數(shù)據(jù)集（Betancourt, 2015）。因此，近年來更流行的做法是在某一近似分布族（通常為高斯分布）中尋找最優(yōu)代理后驗。其中最早的是 MacKay（1992）提出的拉普拉斯近似，他也同時建議使用線性化模型的后驗來近似預(yù)測后驗（Khan 等，2019；Immer 等，2020）。隨著更大規(guī)模神經(jīng)網(wǎng)絡(luò)的普及，能捕捉權(quán)重間相關(guān)性的代理分布因計算不可行而受限；因此，絕大多數(shù)現(xiàn)代方法轉(zhuǎn)而采用平均場假設(shè)（Blundell 等，2015；Hernández-Lobato & Adams，2015；Gal & Ghahramani，2016；Mishkin 等，2018；Osawa 等，2019），但這犧牲了模型表達(dá)能力（Foong 等，2019a）并導(dǎo)致實證性能下降（Ovadia 等，2019；Antorán 等，2020）。Farquhar 等（2020）曾提出：在更深網(wǎng)絡(luò)中，平均場假設(shè)或許并不構(gòu)成限制；但我們的實證結(jié)果似乎與該觀點相悖。我們發(fā)現(xiàn)，通過降低權(quán)重空間維度來擴(kuò)展那些能考慮權(quán)重相關(guān)性的近似方法（如 MacKay，1992；Louizos & Welling，2016；Maddox 等，2019；Ritter 等，2018），其性能優(yōu)于對角近似。由此我們認(rèn)為，該方向仍需進(jìn)一步深入研究。

神經(jīng)線性方法（Neural Linear Methods）：此類方法可視為廣義線性模型，其基函數(shù)由神經(jīng)網(wǎng)絡(luò)前 l ? 1 層定義；即僅對神經(jīng)網(wǎng)絡(luò)最后一層進(jìn)行推斷，其余層保持固定（Snoek 等，2015；Riquelme 等，2018；Ovadia 等，2019；Ober & Rasmussen，2019；Pinsler 等，2019；Kristiadi 等，2020）。它們也可被視作子網(wǎng)絡(luò)推斷的特例——其中子網(wǎng)絡(luò)被簡單地定義為網(wǎng)絡(luò)的最后一層。

子空間推斷：神經(jīng)網(wǎng)絡(luò)剪枝這一子領(lǐng)域旨在通過識別實現(xiàn)準(zhǔn)確預(yù)測所需的最小權(quán)重子集來提升計算效率（例如 Frankle & Carbin，2019；Wang 等，2020）。我們的工作與其不同：我們保留全部網(wǎng)絡(luò)權(quán)重，但目標(biāo)是找到一個用于概率推理的小型權(quán)重子集。與我們更密切相關(guān)的是 Izmailov 等（2019）的工作，他們提出在低維權(quán)重子空間（例如由 SGD 軌跡主成分構(gòu)造的子空間）上進(jìn)行推斷。此外，若干近期方法在變分推斷框架下采用低秩參數(shù)化來近似后驗（Rossi 等，2019；Swiatkowski 等，2020；Dusenberry 等，2020），這亦可視為在權(quán)重空間的某種隱式子空間上進(jìn)行推斷。相比之下，我們提出了一種顯式識別與預(yù)測不確定性相關(guān)權(quán)重子集的技術(shù)——即尋找坐標(biāo)軸對齊的子空間（axis-aligned subspaces）。

1. 結(jié)論

本研究得出三項主要結(jié)論：
1）在神經(jīng)網(wǎng)絡(luò)中建模權(quán)重相關(guān)性對獲得可靠的預(yù)測后驗至關(guān)重要；
2）在考慮此類相關(guān)性的前提下，單峰后驗近似即可與多峰近似（如深度集成）相媲美；
3）為獲得可靠的預(yù)測后驗，無需對全部權(quán)重進(jìn)行推斷。

基于上述洞見，我們構(gòu)建了一種將貝葉斯推斷擴(kuò)展至大規(guī)模神經(jīng)網(wǎng)絡(luò)的框架：僅對權(quán)重子集進(jìn)行后驗近似，其余權(quán)重保持為確定性點估計。該框架將計算成本與總參數(shù)量解耦，從而可靈活權(quán)衡計算開銷與近似質(zhì)量，并得以采用更具表達(dá)力的后驗近似（如全協(xié)方差高斯分布）。

線性化拉普拉斯子網(wǎng)絡(luò)推斷方法可事后應(yīng)用于任意預(yù)訓(xùn)練模型，極具實用價值。實證分析表明，該方法：
1）相較于全網(wǎng)絡(luò)采用粗糙近似的方法，表達(dá)能力更強(qiáng)，能保留更多不確定性；
2）允許我們使用容量更大、函數(shù)擬合能力更廣的神經(jīng)網(wǎng)絡(luò)，而不犧牲不確定性估計質(zhì)量；
3）性能與當(dāng)前最先進(jìn)的不確定性量化方法（如深度集成）相當(dāng)。

我們期待未來進(jìn)一步探索：將子網(wǎng)絡(luò)推斷與不同近似推斷方法結(jié)合、開發(fā)更優(yōu)的子網(wǎng)絡(luò)選擇策略，并深入研究子網(wǎng)絡(luò)對預(yù)測分布特性的影響。

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