Fast Approximate Solution of Stein Equations forPost-Processing of MCMC

用于MCMC后處理的Stein方程快速近似求解方法

https://arxiv.org/pdf/2501.06634

摘要

貝葉斯推斷在理論上具有優美的結構，但計算后驗期望往往需要高昂的計算成本。蒙特卡洛方法具有良好的漸近理論保證，且可靠性高，但卻未能利用被積函數的光滑性。求解Stein方程已成為一種可能的替代方法，它為后驗期望的數值近似提供了一個框架，在該框架中可以有效利用被積函數的光滑性。然而，現有的Stein方程數值解法通常伴隨著較高的計算代價，因為需要求解大規模線性系統。本文探討了結合迭代線性求解器與預處理策略來快速近似求解Stein方程的方法。我們的主要發現是這些方法可以有效，但不同預處理策略的效果依賴于具體應用場景，因此沒有哪一種策略能在所有情況下都表現最優。

1 引言

條件概率是統計推斷的基礎，但條件概率的實際計算仍是一個具有挑戰性的計算任務。我們關注的是貝葉斯統計中出現的概率分布：在已知參數 X 與數據 Y 的聯合分布的情況下，我們希望計算在觀測數據 Y=y 給定下參數 X的條件分布。我們忽略可測性方面的考慮，假設該條件分布在 Rd 上具有Lebesgue密度 p(?)。根據貝葉斯定理，該條件密度為：

在第一種情形中，類比于已有的求積方法（cubature methods），人們可能期望通過利用被積函數的光滑性來獲得更快的收斂速率 [38]。然而在此處的挑戰在于，由于后驗分布是通過公式（1）隱式定義的，因此傳統的求積方法無法直接應用于貝葉斯框架中。

在第二種情形中，基于采樣的方法相對較慢的誤差衰減速率意味著為了獲得高精度的積分近似值，可能會付出較高的計算代價。然而這種收斂速率是所有蒙特卡洛方法所共有的特征，很難看出如何能在采樣方法的范疇內實現有意義的加速，除非能改進速率常數（rate constant）。

求解 Stein 方程 已成為解決上述兩個問題的一種有前景的解決方案。簡而言之，Stein 方程將計算后驗期望的任務轉化為求解一個偏微分方程（PDE）的任務。盡管求解 PDE 看起來是一個更具挑戰性的計算任務，但一個關鍵的觀察是：PDE 求解器可以直接利用函數的光滑性。這使得數值近似可以達到更好的收斂速率，從而降低在逼近感興趣的后驗期望量時所需的計算成本。

本文的目標有三個。首先，我們旨在對用于求解 Stein 方程并計算后驗期望的配置方法進行自成一體的闡述。其次，我們旨在綜述現有的、可用于將配置類方法擴展到大規模數據集的預處理技術，并對其在我們所關注的應用背景下的適用性進行批判性討論。第三，我們旨在對使用這些策略所能帶來的性能提升進行客觀的實證評估。

2 背景知識

我們現在回顧如何利用 Stein 方程對 MCMC 輸出進行后處理。第 2.1 節介紹 Stein 方程，第 2.2 節討論如何通過求解 Stein 方程的近似解來近似計算后驗期望。數值方法，包括配置法（collocation methods），在第 2.3 節中進行討論。我們主要考慮的數值技術是迭代線性求解器和預處理方法，并在第 2.4 節中簡要回顧這些方法在配置法中的應用方式。

2.1 Stein 方程

Stein 方程在后驗期望與偏微分方程（PDE）解之間建立了數學聯系。它最初由 Charles Stein 在 [50] 中以一種特殊形式提出，起初僅作為一種理論工具。如今已有大量不同的構建 Stein 方程的方法，既用于理論研究也用于計算目的 [2]，但在本文中我們僅關注標準形式的 Stein 方程。

2.2 利用 Stein 方程的近似解

2.3 求解標準 Stein 方程

2.4 共軛梯度法與配置方法的預處理

3 Stein 方程的快速近似解法

Stein 方程（3）的數值解法將在第 3.1 節以變分形式進行描述。在考慮第 3.2 節中討論的實際因素的前提下，這將導出一個線性方程組，之后可以進一步對其進行預處理。多種不同的預處理技術將在第 3.3 節中進行介紹。

3.1 變分公式化

本節解釋了如何按照 [42, 41, 3] 的方法，將第 2.3 節中對稱配置法的基本形式擴展到 Stein 方程（3）的情形。

3.2 實用考慮

