Scaling Continuous Latent Variable Models asProbabilistic Integral Circuits

連續隱變量模型的概率積分電路擴展

https://arxiv.org/abs/2406.06494v2

摘要

概率積分電路（Probabilistic Integral Circuits, PICs）最近被提出作為一種概率模型，它具備表達力強的生成模型背后的關鍵要素：連續潛在變量（LVs）。PIC 是符號化的計算圖，將連續潛在變量模型定義為由函數組成的層級結構，這些函數通過相加、相乘，或對某些潛在變量進行積分來組合。如果潛在變量可以解析地積分掉，那么 PIC 是可追蹤的；否則，可以用一種稱為 QPC（Quadrature PC）的概率電路來近似，該電路編碼了一種分層數值積分過程。到目前為止，只有樹狀結構的 PIC 被研究過，而通過數值積分訓練它們在大規模數據上需要大量內存。本文中我們解決了這些問題，并提出了：(i) 一個基于任意變量分解構建有向無環圖（DAG）結構 PIC 的流程；(ii) 一個使用張量化電路架構訓練 PIC 的方法；以及 (iii) 神經函數共享技術以實現可擴展訓練。在廣泛的實驗中，我們展示了函數共享的有效性以及 QPC 相比傳統 PC 的優越性。

1 引言

連續潛在變量（LVs）可以說是許多成功生成模型的關鍵要素，從變分自編碼器 [24] 到生成對抗網絡 [20]，再到最近的擴散模型 [56]。雖然這些模型能夠從數據中學習復雜的分布，但它們通常僅限于采樣，在執行諸如計算邊緣分布或條件分布等基本概率推理任務時并不高效，甚至不可行。另一方面，對于（分層）離散潛在變量模型來說，這些任務可能是可追蹤的 [4,3]，但在大規模數據上學習這些模型更具挑戰性 [9,10,32,33]。

這種在可追蹤性、學習難度和表達能力之間的權衡，可以通過一種新提出的深度生成模型——概率積分電路（PICs）[18] 來分析和探索。PIC 使用符號函數電路定義連續潛在變量的層次結構。當其連續潛在變量可以解析地積分掉時，PIC 是可追蹤的。而對于不可追蹤的 PIC，也可以系統性地用（張量化）概率電路（PCs）[53,4] 進行近似，后者是離散潛在變量模型的表示語言。這類 PC 實例編碼了用于近似 PIC 的分層數值積分過程，因此被稱為“積分 PC”（QPC）。

從 PIC 中提取 QPC 已被證明是一種有效的訓練 PC 的替代方式，但目前僅針對樹狀結構的 PIC 進行了探索。如何構建并擴展到更豐富的潛在變量結構仍然是一個開放的研究問題，需要新的工具 [18]。在本文中，我們填補了這一空白，通過重新定義 PIC 的語義，并將其擴展到 DAG 形式的連續潛在變量層次結構。具體來說，我們將 PIC 設計為一種表示分層擬張量分解的語言 [49]，并通過輕量級多層感知機進行參數化。

貢獻如下：

1. 我們提出了一個系統的流程，可以從任意變量分解出發構建 DAG 結構的 PIC（第 3.1 節）。

2. 我們展示了如何通過一種分層積分過程來學習和近似 PIC，并將其編碼為與已有電路架構 [43,42,31,36] 匹配的張量化 QPC（第 3.2 節）。

3. 我們提出了函數共享技術以實現 PIC 訓練的可擴展性，使我們可以用多頭多層感知機（MLP）來參數化 PIC，其資源需求與 PC 相當（第 3.3 節）。

4. 在廣泛的實驗中（第 4 節），我們展示了：(i) 函數共享在擴展性方面表現出色；(ii) 與通過 EM 或 SGD 常規訓練的 PC 相比，QPC 表現更優，而且它們是從參數數量少達 99% 的 PIC 中提煉出來的。

2 概率積分電路

符號說明。我們將輸入變量表示為X，將潛在變量（LVs）表示為Y和Z，它們的取值分別記作x、y和z。我們用小寫字母表示標量（例如，w ∈ ?），用粗體小寫字母表示向量（例如，），用粗體大寫字母表示矩陣（不包括 X、Y、Z，例如 ），用粗體花體字母表示張量（例如，。

