Disentangling impact of capacity, objective, batchsize, estimators, andstep-size on flow VI

分離容量、目標函數、批量大小、估計器和步長對流變分推斷（flow VI）的影響

https://arxiv.org/pdf/2412.08824?

摘要

基于歸一化流的變分推斷（flow VI）是一種有前景的近似推斷方法，但其在不同研究中的表現仍不一致。許多算法選擇會影響 flow VI 的性能。我們進行了逐步分析，以厘清一些關鍵因素的影響：容量（capacity）、目標函數（objectives）、梯度估計器（gradient estimators）、梯度估計的數量（batchsize）以及步長（step-sizes）。每一步都單獨考察一個因素，同時利用前一步驟的見解或通過大量并行計算來中和其他因素的影響。為了實現高保真評估，我們整理了一個合成目標分布的基準測試集，這些目標分布代表了常見的后驗病理情況，并允許進行精確采樣。我們針對各個因素提供了具體的建議，并提出了一種 flow VI 的實施方案（recipe），其性能達到或超過了當前主流的即插即用哈密頓蒙特卡洛（HMC）方法的表現。

1 引言

基于歸一化流的變分推斷（flow VI）[80] 通過一系列可逆變換構建靈活的變分族，以近似后驗分布 [96, 3, 5, 38, 31] [73, 55]。盡管這一方法具有前景，但其表現仍不一致——一些研究報告了成功案例 [3, 5, 38, 91, 31]，而另一些則強調優化困難和較差的結果 [43, 10, 8, 24, 47, 60, 6]。

多個因素影響著 flow VI 的性能，使得難以確定不一致性的原因：是容量不足、目標函數不合適、梯度方差過高，還是步長設置錯誤？這些因素之間的“糾纏”限制了我們對 flow VI 能力的理解和最佳實踐的形成，阻礙了其更廣泛的應用。

本文采用逐步分析的方法，分別探討幾個關鍵因素的影響——容量（第3節）、目標函數（第4節）、梯度估計器與批量大小（第5節），以及步長（第6節）。每一步都回答一個關于某個因素的基本問題，同時通過應用前一步驟的見解或利用現代 GPU 集群的計算能力來中和其他因素的影響。例如，在第4節中我們提出問題：是否需要復雜的覆蓋多模態的目標函數？為了測試這一點，我們通過借用第3節中驗證過的流模型結構來中和容量的影響，并通過使用極大的批量大小和廣泛的超參數掃描來中和優化選擇的影響。

這種方法要求有精確的準確性衡量標準。然而，由于真實后驗通常是難以處理的，這使得高保真評估變得復雜：在沒有真實樣本的情況下，適用于變分推斷（VI）和哈密頓蒙特卡洛（HMC）的通用評估指標非常有限——證據下界及其變體 [16, 25, 29, 30] 僅適用于 VI，而像有效樣本數量這樣的收斂性度量則主要用于 HMC。雖然可以使用已有方法生成的樣本作為真實樣本的代理 [61]，但通用且可靠的評估指標仍然有限——Wasserstein 距離 [93] 雖常被使用，但其計算成本隨樣本數量增長較快 [21]，不適合用于高保真評估（見第2.2節）。

為了解決評估難題，我們整理了一組合成目標分布基準（第2.1節），這些目標反映了常見的后驗病理特征：條件數差、非線性曲率、重尾分布以及參數間的相互依賴關系。這些目標提供了真實的樣本，從而確保了高保真評估（評估策略詳見第2.3節）。此外，我們使用了一個可擴展的 Wasserstein 距離替代指標——邊緣-Wasserstein 距離（公式2），它通過對各個一維邊緣分布的 Wasserstein 距離進行平均得到。該距離度量在樣本數量上呈對數線性增長，能夠高效而準確地比較 VI 和 HMC 方法（第2.2節）。

我們在下面展示了主要的研究發現。

容量（Capacity） 。Real-NVP 流 [27] 可以準確地逼近具有挑戰性的目標分布（圖3）。盡管 Real-NVP 被廣泛使用 [96, 3, 38, 99]，但一些研究也報告了其表現不佳的情況 [10, 24, 47, 6]。我們通過中和其他因素的影響后表明，只要網絡結構設計得當，Real-NVP 具備足夠的容量來準確表示復雜的后驗分布（第3節）。