本節總結了與 Stein 方程數值求解相關的幾個主要實際考慮因素：核函數的選擇、性能監控，以及內存高效的數值預處理方法。

3.2.1 核函數選擇

與所有基于核的方法一樣，為當前任務選擇一個合適的核函數是非常重要的。在我們的場景中，這意味著需要選擇一個適當的基核 k，即公式（11）中的核函數。

例如，我們可以考慮如下形式的基核：

其中 ?>0 是一個需要指定的長度尺度參數。標準技術（如交叉驗證）為如何選擇 ? 提供了有用的依據。然而，這些技術需要在配置節點的子集上求解 Stein 方程，這又再次回到了我們最初所面對的計算任務。特別是，較大的長度尺度 ? 往往會導致矩陣 Kp 更加病態（ill-conditioned）。

因此，我們在第 3.3 節中將要討論的預處理技術所帶來的任何計算優勢，也可以被用于核函數選擇這一任務之上。

3.2 實用考慮

3.2.3 對標準 Stein 方程的內存高效預處理

受第 2.4 節中關于配置法預處理策略討論的啟發，我們在本節中考慮對 Stein 方程進行類似的預處理。具體來說，我們首先求解：

3.3 預處理矩陣（Preconditioners）

在配置法和核方法的相關文獻中，存在多種不同的預處理策略。我們將討論一些具有代表性的示例，但不求窮盡。

3.3.1 Jacobi 預處理

3.3.3 Nystr?m 預處理與對角采樣

基于 Nystr?m 的方法的性能可能對所選擇的 n 個誘導點（inducing points）非常敏感，已有多種節點選擇方法被提出。這些方法可分為固定采樣 （fixed sampling）和自適應采樣 （adaptive sampling）兩類。

• 固定采樣

  方法從一個固定的概率分布中進行無放回采樣以選取 n 個節點，其中最常見的是均勻分布，對應的就是經典的 Nystr?m 方法。

• 自適應采樣

  方法則會動態更新采樣策略，以確保對定義域有較好的覆蓋，通常是通過降低那些已經選中的節點附近的其他節點的采樣權重來實現的。

這些采樣方法在效率與精度之間存在權衡：更復雜的采樣方法雖然耗時更多，但通常可以為 Nystr?m 方法選出質量更高的節點。

在本文中，我們重點關注一種稱為對角采樣 （diagonal sampling）的固定采樣方法。

3.3.4 完全獨立訓練條件（Fully Independent Training Conditional, FITC）

完全獨立訓練條件 （FITC）方法由 [44] 提出，作為一種用于高斯過程計算的近似技術，其中將矩陣 Kp 視為一個協方差矩陣。為了避免繁瑣的記號，我們再次用簡寫符號 K 表示 Kp，省略下標 p。

FITC 方法可以看作是 Nystr?m 近似（18）再加上一個額外的修正項：

3.3.5 隨機 Nystr?m 預處理

隨機 Nystr?m 預處理 方法 [23] 借鑒了 Nystr?m 方法的相同思想，但與從核矩陣中采樣一小部分行和列不同，它使用隨機投影 （random projections）來實現低秩近似。

令矩陣 的元素獨立地服從標準高斯分布。那么核矩陣將被近似為：

3.3.6 隨機截斷奇異值分解（Randomised Truncated SVD）

由 [26] 提出的隨機奇異值分解 （SVD）預處理方法包含兩個階段：
首先，將原始核矩陣投影到一個低維子空間中；然后對投影后的矩陣進行標準的奇異值分解（SVD），并將結果在第二階段映射回原始維度。

3.3.7 譜預處理（Spectral Preconditioning）

我們最后考慮的一種預處理方法是 [34] 提出的譜預處理 （spectral preconditioner）。

為了引出這種預處理的思想，需要注意的是，我們的線性系統（14）實際上是以下最小二乘任務的一個離散版本：

該問題也可以等價地表述為在再生核希爾伯特空間 H(kp) 上定義的一個二次函數的最小化問題，該函數由一個正定（經驗協方差）算子所刻畫。

像共軛梯度法（CG）這樣的迭代方法的收斂速度依賴于該算子的條件數。為了加速梯度下降的收斂，[33] 提出了一種預處理算子，該算子將算子的前 r?1 個最大特征值替換為第 r 個最大特征值。

我們實際實現的預處理子是通過對該預處理算子進行子采樣 （subsampling）得到的一個近似 [34, 1]，并用于在 RKHS（再生核希爾伯特空間）中直接執行共軛梯度法。