一個概率積分電路（Probabilistic Integral Circuit, PIC）c是一個符號化的計算圖，它表示關于觀測變量X和連續潛在變量Z的非負函數c(X) = ∫ c(X, z) dz，即對于所有x，有c(x) ≥ 0。與概率電路（Probabilistic Circuits, PCs）[53, 4] 類似，PIC 也包含輸入單元、求和單元和乘積單元。然而，與 PC 不同的是，PIC 操作的是函數而非標量，并且使用了一種新型的單元：積分單元 ∫，它可以用來表示連續潛在變量。為了滿足非負性，我們使用非負函數對 PIC 進行參數化，并采用正的求和權重，具體如下：

圖 1(b) 和圖 2(b) 展示了概率積分電路（PIC, Probabilistic Integral Circuits）的示例。請注意，我們使用f來表示附加在輸入單元和積分單元上的函數，這些函數本質上是模型的參數；而我們使用g來表示所有類型單元所輸出的函數。一個 PIC c c的輸出是其根單元 u u返回的輸出函數，該函數僅在輸入空間 X X上定義，即 c ( X ) = g u ( X ) c(X)=gu(X)，并且 Z u = ? Zu=?。

與概率電路（PCs）類似，在 PIC 上施加結構性約束可以實現可處理的推理 [4]。因此，我們假設如下結構性條件：
(i) 所有加法單元（+ 單元）接收的函數都定義在相同的輸入變量上（即“平滑性”）；
(ii) 所有乘法單元（× 單元）接收的函數都定義在不相交的輸入變量集合上（即“可分解性”）。

當一個 PIC 的局部變量（LVs）可以被解析地積分出去時，該 PIC 被認為是可處理的。這意味著我們可以穿過積分單元，計算它們所定義的積分問題，最終輸出一個函數。
值得注意的是，當局部變量之間具有線性-高斯關系[27] 或當這些函數是多項式時，這種積分是可行的。

對于不可處理的 PIC，我們可以通過一種分層數值求積過程來進行近似，該過程可以被編碼成一種稱為求積概率電路（Quadrature PC，QPC）的 PC。直觀上，PIC 中的每一個積分單元都可以用 QPC 中的一組加法單元來近似，這些加法單元以一些預先計算的求積值為條件，并包含大量但有限個輸入單元 [18]。將一個 QPC 具體實現（materialize）出來，可以用于通過近似最大似然方法來訓練 PIC：給定一個 PIC，其附加在輸入、加法和積分單元上的參數的梯度，可以通過對應的 QPC 進行反向傳播 [18]。這種方法也為訓練 PC 提供了一種替代方案，其效果可以媲美傳統的學習方法。

到目前為止，PIC 的構建僅限于從具有連續潛變量（LVs）的概率圖模型（PGMs）[27] 進行簡單的編譯過程 [18]。簡而言之，PGM 中的潛變量節點成為 PIC 模型中的積分單元，而 PGM 中的（條件）分布則成為與 PIC 的輸入和積分單元相連的函數 fu，如我們在圖 1 中所示。然而，為了避免分層求積過程導致 QPC 呈指數級增長從而阻礙學習過程，PGM 的結構必須被限制為一棵樹。這使得當前的 PIC 具有了作為簡單潛變量樹模型 [3] 的語義，也明顯限制了其表達能力，因為更復雜的潛變量之間的交互無法實現。要構建更具表達能力的 PIC，就需要重新詮釋這一語義并引入新的工具，我們將在下文中完成這一點。

3 構建、學習與擴展 PICs

在第 3.1 節中，我們系統化了 DAG 形狀 PIC 的構建過程，展示了如何從任意變量分解出發來構建它們，從而超越當前最先進的方法 [18]。隨后，在第 3.2 節中，我們展示了如何使用編碼分層求積過程的 QPC 來學習和近似此類 PIC，并復現了先前工作中提出的 PC 架構 [36]。最后，在第 3.3 節中，我們提出了（神經）函數共享技術，用以參數化 PIC，使得其 QPC 實例化過程快速且高效，從而實現更大模型和更大數據集上的擴展。