目標函數（Objectives） 。在具備高容量流模型的前提下，復雜的目標函數是不必要的（圖4）。一些工作建議使用覆蓋多模態的目標函數，因為它們能更好地覆蓋整個分布 [59, 95, 68]，但也有一些研究表明這些目標難以優化 [35, 37]。我們直接使用真實樣本優化覆蓋多模態的 KL 散度（KL(p ∥ q)），并驗證這種方法確實可以取得很好的效果（圖4）。然而，當模型容量足夠時，更易于優化的尋找模態的 KL 散度（KL(q ∥ p)）已經足夠有效（圖4）。

梯度批量大小與估計器（Gradient batchsize and estimators） 。大批量大小顯著提升了高容量流模型的表現（圖5）。已有研究探索了用于流模型的低方差梯度估計器 [3, 91, 92]。雖然這些方法有所幫助，但我們發現僅靠它們還不夠，更大的批量大小能夠顯著提升高容量流模型的性能，因此建議在可能的情況下將兩者結合使用（第5節）。

步長與優化（Step-size and optimization） 。保持步長在一個狹窄范圍內對收斂至關重要（圖6）。自動選擇合適步長仍然是一個開放性問題 [56, 97]。我們展示出即使在使用高容量流、低梯度方差和自適應優化器 [51] 的情況下，優化過程也可能在看似穩定運行數千次迭代之后突然發散（第6節）。實驗表明，將步長控制在 之間，可以在不同目標上實現穩定的長期訓練結果（圖6）。

總體來看，有兩個關鍵因素至關重要：高容量的流模型 （解決表達能力限制）和大的梯度批量大小 （簡化優化過程）。有了這兩個要素，一個簡單的實施方案就已足夠：使用標準的變分推斷目標函數、盡可能使用降低方差的梯度估計器，并采用固定在小范圍內的步長；合理初始化并參數化流模型架構，并使用像 Adam 這樣的自適應優化器進行長時間優化 [51]（詳見附錄B）。

基于這一方案，我們展示了 flow VI 在具有挑戰性的目標分布上，僅需較少的模型評估次數即可達到或超越主流 HMC 方法的表現（見第7節以及圖1和圖7）。

2 設定（Setup）

給定一個模型 p(z,y)，其中 z 是潛在變量（latent variables），y 是觀測變量（observed variables），推斷的目標是近似后驗分布 p(z∣y)。變分推斷（Variational Inference）通過最大化證據下界（ELBO）來學習一個近似分布 q [86, 49, 94, 12]，其中：

2.1 目標分布（Targets）

我們設計了一組合成目標分布的基準測試，用以捕捉現實世界后驗分布中常見的病理特征（參見附錄 L 獲取目標細節）。所有這些目標分布都允許精確采樣 ，使得原本無法進行的分析成為可能（見第3節和第4節）。精確樣本還使得對 VI 和 HMC 方法的公平比較 成為可能（評估策略詳見第2.3節）。圖2展示了這些三維目標分布的樣本。

2.2 評估指標（Metrics）

有一些性能評估指標僅適用于 VI 或 HMC 中的一種，而非兩者通用：例如 ELBO、證據上界 [48] 和重要性加權 ELBO [16] 等標準指標并不適用于 HMC；而像有效樣本數量這樣的收斂診斷指標則不適用于 VI。

還有一些指標使得精確比較變得困難：預測測試似然度由于測試數據集的不確定性 [2] 而不可靠；而對于像自由度為 1.5 的學生 t 分布（Student-t）這類重尾目標分布來說，矩甚至可能不存在。

Wasserstein 距離 [93] 的優勢在于它可用于對 VI 和 HMC 進行公平比較，但它需要參考樣本，并且在樣本數量增加時計算效率較低 [21]。在初步實驗中，我們發現即使使用了數千個樣本，Wasserstein 距離估計中的噪聲仍然使得精細比較難以實現（參見附錄 C 中圖8）。

幸運的是，在單變量（一維）情況下，計算 Wasserstein 距離可以簡化為排序操作。基于這一點，我們采用了邊緣-Wasserstein （marginal-Wasserstein）度量——即各對應的一維邊緣分布之間的 Wasserstein 距離的平均值。

設 A 和 B 為兩個 S×d 的樣本矩陣，其中 S 是樣本數量，d 是維度數量。令 表示對 A 按列排序后的版本，使得 是第 j 列中的最小值。

與 Wasserstein 距離一樣，當樣本量相對較小時，這種度量也會受到蒙特卡洛噪聲的影響，但由于它具有更好的可擴展性（scalability），我們可以使用更大的樣本量來降低噪聲（見附錄 C）。

當然，邊緣-Wasserstein 度量僅關注邊緣分布，因此它無法捕捉維度之間的相關性。然而，由于其良好的可擴展性，它成為了高保真評估中唯一可靠的選擇。