其中 再次表示所選取的配置節點的索引。

需要注意的是，矩陣 M 并不是對稱的，因為它源自 RKHS 中的共軛梯度（CG）算法；而底層的預處理算子在 RKHS 的內積意義下是對稱的。

4 實證評估

本節對第 3.3 節中介紹的各類預處理方法進行實證比較。我們使用的測試平臺是一個在第 4.1 節中描述的邏輯回歸問題。實驗結果展示于第 4.2 節，而這些方法的可擴展性則在第 4.3 節中討論。用于復現實驗結果的 Python 代碼可以從以下網址下載：
https://github.com/MatthewAlexanderFisher/FastStein

4.1 實驗設置

邏輯回歸測試平臺 ：邏輯回歸是解析不可解后驗分布的一個簡單示例，需要使用數值方法進行求解。完整的細節見在線附錄，但我們可以指出后驗分布的維度最初被固定為 d=4。

近似樣本是通過隨機游走 Metropolis–Hastings MCMC 方法從后驗分布中生成的，我們將前 N 個不重復的樣本作為我們的配置節點。

4.2 預處理方法的比較

我們研究的結果總結在圖1中，其中展示了不同預處理策略的“增益”如何隨著長度尺度參數 ? 的變化而變化，并且也展示了其如何依賴于各預處理方法的相關參數。（較大的 ? 值對應更病態的矩陣。）

首先，我們注意到并非所有預處理方法表現相同，也沒有哪一種策略在所有考慮的問題設置下都優于其他策略。

其次，我們觀察到，在 ? 取較小值時，所有的預處理方法都導致了性能惡化。

另一方面，除了 Jacobi 預處理之外，其余所有預處理方法在某些特定的 ? 和 η 參數區間內均帶來了性能提升，盡管這些區間在初始階段并不容易直觀預測。

在有利的設置下，我們觀察到了 3 到 4 倍的增益，這意味著所需迭代次數減少了 20 到 50 倍，這是非常顯著的改進。

隨機 SVD 預處理方法在不同的 ? 和 η 區間中表現出最好與最差 的性能；相比之下，Nystr?m 和譜預處理方法則提供了較為穩定的、幅度為 0 到 2 倍的增益（即所需迭代次數減少 1 到 7 倍），前提是 ? 不至于過小。

在在線附錄中，我們進一步展示了當前報告的結果不會因定義數值任務的后驗分布的維度 d 而發生顯著變化。

4.3 向大規模 N 的擴展

雖然在 N 較大時我們無法獲得“真實值”作為基準，但我們仍好奇 Jacobi 預處理方法在此類設置中是否可能帶來性能增益。這是因為在偏微分方程（PDE）背景下，該方法此前已被推薦使用 [20]。

為了進行研究，我們固定了一個邏輯回歸任務，并運行了 次 MCMC 迭代——這在實際應用中是一個典型的 MCMC 輸出樣本量，最終共獲得 N=23,282 個不同的配置節點。

我們分別對使用 Jacobi 預處理（塊大小 b=1）和標準 CG 方法的情況，計算了最壞誤差 σ(wm) 隨迭代次數 m的變化情況。

結果顯示在圖 2 中，可以看出，在初始迭代階段確實實現了對近似誤差的非平凡降低，但達到收斂所需的總迭代次數基本沒有變化。

5 討論

盡管近年來對 Stein 方程的數值求解產生了濃厚興趣，并且已有不少令人鼓舞的概念驗證研究，但這一任務仍然具有相當大的挑戰性。本文揭示了其中的困難，并探討了數值預處理技術在改善計算成本與精度之間權衡關系方面的潛力。

我們的主要發現是：預處理可以是一種有效的策略 ，但不同預處理方法的表現高度依賴于具體的應用場景。

同時，我們也承認本次實證研究存在一定的局限性——我們僅關注了一個特定的邏輯回歸示例、一個特定的核函數，以及由一種特定 MCMC 方法生成的樣本。要判斷我們的結論是否具有普遍適用性，還需要更多后續實證證據的積累。

作為未來研究的一個可能方向，我們注意到：限制函數空間 （restricting the function space）可以提供另一條與預處理方法互補的降低計算成本的路徑，例如 [45, 35] 中的研究所示。

另一個潛在方向是：我們的結果表明，將不同的預處理方法結合使用 ，利用它們在不同參數區間中的互補性能，可能會帶來更優的整體表現。

原文鏈接： https://arxiv.org/pdf/2501.06634