3.1 從任意變量分解構建 PICs

標準的概率電路（PCs）可以根據已建立的流程進行構建，這些流程能夠靈活地表示任意變量分解以及豐富的離散潛變量（LV）交互 [36, 41]。在下文中，我們推導出一個類似的用于構建 PIC 的流程，使其能夠處理連續 LV，從而超越當前僅限于樹狀語義的 PIC（見第 2 節），同時仍能在不使實例化的 QPC 規模爆炸的前提下執行分層求積。為此，我們首先形式化“層次變量分解”或“區域圖”的概念，基于此我們將構建我們的 PIC 結構。

定義 1（區域圖（RG）[15]）：
一個關于輸入變量 X 的區域圖 R 是一個二分的、有根的有向無環圖（DAG），其節點要么是表示 X 子集的區域（regions），要么是指定某個區域如何被劃分為其他區域的劃分（partitions）（見圖 2(a)）。

區域圖可以（i）從概率圖模型（PGMs）中編譯而來 [5, 27, 3, 31]，（ii）隨機初始化 [43, 16]，（iii）從數據中學習得到 [15, 19, 40, 57]，或者（iv）根據數據模態（如圖像）構建 [45, 42, 36]。如果我們像圖 1 那樣從一棵 PGM 樹中編譯出來，那么生成的 RG 也將是一棵樹，從而產生一個類似樹結構的 PIC [18]。

我們在算法 1 中詳細描述的流程接受任意 DAG 形狀的 RG 作為輸入，并能輸出類似 DAG 的 PIC。不失一般性，我們假設有一個 RG R，它只允許（i）對區域進行二元劃分（即所有乘積單元將有兩個輸入單元），以及（ii）單變量葉子節點，如圖 2(a) 所示。

3.2 通過張量化 QPC 學習 PIC

給定一個 DAG 形狀（難以處理的）PIC，我們現在展示如何使用編碼分層求積過程的張量化 PC 來近似它，即 QPC。直觀上，我們將 PIC 解釋為編碼一組準張量 [49]，這是一種張量的推廣形式，其每個維度可能對應于連續潛變量（LV），從而在每個維度中具有潛在無限的條目。我們通過數值求積將其具體化為經典的張量。

我們從張量化電路的定義開始，并簡要回顧數值求積的相關知識

從 PIC 到 QPC 給定一個候選的 PIC，我們將以后序遍歷的方式對其進行遍歷，并迭代地將每個 PIC 單元與一個電路層（定義 2）相關聯，這一過程我們稱之為實例化（materialization）。
我們在算法 3 中詳細描述了這一過程，其本質上是在 PIC 單元上層次化地應用公式 (1)

每個單元編碼了一個可能作用于連續和離散變量上的函數，因此表示一個準張量（quasi-tensor），我們通過僅在求積點上對該函數進行評估來近似它，從而將其實例化為一個經典張量（如圖 3(a,b) 所示）。
為了便于進行求積計算，我們假設所有 PIC 的潛變量（LV）都具有有界的定義域Δ:=[?1,1]。
這樣一來，我們就可以對每一個所需的（多變量）近似使用相同的求積規則 (z~,w~)，從而簡化處理和表述。

3.3 利用神經函數共享擴展 PIC

QPC 的實例化過程可能會非常耗費內存和時間，這取決于以下三個因素：
(i) 我們需要實例化的函數的計算代價，
(ii) 所需函數評估的并行化程度，
(iii) 積分點的數量 K。

為了解決這些問題，我們引入了神經函數共享（neural functional sharing）[47]，即通過共享多層感知機（MLP），一次性參數化多個 PIC 單元。
這種方法使我們能夠擴展到更大的模型和數據集，相比之前的工作 [18]，我們的實例化過程更快、更節省內存。

PIC 函數共享

函數共享之于 PIC，就如同參數共享之于 PC。
這種類型的共享可以應用于一組輸入/積分單元——這些單元根據某些標準被分組在一起，且它們的函數都具有相同數量的輸入變量和輸出變量。