2.3 評估策略

為了進行可靠的性能評估，我們使用邊緣-Wasserstein 度量（公式2）。即使在精確推斷的情況下，該度量的最小可實現值也依賴于樣本數量和目標分布的幾何特性 [21]。為了體現理想性能的參考標準，我們繪制了從目標分布中獨立抽取的兩組樣本之間的邊緣-Wasserstein 度量，并用黑色虛線表示（如圖1所示）。當一種推斷方法的結果接近這條線時，說明其生成的樣本質量與目標分布的樣本相當。

對于所有合成目標的評估，我們都使用同一組一萬個參考樣本（見圖1和圖7）。在第7節中，我們還在一些真實目標上進行了評估，并通過長時間運行 HMC 方法生成參考樣本（詳見附錄K）。

2.4 實驗細節

為了模擬不同的規模，我們使用了三個維度設置：2維、10維和100維（除了 Funana 目標，它至少需要3個維度）。我們將 flow-VI 在 JAX [14] 中實現，并使用 TensorFlow Probability [58] 和 NumPyro [76] 中的實現來進行 HMC 方法的實驗。所有方法均在 Nvidia A100 GPU 上運行。如需完整細節，請參見附錄 G 到 K。

3 Real-NVP 流是否具備足夠的容量來表示具有挑戰性的目標分布？

Real-NVP 流是推斷研究者中一種廣受歡迎的選擇 [96, 3, 89, 24, 85, 38, 87, 99, 6]。然而，也有一些研究報告了其表現不穩定或效果不佳的情況 [24, 47, 6]。目前尚不清楚這些結果是由于 Real-NVP 流在表達能力上的不足，還是優化過程中的失敗所導致。

為了解決這一問題，我們在中和其他因素的前提下 ，研究容量（capacity）的影響，并提出如下問題：Real-NVP 流是否具備足夠的容量來準確表示具有挑戰性的目標分布？

為了回答上述問題，我們直接優化 KL(p ∥ q)。具體來說，我們使用來自目標分布的大批量樣本進行優化，從而大大簡化了訓練過程。（實際中用于優化 KL(p ∥ q) 的算法通常會受到基于自歸一化重要性采樣的高方差梯度的影響 [35, 37]，但在使用精確樣本的情況下，這些問題可以被規避。）

此外，優化 KL(p ∥ q) 通常是更優的選擇，因為由此得到的近似分布相較于通過優化標準變分推斷目標函數 KL(q ∥ p) 得到的近似結果，預期能夠更好地覆蓋整個目標分布 [67, 68, 71, 50]。

我們通過掃描不同的步長和學習率調度策略、并進行大量迭代優化，來中和其他優化因素的影響（詳見附錄 G）。

圖3 展示了當我們通過增加耦合層的數量或神經網絡中的隱藏單元數量來提升模型容量時，邊緣-Wasserstein 度量的變化情況（詳見附錄 E）。關鍵發現是：只要合理選擇這些網絡結構參數，Real-NVP 流具備足夠的容量來準確表示復雜的后驗分布 。

唯一的例外是自由度為 1.5 的學生 t 分布（Student-t），在這種情況下，KL(p ∥ q) 的優化因重尾特性而遇到困難 [47]（相比之下，KL(q ∥ p) 表現良好，見圖4）。

主要發現包括以下幾點：

1. 層數 ：增加層數非常有效。當層數達到10層或更多時，性能往往能夠接近精確推斷的結果（隨著層數增加，曲線逐漸靠近黑色虛線）。

2. 隱藏單元數 ：增加隱藏單元對高維問題更有幫助（紅色曲線與藍色曲線之間的差異隨維度增加而增大）。

盡管人們通常認為增加容量會提升表達能力，但這些結果表明，現有的流模型（如 Real-NVP）已經具備表示復雜后驗幾何結構的能力。因此，只要做出適當的優化選擇，即使是“相對簡單”的流模型，也有望取得令人印象深刻的結果。

4 Flow-VI 是否需要覆蓋多模態的目標函數？

前一節表明，Real-NVP 流具備足夠的容量來表示具有挑戰性的目標分布。然而，第3節依賴于使用精確樣本 進行優化，這使得我們有必要重新考慮所使用的優化目標。

選擇合適的優化目標是一個活躍的研究領域。一些研究者提倡使用“覆蓋多模態”的目標函數（mode-spanning objectives），例如第3節中使用的 KL(p ∥ q)，因為它們能夠更好地覆蓋整個目標分布 [59, 95, 25, 42, 64, 19, 101]。盡管這些方法前景良好，但其實用算法可能面臨梯度方差大的問題 [78, 32, 37, 35]。

