神經網絡高斯過程

神經網絡高斯過程（英語：neural network Gaussian process，簡稱NNGP) 是一種特殊的高斯過程，可以看作一類特定人工神經網絡序列的極限。具體而言，當多種神經網絡架構的寬度趨於無窮時，其函數分布會收斂到一個高斯過程。^{[1][2][3][4][5][6][7][8]}

背景

貝葉斯網絡是一種建模工具，它通過為事件分配概，來量化模型預測中的不確定性。深度學習和人工神經網絡則是機器學習中的主流方法，用於構建能從訓練樣本中學習的計算模型。貝葉斯神經網絡則將二者相融合，是一種參數與預測都具有概率性的神經網絡。^[9][10]標準的神經網絡常會對錯誤的預測賦予較高的置信度^[11]，而貝葉斯神經網絡能夠更準確地評估自身預測正確的可能性。

左圖：包含兩個隱藏層的貝葉斯神經網絡，該網絡將三維輸入（下）轉換為二維輸出${\displaystyle (y_{1},y_{2})}$（上）。右圖：網絡輸出的概率密度函數${\displaystyle p(y_{1},y_{2})}$，由網絡的隨機權重所決定。視頻：隨著網絡寬度的增加，輸出分布變得簡單，最終在無限寬度極限下收斂到多元常態分布。

人工神經網絡的計算可以表示成人工神經元構成的層序列，其中每一層的神經元數量稱為層的寬度。當我們考察一個貝葉斯神經網絡序列（見圖），其所有層的寬度不斷增加時，這個序列在函數分布上最終會收斂到一個神經網絡高斯過程。這種無限寬度下的極限具有現實意義，因為在實踐中更寬的網絡通常會有更好的性能。^[12][4][13]同時，這一極限也為我們提供了一種評估網絡性能的閉式方法。

除了作為貝葉斯神經網絡的極限，神經網絡高斯過程還出現在其他幾種場景下：它描述了一個非貝葉斯寬人工神經網絡在參數隨機初始化之後、訓練之前的輸出函數分布；它可以作為神經正切核預測方程中的關鍵組成部分；它在深度信息傳播中用以表徵超參數和架構是否可以訓練。^[14]它還與神經網絡的其他大寬度極限有關。

適用範圍

關於神經網絡與高斯過程的第一個對應關係，最早可以追溯到Radford M. Neal在其1995年博士論文中的結果^[15]，當時他的導師是多倫多大學教授傑弗里·辛頓。Neal表示他的研究靈感來源於從事貝葉斯學習研究的戴維·J·C·麥凱。

如今，這一對應關係已被證明適用於多種架構，如單隱藏層貝葉斯神經網絡^[15]、層寬趨於無窮時的深度全連接網絡^[2][3]、通道數趨於無窮時的卷積神經網絡^[4][5][6]、注意頭數量趨於無窮時的Transformer網絡^[16]、單元數趨於無窮時的循環網絡^[8]等。事實上，這種對應關係對幾乎所有神經網絡架構都成立。只要一個架構可以完全由矩陣乘法和逐坐標的非線性運算來表達，那麼它就存在一個無限寬度下的高斯過程極限。^[8]這一定義涵蓋了由多層感知器、循環神經網絡（如LSTM和GRU）、（任意維度或圖上的）卷積、池化、跳躍連接、注意力、批量歸一化及層歸一化組成的所有前饋或循環神經網絡。

圖解說明

當無限寬度網絡的參數${\displaystyle \theta }$從其先驗分布${\displaystyle p(\theta )}$中採樣時，網絡輸出的最終分布可以由高斯過程描述。

神經網絡參數的每一組參數${\displaystyle \theta }$對應於由該網絡所計算的特定函數。當我們指定網絡參數的先驗分布${\displaystyle p(\theta )}$時，也就等同於在網絡可能實現的所有函數上確定了一個先驗分布。對於許多網絡架構而言，當其寬度趨於無窮時，這個函數空間上的分布會收斂到一個高斯過程。

圖中直觀地展示了這一概念。對於一維輸出的神經網絡${\displaystyle z^{L}(\cdot ;\theta )}$，該圖的坐標軸表示網絡對兩個不同輸入${\displaystyle x}$和${\displaystyle x^{*}}$的對應輸出值。圖中的每一個黑點都代表了一次隨機採樣：先從${\displaystyle p(\theta )}$中隨機抽取一組參數，然後再計算兩個輸入值對應的輸出對${\displaystyle z^{L}(x;\theta )}$和${\displaystyle z^{L}(x^{*};\theta )}$。而圖中的紅線則描繪了由${\displaystyle p(\theta )}$得到的輸出對的聯合概率分布。這是參數空間中${\displaystyle p(\theta )}$對應於函數空間中的分布。在無限寬的神經網絡中，由於其函數分布是一個高斯過程，因此對任意有限的一組輸入，其輸出的聯合分布必定是一個多元高斯分布。

討論