最后，我們介紹并應用兩種不同的單元分組方式：

1. 第一種方式是將所有輸入單元分組在一起，這種方法只有在所有輸入變量具有相同定義域時才適用。
   在這種分組下，結合 F-sharing，我們只需要實例化 K×P 個參數，并用它們來參數化每一個 QPC 輸入層。

2. 第二種方式是將所有處于 PIC 結構中相同深度的積分單元分組在一起，我們將這種方式與 C-sharing 相結合，并將其實例化為一個折疊的求和-乘積層（folded sum-product layer）。
   盡管將那些會實例化為折疊層的單元進行分組是一個自然且方便的選擇，但請注意，我們也可以對那些不會以這種方式實例化的單元進行分組。
   一旦一個 PIC 中的所有單元都被分組，就可以按組進行實例化操作。

與 [18] 類似，我們使用輕量級的多層感知機（MLPs）來參數化 PIC 的輸入單元和積分單元。
不過，我們不再為每個函數單獨使用一個 MLP，而是應用函數共享，以期使 QPC 的實例化過程更加快速且節省內存。

快速且內存高效的 QPC 實例化 結合PIC 函數共享按組實例化，使得通過數值求積訓練 PIC 成為可擴展的方法，因為我們大幅減少了 QPC 實例化所需函數評估的數量及其計算代價。我們現在可以實例化非常大的 QPC，其規模可以與近期過參數化的 PC 模型相當，但在使用較大 K 的情況下，所需的可訓練參數最多可減少 99%。這在原始的 PIC 公式 [18] 中是無法實現的，原因有兩點：(i) 整個 QPC 是一次性（one-shot）實例化的，而不是按組進行的；(ii) 沒有實現函數共享，因為每個輸入/積分函數都有自己的 MLP。

4 實驗

在我們的實驗中，首先我們評估了函數共享在通過數值積分擴展 PIC 訓練方面的有效性，并將其與標準 PC 以及不使用函數共享的 PIC [18] 進行比較。接著，借鑒以往的研究 [10, 32, 33, 18]，我們將 QPC 和 PC 作為分布估計器，在多個圖像數據集上進行對比。我們始終使用梯形積分規則進行訓練。所有實驗都在 NVIDIA A100 40GB 顯卡上完成。我們的代碼可在 github.com/gengala/tenpics 獲取。

得益于我們的流程，我們現在可以使用兩種最近為圖像數據設計的 RG（區域圖），它們相比傳統 RG 構建的結構具有更好的擴展性 [45, 43, 31]：

• 四叉樹

  （QTs）是樹狀的 RG，

• 四圖

  （QGs）是 DAG 形狀的 RG [36]。

這些 RG 是完全平衡的，因此使用算法 1 對其進行構建時，會得到平衡的 PIC 結構，適合對積分單元進行按深度方向的 C-sharing（組合共享）。我們在附錄 B 中詳細介紹了 QT 和 QG 的結構。

我們將一個張量化架構表示為[RG]-[求和-乘積層]-[K]，例如QT-CP-16，它可以像標準 PC 一樣訓練，也可以從 PIC 實例化為 QPC。我們將像素視為分類變量，因此我們的架構建模的是概率質量函數。

擴展 PIC 對于每種模型類型{PC, PIC}，我們指定一個二元組(·, ·)，其中第一個（第二個）參數表示輸入（內部）層/組所使用的共享技術，包括{F, C, N}，其中 N 表示不共享。

在圖 5 中，我們報告了使用 PC ( )、使用帶 ( ) 和不帶 ( ) 積分單元組函數 C-sharing 的 PIC 進行一次 Adam [23] 優化步驟所需的時間和 GPU 內存。

我們注意到，使用函數共享的 PIC ( ) 在擴展方面非常有效，所需的資源與標準 PC ( ) 相當；而未使用函數共享的 PIC ( ) ——如先前的工作 [18]——則慢了幾個數量級，并且在 K>64 時很快出現內存不足（OOM）。

值得注意的是，某些 PIC 的QG-TK配置 ( )，見表 D.2，所需的 GPU 內存甚至比 PC 更少，這是因為可訓練參數的數量存在顯著差異，因為 Adam 在優化過程中需要存儲這些參數的副本。