一個主要的替代方案是使用“更簡單”的標準變分推斷目標函數。然而，也有研究表明，在與流模型結合使用時，這種方法也可能遇到困難 [10, 24, 6]。

總體而言，目標函數的選擇仍不明確，引發了一些關鍵問題：
我們是否真的需要復雜的覆蓋多模態的目標函數？
當模型容量足夠高時，標準的 VI 目標是否已經足夠有效？

為了回答上述問題，我們在中和其他因素影響的前提下 ，專注于優化“更容易”的 KL(q ∥ p) 目標，并檢驗它是否足以實現我們在第3節中通過優化 KL(p ∥ q) 所觀察到的良好性能。

為此，我們采用了大批量大小 （以減少方差），并對梯度估計器、步長和學習率調度策略 進行了詳盡搜索，并進行了大量迭代（詳見附錄 H）。

圖4 展示了在優化 KL(q ∥ p) 和 KL(p ∥ q) 后，邊緣-Wasserstein 度量隨耦合層數量變化的趨勢。關鍵發現如下：

1. 目標函數 ：優化復雜但“覆蓋多模態”的 KL(p ∥ q)（藍色曲線）確實比優化“尋找模態”的 KL(q ∥ p)（紅色曲線）取得了更好的結果，除了 Student-t 分布（自由度 ν = 1.5）的情況 ，這可能是由于其重尾特性導致的 [47]。

2. 容量 ：隨著流模型容量的增加，兩種目標之間的性能差距逐漸縮小，兩者都能達到接近精確推斷的結果。（僅在 Funana 目標的一百維情況下仍有明顯差距，表明該問題具有非常復雜的幾何結構。）

這些結果揭示了 flow VI 與使用簡單變分族（如高斯分布）的傳統 VI 之間的一個重要區別：
與其使用復雜的散度（divergence）來彌補目標分布與變分族之間的不匹配，不如直接增加流模型的容量并優化標準目標函數 。只要容量足夠，散度形式的選擇就不再那么關鍵 。

5 減小方差的估計器是否有幫助？

前幾節表明，如果流模型具備足夠的容量，優化標準的 KL(q ∥ p) 目標函數就可以逼近具有挑戰性的目標分布。然而，第4節所采用的策略依賴于不切實際的超參數窮舉搜索 。為了實現有效的實際應用，我們需要一種更高效的流程。為此，我們首先提出一個問題：如何估計梯度以實現可靠的優化？

優化 ELBO（公式1）的一個主要挑戰在于它依賴于隨機梯度估計 。已有不少研究探索了構建低方差梯度估計器的方法 [79, 82, 36, 66, 83, 90, 9, 33, 100, 17]。一個流行的選擇是“sticking-the-landing”（STL）[83, 90] 估計器，它通常表現良好 [3, 23, 91, 92, 6]。然而，STL 要求對流變換 T 進行求逆 [3, 91, 92]，這會帶來額外的計算成本和潛在的數值問題 [10]。對于像自回歸流這樣求逆代價高昂的流模型來說，STL 也不太實用 [53, 72]。

從原則上講，現代 GPU 可以為梯度方差問題提供一種蠻力解決方案——只需并行抽取大量樣本進行估計并取平均值。這引發了一些自然的問題：在使用高容量流模型時，梯度估計器的選擇與批量大小之間是如何相互作用的？

為了回答這個問題，我們使用不同批量大小 ，分別采用 STL 和標準的總梯度估計器來優化一個高容量流模型。我們對不同的迭代次數進行獨立優化，并對步長和學習率調度策略進行詳盡搜索，以減少優化因素的影響（詳見附錄 I）。

圖5 繪制了在十維目標下，邊緣-Wasserstein 度量隨梯度批量大小變化的趨勢。關鍵發現如下：

1. 估計器選擇 ：在小批量情況下，STL（紅色）始終優于標準的總梯度估計器（藍色），但隨著批量大小增加，兩者之間的差異逐漸縮小，并最終都達到了精確樣本的精度水平。

2. 批量大小 ：對于任意固定的迭代次數（同一行數據），性能隨著批量大小的增大而顯著提升，突出了減小梯度方差的重要性。

這些發現強調了即使對于那些能夠表示復雜后驗分布的高容量流模型而言，減小梯度方差對于實證性能仍然至關重要。大批量大小可以顯著提高性能，在硬件條件允許的情況下幾乎都應該被使用。此外，STL 也能有效減小方差、提升收斂性，在可行的情況下也應被優先考慮（參見附錄 F 獲取效率分析）。