無限寬全連接網絡

本節針對全連接網絡這一具體架構，討論無限寬神經網絡與高斯過程之間的對應關係。我們將提供了一個證明概要，旨在闡述這一對應關係成立的原因，並給出這一架構下神經網絡高斯過程的具體函數形式。此處的證明概要主要遵循了Novak等人提出的方法。^[4]

網絡架構

全連接貝葉斯神經網絡

考慮一個全連接人工神經網絡，其中${\displaystyle x}$為輸入，參數${\displaystyle \theta }$由每一層${\displaystyle l}$的權重${\displaystyle W^{l}}$和偏置${\displaystyle b^{l}}$組成，${\displaystyle z^{l}}$和${\displaystyle y^{l}}$則分別表示每一層的預激活值（非線性運算之前）和激活值（非線性運算之後），${\displaystyle \phi (\cdot )}$為逐點作用的非線性函數，${\displaystyle n^{l}}$為層寬。為簡單起見，輸出向量${\displaystyle z^{L}}$的寬度${\displaystyle n^{L+1}}$取為 1。假設網絡的參數具有先驗分布${\displaystyle p(\theta )}$，其中每一個權重和偏置都獨立地服從一個各向同性的高斯分布，而權重的方差與層寬成反比。該網絡的結構如右圖所示，並由以下方程組描述：

${\displaystyle {\begin{aligned}x&\equiv {\text{input}}\\y^{l}(x)&=\left\{{\begin{array}{lcl}x&&l=0\\\phi \left(z^{l-1}(x)\right)&&l>0\end{array}}\right.\\z_{i}^{l}(x)&=\sum _{j}W_{ij}^{l}y_{j}^{l}(x)+b_{i}^{l}\\W_{ij}^{l}&\sim {\mathcal {N}}\left(0,{\frac {\sigma _{w}^{2}}{n^{l}}}\right)\\b_{i}^{l}&\sim {\mathcal {N}}\left(0,\sigma _{b}^{2}\right)\\\phi (\cdot )&\equiv {\text{nonlinearity}}\\y^{l}(x),z^{l-1}(x)&\in \mathbb {R} ^{n^{l}\times 1}\\n^{L+1}&=1\\\theta &=\left\{W^{0},b^{0},\dots ,W^{L},b^{L}\right\}\end{aligned}}}$

高斯過程${\displaystyle z^{l}|y^{l}}$

首先，我們注意到預激活值${\displaystyle z^{l}}$可以被描述為一個以激活值${\displaystyle y^{l}}$為條件的高斯過程。這一結論即便在有限寬度的網絡中也成立。每個預激活值${\displaystyle z_{i}^{l}}$是一系列高斯隨機變量（權重${\displaystyle W_{ij}^{l}}$和偏置${\displaystyle b_{i}^{l}}$）的加權和，而這一求和中每個高斯變量的係數都是之前的激活值${\displaystyle y_{j}^{l}}$。因為${\displaystyle z_{i}^{l}}$是零均值高斯隨機變量的加權和，所以其本身也是零均值高斯隨機變量。對於任意${\displaystyle y^{l}}$，${\displaystyle z^{l}}$服從聯合高斯分布，因此其可以被定義為以${\displaystyle y^{l}}$為條件的高斯過程。該高斯過程的協方差（即核函數）取決於權重方差${\displaystyle \sigma _{w}^{2}}$、偏置方差${\displaystyle \sigma _{b}^{2}}$以及激活值${\displaystyle y^{l}}$的二階矩矩陣${\displaystyle K^{l}}$：

${\displaystyle {\begin{aligned}z_{i}^{l}\mid y^{l}&\sim {\mathcal {GP}}\left(0,\sigma _{w}^{2}K^{l}+\sigma _{b}^{2}\right)\\K^{l}(x,x')&={\frac {1}{n^{l}}}\sum _{i}y_{i}^{l}(x)y_{i}^{l}(x')\end{aligned}}}$

其中權重方差${\displaystyle \sigma _{w}^{2}}$的作用是縮放來自${\displaystyle K^{l}}$的貢獻，而偏置由於是對所有輸入共享的，因此${\displaystyle \sigma _{b}^{2}}$使得${\displaystyle z_{i}^{l}}$更接近常數矩陣，即讓不同輸入數據點所對應的${\displaystyle z_{i}^{l}}$更相似。

高斯過程${\displaystyle z^{l}|K^{l}}$

預激活值${\displaystyle z^{l}}$對${\displaystyle y^{l}}$的依賴僅取決於其二階矩矩陣${\displaystyle K^{l}}$。正因如此，我們可以說${\displaystyle z^{l}}$是一個以${\displaystyle K^{l}}$為條件的高斯過程，而無需取決於整個激活向量${\displaystyle y^{l}}$：

${\displaystyle {\begin{aligned}z_{i}^{l}\mid K^{l}&\sim {\mathcal {GP}}\left(0,\sigma _{w}^{2}K^{l}+\sigma _{b}^{2}\right).\end{aligned}}}$

層寬趨於無窮時${\displaystyle K^{l}\mid K^{l-1}}$的確定性

如前所述，${\displaystyle K^{l}}$是激活值${\displaystyle y^{l}}$的二階矩矩陣。由於${\displaystyle y^{l}}$是對預激活值${\displaystyle z^{l-1}}$應用非線性函數${\displaystyle \phi }$的結果，可以將其替換為${\displaystyle \phi \left(z^{l-1}\right)}$，從而將${\displaystyle K^{l}}$（${\displaystyle l>0}$）的定義改寫為

${\displaystyle {\begin{aligned}K^{l}(x,x')&={\frac {1}{n^{l}}}\sum _{i}\phi \left(z_{i}^{l-1}(x)\right)\phi \left(z_{i}^{l-1}(x')\right).\end{aligned}}}$