實際上，標準 PC 和無函數共享的 PIC ( , ) 的參數數量達到了數億級別（最高達 20 億+），而使用函數共享的 PIC ( ) 則擴展得更為平滑，最多僅達到 600 萬+ 參數。

我們強調，PIC 的可訓練參數數量與我們實例化時的 K 無關，只與用于參數化的 MLP 大小 M 有關，這也可以理解為評估 PIC 函數的成本。

更多（表格形式的）細節請參見附錄 D.1。

分布估計 借鑒以往研究 [10, 32, 33, 18]，我們廣泛測試了 QPC 和 PC 作為分布估計器在標準圖像數據集上的表現。

我們的完整結果見附錄 D.2，此處我們僅報告表現最佳模型的每維比特數（bits-per-dimension, bpd），這些模型均屬于QG-CP架構，反映了 DAG 形狀 RG 具有更強的表達能力。

我們的完整結果還表明，我們為 PIC 引入的更具表達力但計算代價更高的 Tucker 層在 K 較小時表現最好，但在擴展時較為困難。

所有 QPC 都是從應用了輸入單元 F-sharing 和積分單元組 C-sharing 的 PIC 實例化而來，即 PIC(F, C)。

我們首先以 MNIST 家族數據集開始，它包含 6 個灰度 28x28 圖像數據集：MNIST [29]、FASHIONMNIST [55]，以及帶有 4 個劃分的 EMNIST [7]。

圖 6 顯示，QPC 總體表現最佳，在 6 個任務中有 5 個優于標準 PC，在 6 個任務中有 5 個優于基于復雜啟發式學習的 PC 方法（如剪枝與生長），并在全部 6 個任務中優于一些深度生成模型（DGMs）。

接下來，我們轉向更大的 RGB 圖像數據集：CIFAR [28]、ImageNet32、ImageNet64 [14] 和 CelebA [34]。

為了與之前的工作 [32, 33] 進行比較，我們需要使用YCoCg 變換對數據集進行預處理，這是一種有損的顏色編碼方式，能持續提升 PC 在 RGB 圖像上的性能。我們還報告了使用無損 YCoCg-R 變換 [39] 預處理的數據集結果，有效地將數據集數量翻倍。

關于這些變換的詳細信息請參見附錄 C.3。

從表 1 中可以看出，QPC 再次表現出很強的競爭力，一致優于使用 Adam 常規訓練的標準 PC，也優于文獻中表現最好的 PC 模型 HCLT（后者通過 EM 方案和分塊方法訓練 [32]）。

此外，QPC 的表現接近于通過潛變量蒸餾（LVD，表 1 中 LVD-PG）訓練的 PC [32, 33]，這是一個需要額外監督其潛空間的框架，通過從現有深度生成模型（DGMs）中提取信息實現。

該技術需要預訓練的 DGM、多種啟發式方法，以及通過 EM 或 SGD 進行最終微調階段，而我們的 PIC 訓練方法則是端到端且自包含的。

5 討論與結論

在本工作中，我們系統化了 PIC 的構建過程，將其擴展到了 DAG 形狀結構（第 3.1 節），進行了張量化處理（第 3.2 節），并使用函數共享技術擴展了其訓練能力（第 3.3 節）。

在我們的實驗中（第 4 節），我們展示了當將 PIC 的可追蹤近似模型 QPC 用作分布估計器時，這一流程表現得非常有效。這反過來也成為一種新的、有效的用于大規模學習 PC 的工具。

事實上，之前的研究已經表明，通過 EM 或梯度上升方法直接訓練大型 PC 是具有挑戰性的，而且隨著模型規模的增加，PC 的性能會趨于飽和 [10, 8, 32, 33]。

我們的貢獻超越了這些限制，同時提供了一種簡單、原理清晰且完全可微的流程，其性能可以媲美更復雜的替代方案 [10, 32]（見第 4 節）。我們推測，這種效果的產生是因為通過 PIC 來訓練 PC 極大地縮小了搜索空間，同時允許：
(i) 更平穩的訓練動態，以及
(ii) 實例化任意大的可追蹤模型。