6 優化過程能否可靠地實現自動化？

雖然前一節關注的是可靠地估計 ELBO 的梯度 ，但其余的超參數搜索仍然不切實際。這引出了下一個重要問題：我們如何可靠地進行優化？

自動化變分推斷（VI）的優化仍然是一個開放性問題 [56, 4, 5, 84, 97]。諸如優化器的選擇、步長（step-size）、學習率調度策略（step-schedule）以及迭代次數的設定等決策并不直觀 [56, 3]，而現有文獻中也缺乏針對 flow VI 的具體實用指南。

已有若干使用流模型的研究報告稱，Adam 優化器 [51] 取得了成功。即便如此，在使用 Adam 的前提下，選擇最優的步長調度策略仍然是一個挑戰。

這引發了一系列問題：

• 是否存在一個 適用于不同目標分布 的高容量流模型的通用 步長范圍 ？

• 是否應該根據 維度大小對步長進行縮放 ？

• 能否在 幾百次迭代內的早期表現 基礎上預測出最佳步長？

為了透明地回答這些問題 ，我們使用十個不同的步長值 ，對高容量流模型進行了獨立優化 ，并嘗試了不同數量的迭代次數。我們采用了大批量大小 和 STL 估計器 （詳見附錄 J）。

圖6 展示了邊緣-Wasserstein 度量隨步長變化的情況。從中可以得出一些關鍵結論：

1. 導致發散的步長 ：某些步長在運行數千次迭代時表現良好，但在更長時間運行后會出現發散。例如，在一百維 Funnel 目標中，一個步長為的設置在 1K 次迭代時表現良好，但在 10K 次迭代時卻出現了發散。

2. 穩定的步長范圍 ：步長在 范圍內時，在多個目標分布上均表現出穩定良好的性能（一個顯著例外是一百維的 Funnel 分布，此時稍大的步長反而效果更好）。

這些結果表明，僅憑早期幾百次迭代的表現來預測最佳步長是具有挑戰性的 。當使用梯度方差較低的高容量流模型時，在至少優化 10K 次迭代的前提下，將步長控制在 這一狹窄范圍內，可以在不同目標分布和維度下實現最優性能（參見圖中的紫色、灰色和黃色曲線）。

7 Flow VI 與 HMC 方法的比較如何？

至此，前幾節的內容為我們提供了一個簡單的實施方案（recipe）：

• 使用高容量流模型，

• 采用傳統的 VI 目標函數（以簡化優化），

• 在可行的情況下使用 STL 估計器（以降低梯度方差），

• 使用大批量大小（進一步降低方差），

• 以及在小范圍內使用固定步長（詳見附錄 B 中的詳細建議）。

自然而然地，我們提出一個問題：這一方案的實際效果究竟如何？

在本節中，我們將所提出的方案與當前最先進的即插即用式（turnkey）HMC 方法進行比較。

在現代 GPU 上使用 flow VI 是相對直接的——只需在每次迭代中并行計算更多隨機梯度估計值即可。然而，要將 GPU 有效用于 HMC 方法則更具挑戰性 [63, 45, 46]。對于像 NUTS 這樣領先的即插即用式 HMC 方法來說，其復雜之處在于鏈依賴（chain-dependent）的控制流程使得并行化困難 [57, 75, 77, 45]。而對于一些最新方法，如 CHEES [45] 或 MEADS [46]，高效的跨鏈通信需求也使實現變得復雜。目前，運行多個并行采樣鏈 是將 GPU 與 HMC 方法結合使用的最簡單且最有效的方式 [62, 45]。

在實際問題中，計算瓶頸通常是模型評估 。因此，我們從兩個維度來統計模型評估次數：

• 并行評估次數 （代表加速器規模的代理指標），
• 順序評估次數 （代表運行時間的代理指標）。

對于 HMC 方法來說，并行評估次數對應于鏈的數量，順序評估次數對應于 leapfrog 步數。

對于 VI 方法來說，并行評估次數對應于梯度批量大小，順序評估次數對應于優化迭代次數。

由于當前框架的限制，我們對預熱階段（warmup phase）中的 leapfrog 步數進行了外推處理，參考了預熱后階段的數據（見附錄 K）。

圖7 繪制了邊緣-Wasserstein 度量隨順序評估次數變化的趨勢，針對不同的目標密度。除了合成目標外，我們還在六個現實世界問題上進行了比較（模型細節見附錄 L）。關鍵發現如下：