前文已證明${\displaystyle z^{l-1}|K^{l-1}}$是一個高斯過程。這意味著，上述${\displaystyle K^{l}}$定義中的求和項，實現上是${\displaystyle n^{l}}$個樣本的平均，這些樣本都是從基於${\displaystyle K^{l-1}}$為的高斯過程中採樣得到的，即

${\displaystyle {\begin{aligned}\left\{z_{i}^{l-1}(x),z_{i}^{l-1}(x')\right\}&\sim {\mathcal {GP}}\left(0,\sigma _{w}^{2}K^{l-1}+\sigma _{b}^{2}\right).\end{aligned}}}$

隨著層寬${\displaystyle n^{l}}$趨於無窮大，這${\displaystyle n^{l}}$個高斯過程樣本的均值會收斂到一個高斯過程上的積分：

${\displaystyle {\begin{aligned}\lim _{n^{l}\rightarrow \infty }K^{l}(x,x')&=\int dz\,dz'\,\phi (z)\,\phi (z')\,{\mathcal {N}}\left(\left[{\begin{array}{c}z\\z'\end{array}}\right];0,\sigma _{w}^{2}\left[{\begin{array}{cc}K^{l-1}(x,x)&K^{l-1}(x,x')\\K^{l-1}(x',x)&K^{l-1}(x',x')\end{array}}\right]+\sigma _{b}^{2}\right)\end{aligned}}}$

因此，在無限寬度極限下，輸入對${\displaystyle x}$和${\displaystyle x'}$的二階矩矩陣${\displaystyle K^{l}}$可以通過一個關於二維高斯分布的積分來計算。對許多常見的激活函數${\displaystyle \phi (\cdot )}$，例如ReLU^[17]、ELU、GELU^[18]或誤差函數^[1]等，這一積分都存在解析解。即使在沒有解析解的情況下，由於它只是一個二維積分，通常也可以通過數值方法高效地計算。^[2]由於該積分是確定性的，${\displaystyle K^{l}|K^{l-1}}$也是確定性的。

為了方便表示，我們定義一個泛函${\displaystyle F}$，它代表了上述積分的計算過程，並將前一層的${\displaystyle K^{l-1}}$映射到當前層的${\displaystyle K^{l}}$：

${\displaystyle {\begin{aligned}\lim _{n^{l}\rightarrow \infty }K^{l}&=F\left(K^{l-1}\right).\end{aligned}}}$

神經網絡高斯過程${\displaystyle z^{L}\mid x}$

上節我們得到，${\displaystyle n^{l}\rightarrow \infty }$時${\displaystyle K^{l}\mid K^{l-1}}$是確定性的。遞歸地應用這一結論，最後一層的${\displaystyle K^{L}}$可以寫成輸入層${\displaystyle K^{0}}$的確定性函數：

${\displaystyle {\begin{aligned}\lim _{\min \left(n^{1},\dots ,n^{L}\right)\rightarrow \infty }K^{L}&=F\circ F\cdots \left(K^{0}\right)=F^{L}\left(K^{0}\right),\end{aligned}}}$

其中，${\displaystyle F^{L}}$表示將${\displaystyle F}$連續應用${\displaystyle L}$次。同時，輸入層的二階矩矩陣${\displaystyle K^{0}(x,x')={\tfrac {1}{n^{0}}}\sum _{i}x_{i}x'_{i}}$也是輸入${\displaystyle x}$的確定性函數，加之我們已證明${\displaystyle z^{L}|K^{L}}$是一個高斯過程，最終我們可以將神經網絡的輸出表示為輸入的高斯過程：

${\displaystyle {\begin{aligned}z_{i}^{L}(x)&\sim {\mathcal {GP}}\left(0,\sigma _{w}^{2}F^{L}\left(K^{0}\right)+\sigma _{b}^{2}\right).\end{aligned}}}$

軟體庫

Neural Tangents是一個由Google開發的免費開源Python庫，可以用於計算和推斷與各常見神經網絡架構相對應的神經網絡高斯過程和神經正切核。^[19]