可追蹤模型的發展是機器學習中的一個重要任務，因為它們提供了多種推理機制，并可用于許多下游應用，例如表格數據建模 [8]、生成建模 [42]、無損壓縮 [30]、遺傳學 [11]、知識圖譜 [37]、受限文本生成 [58] 等等。

我們的工作也與張量網絡和**(準) 張量分解** [51, 26, 22, 6, 49] 有一定關聯，正如 [36] 最近所展示的那樣，使用張量化電路的語言可以表示分層張量分解。

此外，我們注意到最近提出的非單調 PC[35]（即具有負求和參數的 PC）也可以看作是從某些函數可能返回負值的 PIC 中進行求積過程的結果。

我們的工作并非沒有局限性。盡管我們展示了使用函數共享訓練 PIC 所需的資源與標準 PC 相當，但像 VAE、流模型（flows）和擴散模型（diffusion models）這類傳統的連續潛變量模型仍然更具擴展性。

此外，目前還無法從 PIC 中進行采樣，因為我們無法對我們的（多頭）MLP 進行可微采樣。

未來的工作可能包括探索更高效的 PIC 訓練方式，例如使用 LVD 或變分推斷 [24] 技術來直接最大化 PIC 的下界，在最后階段僅使用數值求積作為微調步驟，以蒸餾出高性能的可追蹤模型。

我們相信，我們的工作將在生成建模領域，特別是可追蹤模型的研究方向上，激發新的研究進展。

B 區域圖（Region Graphs）

在算法 B.1 中，我們詳細描述了四叉樹（Quad-Tree，QT）和四圖（Quad-Graph，QG）區域圖 [36] 的構建過程。
具體來說，將輸入標志isTree設置為True（對應 QT）或False（對應 QG）即可構建相應的區域圖。

直觀上，這些區域圖通過遞歸地將圖像分割成小塊（patches），直到每個區域恰好對應一個像素為止。
分割操作同時在水平方向和垂直方向進行，后續的小塊可以共享（從而生成非樹狀的區域圖，即 QG），也可以不共享（生成樹狀區域圖，即 QT）。

C.3 YCoCg 顏色編碼變換

在圖 C.1 和圖 C.2 中，我們提供了實驗中所使用的無損和有損版本 YCoCg 變換的 PyTorch 代碼（見第 4 節）。在圖 C.3 中，我們展示了如何應用這些變換，并表明有損版本的平均誤差小于一個比特。最后，在圖 C.4 中，我們展示了將這兩種變換應用于 RGB 圖像時，在視覺效果上的顯著差異。

D.2 附加的分布估計結果

在本節中，我們報告了所有實驗的表格結果。

請注意，標準 PC 的每個輸入層都由一個 K×P 的矩陣進行參數化，其中 P 是類別數量：對于灰度圖像數據集為 256，對于 RGB 圖像則為 768（即 256×3）。

我們發現，在所有像素之間共享一個單一的輸入層，會使灰度圖像的性能（略有）下降（詳見表 D.3）。
相比之下，我們發現這種共享方式對 RGB 圖像數據集的性能有明顯提升。

除了提升 RGB 圖像數據集的性能之外，這種共享方式還將可訓練參數的數量從 D×K×P 顯著減少到僅 K×P，其中 D 是像素數量。

例如，如果我們為一個用于 64x64 圖像、K=256 的張量化架構的所有輸入層分別參數化，則需要

256×64×64×768=805,306,368

個參數；而如果應用共享，則只需 3,145,728 個參數。

因此，如果不采用這種共享方式，我們在 GPU 上甚至無法擴展到較大的張量化架構（如 QG-CP-512）。

所有 QPC 都是從 PIC 實例化而來，其中對輸入單元組使用了F-sharing（全共享），對積分單元組使用了C-sharing（組合共享）。

我們廣泛地將 QPC 和 PC 作為密度估計器在多個圖像數據集上進行了比較，并在表 D.3、表 D.4 和表 D.5 中報告了測試集上的每維比特數（bits-per-dimension, bpd）。

原文鏈接： https://arxiv.org/abs/2406.06494v2