1. 并行評估次數 ：隨著并行計算資源的增加（從上到下穿越各行），所有方法的性能都有顯著提升。

2. 順序評估次數 ：在較小的批量大小下，flow VI 的表現與 HMC 方法相當或更優；而在批量更大時，flow VI 所需的順序評估次數要少得多。

這些結果表明，flow VI 能夠更有效地利用大批量大小，避免了基于鏈式采樣的 HMC 方法中固有的協調開銷。當然，未來 HMC 在現代加速器上的進一步發展可能會削弱這一優勢；但就目前而言，使用我們提出的方案的 flow VI 是一個強有力的替代選擇。

8 結論

本文通過逐步分析的方法，揭示了影響 flow VI 性能的關鍵因素。我們的分析發現，高容量流模型 和大梯度批量大小對于實現優異性能至關重要。

我們為目標函數、梯度估計器以及優化策略的選擇提供了實用建議。此外，我們還展示了 flow VI 在具有挑戰性的目標分布上，能夠達到或超越主流即插即用式 HMC 方法的表現 ，且所需的順序計算步驟要少得多。

A 相關工作（Related Works）

已有若干研究探索將歸一化流應用于黑盒變分推斷 （BBVI）[96, 8, 3, 24, 5, 6]。這些研究通常對流模型的能力做出一些先驗假設，并更側重于具體應用。

我們則專注于通過逐步分析的方法 ，厘清影響 flow VI 的關鍵因素。我們的研究結果應有助于未來此類方法的應用。

許多研究致力于自動化 BBVI 的某些方面 [56, 3, 23, 4, 5, 97]，旨在為概率模型提供無需人工干預的即插即用式解決方案。與這些方法不同，我們的工作重點在于理解不同因素的影響，提供詳細的指導建議，而非追求完全自動化的解決方案。我們認為我們的分析為未來基于流模型的自動推斷工具奠定了必要的基礎。

Agrawal 等人 [3] 提出了一種結合歸一化流、STL [83]、步長搜索方案以及后處理重要性采樣步驟的方法，以提升現成的 BBVI 性能。但他們并未剖析影響 flow VI 的各個因素，使用了相對較小的批量大小，運行在 CPU 上，并且未對其失敗案例進行深入分析。相比之下，我們的工作專門探討了不同因素的影響，充分利用現代 GPU 進行大批量訓練，并展示了適當優化后的 flow VI 可以達到或超越 HMC 方法的表現。

Andrade [6] 研究了如何在高維問題上穩定地優化 Real-NVP 流，并考慮了網絡結構選擇對優化穩定性的影響。我們的研究獨立采用了與其最優選擇一致的結構，從而避免了不穩定性。雖然他們提出的一些技巧可能進一步提升我們的性能，但我們的目標是全面理解各因素的影響，而不僅僅局限于結構設計。

Dhaka 等人 [24] 使用基于重要性采樣的診斷方法評估性能，推薦使用歸一化流，但由于優化困難而報告了較差的表現（見其附錄中的圖 C.2 和 C.3）。根據我們的研究發現，我們推測他們所遇到的問題可能部分源于相對較低容量的流模型（導致表達能力不足）以及較小的批量大小（導致梯度方差較高）。

Jaini 等人 [47] 揭示了具有重尾的目標分布需要近似分布具備相同尾部特性的基礎分布。我們在圖3中的發現支持了這一觀點：優化 KL(p ∥ q) 在重尾目標上的表現不佳。更重要的是，在第4節中，使用標準 VI 目標函數和高容量流顯著提升了性能。

Blessing 等人 [13] 強調了在推斷研究中標準化評估的必要性，并提出了一個包含合成與現實世界問題的基準測試集，特別關注多模態目標。我們也使用了合成密度來構建受控環境，并采用如 Wasserstein 距離等積分指標進行評估。此外，我們還提供了一個理想性能的參考指標（圖中黑色虛線，見第2.3節），以便更好地解讀數值結果。我們還引入了一組六個現實世界問題，以展示 flow VI 的實際性能。

最近的文獻也探索了將流模型作為 MCMC 方法提議分布的用途 [74, 44, 98, 68, 7, 65, 34, 40, 15, 50, 85, 18]，其中一些方法使用 KL(q ∥ p) 優化來初始化提議分布。這些方法與直接使用 flow VI 是正交的，并可能從本研究所揭示的見解中受益（例如學習更好的初始提議分布）。

B 推薦方案（Suggested Recipe）

要實現優異的 Flow VI 性能，需要進行幾項關鍵的選擇。然而，其中一些選擇取決于可用的計算資源。為了便于更廣泛地采用，我們提出一個推薦方案，并鼓勵實踐者根據自身約束條件進行適當調整。

絕對必需項

通常，在使用性能更強的 GPU 時更容易滿足以下要求：

• 容量（Capacity） ：使用高容量的流模型以減少表示能力上的限制。
  我們默認使用至少包含十個耦合層、每層有32個隱藏單元的 Real-NVP 流模型（除非明確研究容量的影響），詳見圖3和圖4，以及附錄 E 中的結構細節。

• 批量大小（Batchsize） ：使用大量的梯度估計值（即大 batchsize）以降低梯度方差并簡化優化過程。
  在非顯式變化批量大小的實驗中，我們使用了 個樣本。不過，即使使用較小的批量大小也能取得不錯的效果（見圖5）。

在具備高容量流模型和大批量大小的前提下，以下建議表現非常出色：

• 目標函數（Objective） ：優化傳統的 VI 目標函數。盡管覆蓋多模態的目標函數有一定幫助，但標準目標更容易優化，并能達到相當的性能（見圖4）。

• 梯度估計器（Estimator） ：使用“sticking-the-landing”（STL）梯度估計器以降低梯度方差。
  對于某些流模型來說，使用 STL 可能具有挑戰性，請參見附錄 F 的相關討論。

• 步長（Step-size） ：選擇一個固定在 范圍內的步長（見圖6）。如果資源允許，可在該范圍內進行搜索。

• 優化策略（Optimization） ：使用 Adam 等自適應優化器進行大量迭代優化。
  實驗表明，至少進行一萬次更新在不同目標和維度下都表現出色（見圖6）。

基于初步實驗，我們還提出以下附加建議：

• 基礎分布（Base distribution）

  ：在優化前評估不同基礎分布在 ELBO 上的表現，并選擇初始 ELBO 最高的那個。
  為了實驗透明，我們默認使用標準正態分布作為基礎分布。但在多個目標分布上，將拉普拉斯近似設為基礎分布有助于優化的熱啟動 [29, 30, 3]。將學生 t 分布設為基礎分布也被證明在重尾目標上具有優異的表現 [47, 60, 6]。我們建議在優化前進行小規模嘗試，選擇 ELBO 更高的那個作為基礎分布。

• 參數初始化（Parameter initialization） ：將神經網絡參數初始化為較小的值，等價于將流變換初始化為恒等變換 [3]。我們建議采用這種方式，以便更好地控制初始分布。

• 尺度函數的非線性選擇（Non-linearity for scale function） ：在仿射變換中的尺度函數部分，我們建議使用雙曲正切函數（hyperbolic tangent）（參見公式3及附錄 E 中的相關討論）。尺度函數所使用的非線性形式會影響仿射耦合流的穩定性 [10, 6]。我們嘗試了多種非線性函數后發現，將雙曲正切函數嵌套在指數函數中（如公式3所示）能夠在穩定性和表達能力之間取得良好的平衡，并在多個目標分布上表現出一致的良好性能。

C n-Wasserstein 評估細節

兩個分布之間的 Wasserstein 距離通常無法以閉合形式表示 [93]。用于計算它的數值方法在樣本數量增加時擴展性較差 [21]。當使用有限數量的樣本時，這些方法會產生噪聲較大的估計值，使得 n-Wasserstein 距離在精細比較中不可靠。為了說明這一點，我們提供了一個簡單的代表性示例。

C.1 Wasserstein 計算示例

考慮一個十維的標準正態目標分布。我們將計算這個目標分布與某些近似分布之間的 2-Wasserstein 距離和邊緣-Wasserstein 度量。對于高斯分布來說，這兩種 Wasserstein 度量都可以用閉合形式表示 [28, 20]。我們將把這些精確值 與通過數值方法計算出的值進行對比，以理解這些度量在精細比較中的有效性。

由于邊緣-Wasserstein 度量具有更好的可擴展性，我們可以使用更大的樣本量，并在相對較短的時間內獲得極其精確的測量結果（見圖8b）。

我們使用 ott-jax [22] 來計算 2-Wasserstein 距離，并在 Nvidia A100 GPU 上運行實驗。

總體來看，圖8 表明：除非使用大量的樣本，否則具有不同程度準確性的近似分布可能會被評估為相同的錯誤 Wasserstein 值。2-Wasserstein 度量不適用于精細比較 ，因為它無法擴展到此類精確評估所需的大量樣本。相比之下，邊緣-Wasserstein 度量具有更強的擴展性，在合理時間內即可提供準確的測量結果 。

D 局限性（Limitations）

我們在實驗中使用了合成目標分布 。這一選擇是有意為之，目的是減少可能的變量干擾，并便于與 HMC 方法進行高保真度的性能比較。未來的研究可以深入探討這些發現對現實世界問題 的影響。然而，我們預計我們的結論仍然適用。事實上，我們也展示了 flow VI 在多個現實問題上能夠達到或超越 HMC 方法的表現（見圖7和圖12）。

我們使用了像 Nvidia A100 這樣的現代 GPU，目前這類硬件尚未普及。盡管存在計算資源上的限制，我們仍相信研究人員和實踐者可以從我們建立的指導原則中受益，從而更好地利用他們現有的資源。

對于合成目標分布，我們最多只測試到了一百維 。雖然這已經是一個相當高的維度，但我們的逐步分析需要大量的計算資源，在 Nvidia A100 GPU 上總共消耗了超過 4000 小時的計算時間。我們預計其中一些發現也適用于更高維度的情形，但將其擴展到更高維將留待未來研究。

一些經驗豐富的專家可能會認為我們的一些發現是顯而易見的。然而，我們認為，除非有人像我們一樣進行直接而全面的研究，否則這些技術的實際效果仍然是假設性的。

E Real-NVP 架構細節（Details of Real-NVP Architecture）

我們使用了具有仿射耦合層 （affine coupling layers）的 Real-NVP [27] 流模型。我們定義每個耦合層由兩個變換組成，其中一次變換對應于對一部分潛在變量進行仿射變換。

例如，如果第 k 層的輸入變量為 z(k)，則第一個變換定義如下：

K.1 非合成模型的參考樣本（Reference Samples for Non-Synthetic Models）

對于非合成模型，我們需要參考樣本以評估邊緣-Wasserstein 距離。為此，我們并行運行了 NumPyro 中帶預處理的 NUTS （根據初步實驗表現最佳）進行采樣，共運行 214 條鏈，每條鏈運行 10K 次迭代，并進行抽稀處理，最終總共收集一百萬個樣本。

然而，由于我們選擇的模型具有較為復雜的幾何結構（見附錄 L.2），我們發現其中若干鏈陷入了低概率區域（表現為在整個運行過程中樣本重復），未能收集到有效樣本。這種情況尤其出現在以下三個模型中：

• Exoplanet 模型
• Pharmacokinetics 模型
• Sparse Linear Regression 模型

（我們使用的是 32 位浮點數精度。在使用 64 位的情況下，性能可能會有所改善。）

在生成參考樣本時，我們剔除了所有陷入低概率區域的鏈。這樣每個模型至少保留了 497,395 個有效樣本。因此，我們統一使用 497,395 個樣本作為所有非合成模型的參考樣本。

在進行評估時，我們盡可能多地使用方法生成的樣本。如果某個方法生成的樣本超過 497,395 個，則從中取前 497,395 個；若少于該數量，則全部使用。

L 目標分布細節（Target Details） L.1 合成目標分布（Synthetic Targets）

病態高斯分布（Ill-conditioned Gaussian）
這是一個零均值的多維高斯分布，其協方差矩陣的特征值從伽馬分布（shape = 0.5, scale = 1）中采樣，并通過一個隨機正交矩陣進行旋轉，從而引入高度相關性。我們使用 Inference Gym 的實現來確保目標分布可復現 [88]。

香蕉分布（Banana）[39]
該分布在二維空間中對標準正態分布進行了非線性變換：對第二維應用一個非線性變換，使其依賴于第一維的值。該變換包含一個曲率參數，它根據第一維的值調整第二維的尺度。當指定維度大于 2 時，其余維度設為標準正態變量。

對于 d 維情況，該分布定義如下：

L.2 非合成模型（Non-synthetic Models）

我們還在一些目標密度無法用閉合形式表示、也無法獲得真實樣本的模型上比較了 flow VI 與 HMC 方法的表現。這些模型的具體細節如下：

Exoplanet 模型（維度 = 7）
該模型描述了一顆系外行星從其宿主恒星前方經過（即發生凌星現象）時，恒星亮度下降的過程。通過分析恒星的光變曲線（light curve），該模型可以推斷出行星的屬性，例如其大小和軌道周期。該模型包括描述凌星形狀的函數及其相關的行星參數。

控制凌星現象和系外行星特性的潛在變量具有以下先驗分布：

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