Scaling Probabilistic Circuits via Data Partitioning

通過數據劃分擴展概率電路的規模

https://arxiv.org/pdf/2503.08141

https://github.com/J0nasSeng/federated-spn

摘要

概率電路（Probabilistic Circuits, PCs）使我們能夠在一組隨機變量上學習聯合分布，并以可追蹤的方式執行各種概率查詢。盡管可追蹤性使概率電路能夠超越諸如貝葉斯網絡之類的不可追蹤模型，但將概率電路的訓練和推理擴展到更大、更真實的 數據集仍然具有挑戰性。為了解決這一問題，我們展示了如何在多臺機器上學習概率電路，方法是對分布式數據集進行遞歸劃分，從而揭示了概率電路與聯邦學習（Federated Learning, FL）之間深刻的聯系。這催生了“聯邦電路”（Federated Circuits, FCs）——一種新穎且靈活的聯邦學習框架，它（1）允許在分布式學習環境中擴展概率電路；（2）加速概率電路的訓練；（3）首次在一個框架中統一了橫向聯邦學習、縱向聯邦學習和混合聯邦學習，通過將聯邦學習重新表述為對分布式數據集的密度估計問題。我們展示了聯邦電路在多個大規模數據集上擴展概率電路的能力。此外，我們在多個分類任務中展示了聯邦電路在統一框架下處理橫向、縱向和混合聯邦學習的靈活性與通用性。

1 引言

概率電路（Probabilistic Circuits, PCs）是一類模型，能夠在各種概率查詢中提供可追蹤的推理能力 [Poon 和 Domingos, 2011；Choi 等，2020]。這是通過在一個計算圖上表示聯合分布并施加某些結構性質來實現的。雖然 PC 在計算效率方面優于傳統的概率模型如貝葉斯網絡 [Pearl, 1985]，但通過優化 PC 表示的緊湊性并針對特定硬件架構進行定制，可以進一步提升其性能 [Peharz 等，2020a；Liu 等，2024]。然而，另一種自然的擴展方式——將模型分布在多臺機器上——目前尚未被充分探索。像神經網絡這樣的模型可以相對容易地在多個機器上進行劃分，而劃分 PC 則更具挑戰性，因為它們必須滿足一定的結構性約束以確保所表示聯合分布的有效性。

有趣的是，我們發現 PC 的結構與聯邦學習（Federated Learning, FL）范式之間存在內在聯系。在 PC 中，和節點 （sum nodes）通過混合的方式結合相同變量集上的概率分布。這類似于橫向聯邦學習 （horizontal FL）場景 [Kone?ny 等，2016；Li 等，2020]，其中所有客戶端擁有相同的特征但樣本不同。相比之下，在縱向聯邦學習（vertical FL）[Yang 等，2019；Wu 等，2020] 場景中，樣本是共享的，但特征在客戶端之間被分割，這可以與 PC 中的積節點 （product nodes）建立聯系，這些節點用于組合不相交變量集的分布。因此，混合聯邦學習（hybrid FL）[Zhang 等，2020] 場景（即樣本和特征都在客戶端間分割）可以用和節點與積節點的組合來表示。因此，PCs 能夠以一種統一的方式來連接三種 FL 設置——這一目標在 FL 社區中被認為難以實現 [Li 等，2023a；Wen 等，2023]。

基于這種聯系，我們提出了聯邦電路 （Federated Circuits, FCs），這是一種新的 FL 框架，它將 FL 重新定義為一個在多個機器（以下簡稱“客戶端”）上分布的數據集集合上的密度估計問題。FCs 自然地處理了所有三種 FL 設置，因此提供了一種靈活的方式來擴展 PC，使其能夠學習任意分布在一組客戶端上的聯合分布（見圖 1）。通過施加與 PC 相同的結構性質，FCs 能夠在多個機器上對諸如邊緣化和條件化等概率查詢進行可追蹤的計算。在此基礎上，我們提出了一種高度通信高效的訓練算法，利用了 FCs 設計中的半環結構。

我們的實驗評估表明，FCs 在大規模密度估計任務上優于 EiNets [Peharz 等，2020a] 和 PyJuice [Liu 等，2024]，展示了擴展 PC 的優勢。此外，在各種分類任務中，FCs 在所有聯邦設置下均優于或達到了與最先進的基于神經網絡和基于樹的方法相當的效果，證明了其在 FL 中的有效性。我們的主要貢獻如下：

1. 我們引入了 FCs，這是一個通信高效且可擴展的 FL 框架，通過將 PC 的語義映射到 FL 上，統一了橫向、縱向和混合聯邦學習。

2. 我們實際構建了 FedPCs，并演示了如何利用 FC 框架將 PC 擴展到大型現實世界數據集。

3. 我們提出了一種適用于多種學習算法的一遍訓練方案（one-pass training scheme）。

4. 我們進行了廣泛的實驗，展示了我們方法在大規模 PC 學習和執行 FL 方面的有效性。我們在表格數據和圖像數據上進行了分類和密度估計實驗。

本文后續安排如下：在簡要回顧相關工作后，我們將從概率視角介紹 FL 并引入 FCs。在結論之前，我們展示了對 FedPCs 的廣泛實驗評估。我們的代碼已公開發布于 https://github.com/J0nasSeng/federated-spn.git 。

2 預備知識與相關工作

接下來，我們將簡要介紹概率電路（Probabilistic Circuits, PCs）和聯邦學習（Federated Learning, FL），并概述相關的研究工作。

已有多個研究致力于擴展概率電路的規模。Peharz 等人 [2020b] 表明可以通過使用大型隨機結構代替學習得到的結構，使概率電路適用于更大規模的問題。此外，模型架構的改進也被證明可以顯著提升計算效率，例如使用 einsum 操作實現的可并行層 [Peharz et al., 2020a] 和減少 I/O 操作的方法 [Liu et al., 2024]。Liu 等人 [2022] 則通過引入潛在變量蒸餾方法，并利用深度生成模型提供額外監督，提升了概率電路的性能。

聯邦學習（Federated Learning） 。在聯邦學習（FL）中，一組數據擁有者（或稱為客戶端）希望在不共享各自數據的前提下協同訓練一個機器學習模型。根據數據劃分方式的不同，聯邦學習可分為橫向聯邦學習、縱向聯邦學習以及混合聯邦學習三種類型：

• 在 橫向聯邦學習 中，數據集 被劃分為若干部分，每個客戶端持有相同的 d 個特征，但對應不同的、互不重疊的樣本集合。
  

• 在 縱向聯邦學習 中，數據集被按特征維度劃分，每個客戶端持有相同的 n 個樣本，但對應不同的、互不重疊的特征子集。

• 混合聯邦學習

  則是橫向與縱向聯邦學習的結合，此時客戶端可能同時持有不同（但可能部分重疊）的樣本集合和特征集合 [Wen et al., 2023; Li et al., 2023a]。

對于這三種聯邦學習（FL）設置，已有專門設計的方法來實現模型的協同學習。

在橫向聯邦學習 中，最常用的方法是在訓練過程中定期對所有客戶端的模型進行平均 [McMahan 等人, 2016; Karimireddy 等人, 2020a,b; Sahu 等人, 2018]。然而，模型平均要求每個客戶端使用相同的模型結構。

在縱向聯邦學習 中，客戶端持有不同的特征集；因此，無法保證客戶端之間可以共享相同的模型結構。在這種情況下，基于樹的方法和神經網絡模型是主要的選擇，通常通過在客戶端之間共享數據統計信息或特征表示來進行訓練 [Kourtellis 等人, 2016; Cheng 等人, 2021; Vepakomma 等人, 2018; Ceballos 等人, 2020; Chen 等人, 2020; Liu 等人, 2019]。

類似于基于樹的縱向聯邦學習方法，基于樹的混合聯邦學習 方法也在客戶端之間共享數據統計信息（如直方圖）或模型屬性（如分裂規則）[Li 等人, 2023b, 2024]。然而，基于樹的方法通常需要復雜的訓練過程。

3 聯邦電路（Federated Circuits）

本工作的目標是通過將數據和模型分布到多個機器上，來擴展概率電路（PCs）的規模。這樣可以利用計算集群的資源，以聯邦學習的方式訓練概率電路。接下來，我們將介紹一種優雅且有效的方式來實現這一目標，即使用我們提出的新型聯邦學習框架——聯邦電路（Federated Circuits, FCs） ，該框架統一了橫向、縱向和混合聯邦學習。

3.1 問題陳述與建模假設

聯邦學習中的一個關鍵假設是： 客戶端之間不能交換原始數據 。然而，在不同客戶端上的變量之間可能仍然存在依賴關系。為了在保護數據隱私的前提下學習這些“隱藏”的依賴關系，我們做出以下假設：

請注意，獨立性僅假設在數據的簇內成立 。因此，潛在變量（可以將其視為“簇選擇器”）能夠捕捉分布在不同客戶端上的變量之間的依賴關系。像假設2 中所描述形式的分布，其表達能力嚴格強于乘積分布（product distribution），并允許進行更復雜的建模。

此外，我們想強調的是，假設1和假設2在整個概率電路（PC）文獻中是常見的 。假設1 是邊緣化有效性的基礎，而假設2 在概率電路的結構學習算法構建中起著關鍵作用。更多細節請參見附錄 A。

3.2 概率電路與聯邦學習的聯系

接下來，我們展示概率電路（PC）語義與聯邦學習（FL）之間的一種內在聯系，這種聯系使我們可以通過將數據劃分到多個客戶端上來擴展概率電路在大規模數據集上的應用。

求和節點與橫向聯邦學習 。在橫向聯邦學習中，假設每個客戶端持有相同的特征集合，即對于所有客戶端 c,c′∈C，都有 Xc=Xc′。然而，每個客戶端持有的樣本不同。典型的橫向聯邦學習方法會在訓練過程中定期聚合各本地模型的參數。

但事實上，橫向聯邦學習的設定恰好對應于概率電路中求和節點（sum node） 的語義解釋：
一個求和節點將數據集劃分為多個互不重疊的簇，從而形成一個由這些簇上學習得到的混合分布來表示整體數據。因此，我們不是聚合模型參數，而是聚合每個客戶端在其數據分區上學習到的分布。

這種對橫向聯邦學習的理解帶來了令人滿意的積極副作用：
如果各客戶端的數據分布差異較大，聚合模型參數 可能會在訓練過程中導致模型發散。而由于我們是在混合分布中聚合分布 ，因此可以自然地處理客戶端之間分布異構的問題。此外，由于客戶端可以獨立地訓練模型，因此訓練過程中的通信開銷也被最小化。

乘積節點與縱向聯邦學習 。在縱向聯邦學習中，假設每個客戶端持有互不重疊的特征集合，即對于所有客戶端 c,c′∈C，都有 Xc∩Xc′=?。與橫向聯邦學習不同的是，所有客戶端持有的是屬于相同樣本實例 的不同特征。

如同在橫向聯邦學習中一樣，縱向聯邦學習與概率電路之間也存在語義上的聯系：
概率電路中的乘積節點計算的是定義在一組互不重疊隨機變量上的乘積分布。因此，乘積節點沿特征維度對數據進行劃分，這正好對應于縱向聯邦學習的設定。

然而，乘積節點假設其子節點所對應的隨機變量之間是相互獨立的。顯然，在縱向聯邦學習中這是一個不現實的假設，因為不同客戶端所持有的特征可能是統計相關的。我們可以利用假設2 來捕捉這些依賴關系，并構建一個由多個獨立簇的乘積組成的混合分布。詳見第 3.3 節以了解細節。

聯邦電路（FCs）在概率電路（PCs）的基礎上進行了擴展，其核心在于：FCs 表示一個分布在多個機器上的計算圖 G=(V,E)，其中每個節點 N∈V 可以執行任意的操作。請注意，通過圖 G，FCs 同時也定義了參與計算的各機器之間的通信網絡結構。

此外，FCs 并不限于上述的概率解釋。例如，若將葉節點參數化為決策樹，并引入一個執行平均操作的節點 N，那么整個模型就變成了一種Bagging 模型 。

3.3 聯邦概率電路（Federated Probabilistic Circuits）

接下來，我們將深入探討 FCs 的概率解釋。為此，我們提出一種具體的實例化方法：使用概率電路（PCs）作為葉節點模型來構建聯邦概率電路（Federated PCs, FedPCs） 。

基于第 3.2 節中的概率解釋，我們將 PC 的結構與通信網絡的結構進行對齊，從而構建出一個聯邦形式的概率電路。

請注意，只有客戶端節點 C 持有數據集 ，并且我們只要求客戶端由密度估計器進行參數化。為了使 FedPCs 具備計算效率，這些密度估計器應當是可追蹤的（tractable）。在下文中，我們將葉節點 C 參數化為概率電路（PCs）。

分配函數 ψ將 PC 的計算圖轉換為分布式計算圖，從而誘導出一個通信網絡。這就在 PC 語義（即計算圖）與 FedPCs 中的通信網絡結構之間建立了直接對應關系。

在 FedPCs 中，推理過程與在傳統 PC 中一樣，通過將似然值從葉節點傳播到根節點來執行。唯一的區別在于：如果某個節點 N′ 是節點 N 的父節點，并且 ψ(N)=ψ(N′)，那么節點 N 的計算結果必須通過通信網絡發送給其父節點 pa(N)。

訓練 FedPCs 需要對標準 PC 的訓練流程進行調整，因為在聯邦學習（FL）設置中，客戶端無法訪問其他客戶端的數據。例如，使用期望最大化（EM）算法進行訓練需要所有客戶端訪問相同的樣本，這與橫向和混合聯邦學習不兼容。同樣地，LearnSPN [Gens and Domingos, 2013] 在訓練過程中依賴于特征之間的獨立性檢驗，因此要求訪問所有特征，這也無法滿足 FL 的設定。

為了解決這個問題，我們提出了一種單遍訓練（one-pass training）方法 用于 FedPCs。

單遍訓練

我們的單遍學習算法可以獨立地訓練 FedPCs 的結構和參數（見算法1，圖2）。訓練開始前，所有客戶端 c∈C向服務器共享它們所持有的、具有唯一標識的特征/隨機變量集合 Xc，從而構建出一個大小為 ∣C∣×∣X∣ 的特征存在指示矩陣 M（第1–2行）。

特征標識符可以是特征的名稱，如“賬戶余額”，并且在所有客戶端上必須對應于相同的隨機變量（從而具備唯一可識別性）。

然后，服務器將聯合特征空間 X 劃分為若干互不重疊的子空間 S(j)。為此，我們考慮矩陣 M 中的所有不同列向量集合 U，記其中每個不同的向量為 u。由于矩陣 M 的每一列表示一個特征所在的客戶端集合，我們可以利用每一個 u∈U 來計算一組被相同客戶端集合所持有的特征。這樣就得到了 ∣U∣ 個不同的特征集合，記作 表示持有該特征集合的客戶端集合（第3–7行）。

這一過程如圖2（頂部）所示。

因此，網絡側的權重可以在不需要任何前向或反向傳播的情況下直接推斷出來。

請注意，這種方法將橫向聯邦學習簡化為對客戶端數據分布的混合建模，而將縱向聯邦學習簡化為對 P 個乘積節點的混合建模。

3.4 通信效率分析

在對大規模劃分數據上的模型進行訓練時，通信效率 是一個關鍵要求。接下來，我們分析 FedPCs 的通信效率。

橫向聯邦學習（Horizontal FL）

假設有一組客戶端 C，每個客戶端持有一個包含 M 個參數的模型。再假設在整個訓練過程中模型被聚合 K 次（即有 K 輪通信）。那么，在橫向聯邦學習中常用的基于模型聚合的算法（如 FedAvg）通常在網絡上傳輸的信息量為 O(M?∣C∣?K)，因為每個客戶端在每輪通信中都要向服務器發送 M 個模型參數。

相比之下，使用單遍訓練方法訓練 FedPCs 只需在網絡上傳輸 O(∣C∣?(M+1)) 條消息。這是因為模型是在本地獨立訓練的，之后只需設置求和節點的參數（共 O(∣C∣) 條消息），并將模型在服務器上聚合（共 O(M?∣C∣)條消息）。

縱向聯邦學習（Vertical FL）

在縱向聯邦學習場景中，常使用類似 SplitNN 的架構。假設訓練一個 SplitNN 架構共計 E 個訓練輪次，對于一個擁有 S 個樣本的數據集，每個樣本輸出一個大小為 F 的特征向量，并且這些特征在客戶端 C 上縱向分布。這種訓練方式在網絡上傳輸的消息數量為 O(E?∣C∣?F?S)。

而使用 FedPCs 的單遍訓練方法，每個客戶端會為所學習到的 K 個簇各自訓練一個具有 M 個參數的專用 PC。FedPC 的最后一層是由 P 個簇的乘積組成的混合模型?；旌夏Ｐ偷膮翟诿總€客戶端模型訓練完成后設定。聚合所有已學習模型并設定網絡側的混合參數所需的總消息數為 O(K?M?∣C∣+P)。

如果滿足不等式 (K?M+P∣C∣)<(E?F?S)，則 FedPCs 的訓練比訓練 SplitNN 類架構更具通信效率。在實際應用中，這一條件通常成立：

• 簇的數量通常小于 100；

• 特征向量可能具有數百維（即 F>100 ）；

• 模型參數數量應小于數據集中的樣本數以保證泛化能力（即 M

• 參數 P 可根據 ∣C∣ 和數據靈活設定。

附錄 C 提供了更多關于通信成本的細節和直觀解釋。

混合聯邦學習（Hybrid FL）

在混合聯邦學習中，FedPCs 在多個子空間上進行訓練：

• 一些子空間存在于所有或部分客戶端上（記作 Rs ）；

• 一些子空間僅存在于單個客戶端上（記作 Rd ）。

進一步將 FedPCs 在橫向聯邦學習和縱向聯邦學習中的通信成本分別記為 Ch 和 Cv。由于混合聯邦學習的訓練過程本質上是對共享特征空間執行橫向聯邦學習、對不重疊特征空間執行縱向聯邦學習，因此整個訓練過程中在網絡上傳輸的消息總數為 O(∣Rs∣?Ch+∣Rv∣?Cv)。

4 實驗

我們的實驗證明，FedPCs 可以通過數據和模型劃分有效擴展概率電路（PCs）。在同一個統一框架下執行橫向、縱向和混合聯邦學習，我們獲得了與主流聯邦學習基線方法相同或更優性能的高性能模型。

我們旨在回答以下問題：

• (Q1)

  FedPCs 是否能夠減少訓練時間，并成功學習分布式數據上的聯合分布？

• (Q2)

  FedPCs 是否能有效擴展 PC 模型，從而獲得更具表達能力的模型？

• (Q3)

  不同參數化的 FCs 在分類任務中相比現有聯邦學習方法表現如何？

• (Q4)

  我們的單遍學習算法與使用 EM 算法的訓練方式相比如何？

為了驗證 FedPCs（FC 的一個實例）是否能夠成功擴展 PCs，我們參考 Liu 等人 [2024] 的做法，在三個大規模、高分辨率圖像數據集上進行密度估計實驗：ImageNet、ImageNet32（各約 120 萬個樣本）和 CelebA（20 萬個樣本）。這些數據集被橫向劃分到 2 至 16 個客戶端上。我們將 FedPCs 與 EiNets 和 PyJuice 進行比較，并對所有方法運行了 5 次不同隨機種子的實驗。

為了評估 FCs 在聯邦學習場景中的表現，我們選擇了三個表格數據集，它們覆蓋了現實世界中的多種應用場景和數據規模：一個是信用卡欺詐檢測數據集（約 30 萬個樣本），一個是醫療數據集（乳腺癌檢測；不到 1000 個樣本），還有一個是廣受歡迎的 Income 數據集（超過 100 萬個樣本）。所選的數據集涵蓋了小數據、中等數據和大數據三種情況（詳見附錄 D）。我們的評估既包括平衡數據集（乳腺癌），也包括不平衡數據集（Income 和 Credit）。

我們選擇表格數據集是因為它們非常適合用于研究 FCs 在橫向、縱向和混合聯邦學習設置下的表現，并且代表了多種現實應用。我們將其與多個廣泛使用的強基線方法進行了比較。作為神經網絡架構的參數化，我們使用了專為表格數據設計的 TabNet [Arik 和 Pfister, 2020]。我們在廣泛使用的 FedAvg（橫向 FL）和 SplitNN（縱向 FL）框架下訓練這些網絡。此外，由于樹模型在表格數據集上表現優異，我們還將 FCs 與 FedTree [Li 等人, 2023b] 進行了比較。詳細信息請參見附錄 D。

(Q1) FedPCs 能在更短的時間內學習分布式數據的聯合分布

首先，我們驗證了 FedPCs 是否能在多個客戶端分布的數據上正確且高效地執行密度估計。為此，我們將多個數據集分別分配給對應于橫向（5 個客戶端） 、縱向（2 個客戶端） 和 混合聯邦學習（2 個客戶端） 的客戶端集合上。

為了展示 FedPCs 對標簽偏移（FL 中常見現象）的魯棒性，在橫向設置中，每個客戶端只接收來自部分類別的數據，并在其本地樣本上獨立訓練 PC。

在縱向設置中，我們對數據進行了劃分，使得客戶端的特征空間互不重疊，但每個客戶端持有相同的樣本。

在混合設置中，數據被劃分為客戶端之間的特征空間和樣本空間存在部分重疊（但不是完全重疊）。

對于所有表格數據集，FedPC 的葉節點使用了 MSPNs [Molina et al., 2018] 參數化，這是一種屬于 PC 模型家族的方法，能夠在混合數據域（即連續和離散隨機變量）上執行密度估計。

我們選用 MSPNs 作為集中式模型，其結構使用 LEARNSPN（一種遞歸貪心結構學習算法）進行訓練 [Gens 和 Domingos, 2013]。對于 MNIST 數據集，我們在所有設置中都使用了具有高斯密度的 EiNets 作為 PC 的實現。

請注意，FedPCs 的大小被設定為大致匹配集中式模型，即沒有進行模型放大。

表1 將集中式 PC 在完整數據集上的訓練對數似然值和相對運行時間，與 FedPC 在不同聯邦學習設置下的結果進行了對比。結果顯示，FedPCs 在表格數據集上實現了與集中式 PC 相當的對數似然值，但訓練速度顯著更快。因此，我們肯定地回答了問題 (Q1)。

(Q2) FedPCs 能有效擴展 PC 模型

為了檢驗 FedPCs 是否能夠有效擴展 PC 模型，我們在 CelebA、ImageNet32 和 ImageNet 數據集上分別訓練了 EiNet、PyJuice 和 FedPC。所有模型均采用 Poon-Domingos（PD）架構。

FedPC 使用 EiNets 參數化，數據分布在 2、4、8 和 16 個客戶端之間。

FedPC 模型和基線模型（EiNets 和 PyJuice）的選擇確保了它們都能適應單個 GPU（系統細節見附錄 D）。

EiNets 和 FedPC 使用高斯分布作為葉節點分布，而 PyJuice 使用分類分布。

這些參數化選擇基于經驗觀察：EiNets 和 FedPC 使用高斯分布效果最好，而 PyJuice 使用分類分布效果最佳。

在 FedPC 訓練中，圖像被隨機水平劃分，每個客戶端擁有大致相同樣本數量的數據集。葉節點和所有基線模型均使用 EM 算法進行訓練。

表2 展示了 EiNets、PyJuice 和 FedPC 在測試集上每樣本每維度的 nats 分數?？梢钥闯?，隨著參與客戶端數量的增加，模型規模也隨之增大，三組數據集上的密度估計性能均有提升。

我們認為這是因為更大規模的模型具有更高的表達能力，使其能夠更好地捕捉數據的統計特性。同時，更大的模型在測試集上取得更高的 nats 分數也表明并未因更多參數而導致過擬合。

然而需要注意的是，如果進一步擴大模型規模，可能會導致過擬合。以原則性的方式找到最優模型大小/客戶端數量超出了本文范圍，留待未來研究。

除了更好的建模性能外，更多的客戶端還顯著減少了訓練時間（見圖3）。因此，FedPCs 能夠高效地將可追蹤的概率模型擴展到大規模數據集上。

(Q3) FCs 在聯邦學習中實現了最先進的分類結果

聯邦電路（FCs）可以在葉節點使用不同的模型進行參數化。我們考察了兩種參數化方式，在三個表格數據集上執行聯邦分類任務。

首先，我們使用了在 (Q1) 中介紹的 FedPC（即 FC[PC]），它可以通過概率電路中條件分布的可追蹤計算來解決判別式任務。

其次，我們考察了另一種 FC 參數化方式：決策樹（FC[DT]），這代表了一種 bagging 模型的實例化。

為了評估 FCs 在聯邦分類任務中的表現，我們將其與已知的橫向聯邦學習和縱向聯邦學習方法進行了比較。實驗是在覆蓋多種現實應用場景和分布特性的表格數據集上進行的。

我們采用了 TabNet 和 FedTree 作為強基線方法。

• 在橫向聯邦學習設置中，TabNet 使用 FedAvg 進行訓練；

• 在縱向聯邦學習設置中，TabNet 以 SplitNN 的方式進行訓練 [Ceballos 等人, 2020]。

我們將這些方法的結果與我們的單遍訓練 方法進行對比。結果顯示，FCs 在所有數據集上的表現與所選基線方法相當甚至更優（見圖4；附錄E），同時相較于基線方法具有更高的靈活性。

(Q4) 單遍訓練保持性能優勢

為了評估我們提出的單遍訓練 方法與使用標準優化算法（如 EM）訓練 PC 的方法之間的差異，我們定義了一個允許數據交換的聯邦學習設置。這是必要的，因為我們需要使用 EM 來訓練 PC 和 FedPC 架構，以便與我們的單遍訓練流程進行對比。

我們使用 RAT-SPNs [Peharz 等人, 2020b] 作為 FedPC 的葉節點參數化方式。

隨后，我們使用標準的 EM 算法（即允許數據交換）訓練了一個 FedPC，并在同一垂直劃分的數據集上使用我們的單遍訓練方法訓練了另一個結構相同的 FedPC。

我們報告了 EM 訓練和單遍訓練在測試數據集上的平均對數似然值（見表3）?？梢钥闯?，在任何情況下，單遍訓練都沒有顯著降低對數似然值。

因此，我們的結果表明，單遍訓練是更優的選擇，因為它在保持模型性能的同時具備更高的通信效率。

5 結論

在本工作中，我們提出了聯邦電路（federated circuits） ，其核心在于概率電路（PCs）與聯邦學習（FL）之間的一種內在聯系。我們展示了通過在劃分的數據上大規模訓練概率電路，可以顯著提升其訓練速度和模型表達能力。

由于我們的框架允許集成多種類型的密度估計器，因此可以無縫整合 PC 領域及其他領域中的其他模型和最新進展，從而保持聯邦學習方法在擴展性方面的持續相關性。

局限性與未來工作 盡管我們的實驗表明，擴展概率電路可以在訓練速度和性能方面帶來顯著提升，但要訓練如此大規模的模型，仍需要充足的計算資源。

未來的工作方向包括：探索除概率電路以外的其他聯邦電路（FCs）參數化方式，這具有很大的研究潛力。此外，也很值得進一步研究：這種用于混合聯邦學習的概率框架，除了用于擴展 PCs 外，是否也能為更傳統的聯邦學習應用帶來益處。

根據命題1（Prop. 1） ，我們可以看到，任何可分解且平滑的概率電路（PC）都可以表示為一個無層次結構的混合模型，即：可以將 PC 的結構壓縮為一個深度為一的結構。由于對一個可分解且平滑的 PC 進行邊緣化操作后，仍會得到另一個可分解且平滑的 PC，并且該邊緣化后的 PC 可以表示為一棵誘導樹（induced tree），因此假設1 在 PC 文獻中是一個標準假設。

從概率電路的角度也可以理解假設2 。在流行的結構學習算法中，如 LearnSPN [Gens 和 Domingos, 2013]，PC 是通過交替進行數據聚類與特征子集的獨立性檢驗來學習的。因此，像 LearnSPN 這類算法的最終目標是尋找一些簇，在這些簇中，某些隨機變量子集被認為是相互獨立的，從而最大化對數似然值。因此，假設2 與 LearnSPN 密切相關，也是 PC 建模中常見的一個假設。

B 證明

在本節中，我們為論文中的命題提供完整的證明。

B.1 FedPCs 與最大熵原理

假設2 與最大熵原理（principle of maximum entropy） 是一致的：我們的目標是在簇內尋找具有最大熵的聯合分布，同時允許客戶端所持有的隨機變量之間存在依賴關系，并確保每個客戶端的邊緣分布得以保留。

雖然可能存在多個聯合分布能夠保持邊緣分布不變，但非最大熵的解會引入額外的假設或先驗知識，從而限制了模型的靈活性。通過假設一個簇內所有變量相互獨立，我們可以通過一個由乘積分布組成的混合模型，高效地構造出最大熵分布。

對于相互獨立的變量來說，乘積分布可以最大化熵。這一點可以通過聯合熵和條件微分熵來加以證明。

C 通信效率

在聯邦學習（FL）中，跨多臺機器學習模型時，通信效率 是一個關鍵屬性。在這里，除了我們的理論結果外，我們還更直觀地提供了關于 FCs（聯邦電路）通信效率的進一步細節。

為此，我們繪制了在橫向/縱向聯邦學習設置下，訓練 FedPC 與 FedAvg/SplitNN 所需的通信成本（以 MB（兆字節） 為單位），數據集規模分別為 100 萬（1M）和 1 億（100M）樣本。

無論數據集中樣本數量多少，在橫向和縱向設置中，FedPCs 的通信效率都優于我們的基線方法（見圖5）。

D 實驗細節 D.1 數據集

以下是我們實驗中使用到的數據集。除非另有說明，數據集在客戶端上的分布方式如下：

• 在橫向設置中 ，我們或按樣本隨機劃分（用于所有二分類任務），或按標簽劃分，例如將 MNIST 中某一特定標簽（如“0”）的數據分配給一個客戶端。

• 在縱向設置中 ，我們將表格數據集沿特征維度進行隨機劃分，即每個客戶端擁有全部樣本，但僅被分配一組隨機子集的特征。對于圖像數據，我們將圖像劃分為不重疊的圖像塊，并將這些圖像塊分發給不同的客戶端。

• 在混合設置中 ，我們在特征和樣本兩個維度上對表格數據集進行劃分。具體做法是：確保至少有兩個客戶端共享至少一個隨機選擇的特征（但各自持有該特征的不同樣本）。對于圖像數據，我們將圖像劃分為有重疊的圖像塊，然后從中抽取部分樣本組成子集并分發給各客戶端。

我們使用的 Income 數據集來自：
https://www.kaggle.com/datasets/wenruliu/adult-income-dataset

該數據集是一個二分類問題，包含 14 個特征，訓練集中約有 45 萬個樣本，測試集有 900 個樣本。

我們使用 sklearn 中的 TargetEncoder 將離散變量編碼為數值型，并使用相應特征的中位數填補缺失值。此外，我們對所有特征進行了標準化處理。

Breast Cancer 數據集

我們使用的 Breast Cancer 數據集來自：
https://www.kaggle.com/datasets/uciml/breast-cancer-wisconsin-data

這是一個二分類問題，包含 31 個特征和 570 個樣本。我們將數據集劃分為 450 個訓練樣本和 120 個測試樣本，并對所有特征進行了標準化處理。

Credit 數據集（Give Me Some Credit）

我們使用的 Credit 數據集來自：
https://www.kaggle.com/c/GiveMeSomeCredit

這是一個二分類任務，包含 10 個特征，訓練集有 150 萬個樣本，測試集有 10 萬個樣本。

同樣地，我們使用 TargetEncoder 對離散變量進行編碼，并用對應特征的中位數填補缺失值，最后對所有特征進行了標準化處理。

MNIST 數據集

我們使用的是 pytorch 提供的 MNIST 數據集。它包含 70,000 張手寫數字圖像（0 到 9），尺寸為 28x28 像素（60,000 張訓練圖像，10,000 張測試圖像）。預處理階段我們對所有特征進行了標準化。

ImageNet / ImageNet32 數據集

我們使用了 pytorch 提供的 ImageNet 數據集。它包含大約 120 萬張圖像，涵蓋 1000 個類別。圖像分辨率各異，我們統一將其縮放為 64x64（ImageNet）和 32x32（ImageNet32）像素，隨后進行中心裁剪，并對所有特征進行標準化處理。

樣本被隨機分配到各個客戶端，作為簡單的數據集劃分方式。

D.2 離散化（Discretization）

在我們的實驗設置中，聯邦電路（FCs）和 EiNets 使用高斯葉節點進行參數化，并擬合 RGB 圖像數據。由于圖像數據是離散的 （像素值為 0 到 255 的整數），而高斯分布定義在連續域上，因此它表示的是概率密度函數 而非概率質量函數 。

為了獲得給定圖像 x 的概率，我們需要對高斯葉節點進行離散化處理 。具體做法是構造 255 個“桶”（buckets），通過以下公式計算概率質量：

其中 Φ 是標準正態分布的累積分布函數（CDF）。

計算圖結構將保持不變，因為概率電路（PCs）的概率語義既適用于密度函數，也適用于概率質量函數。

D.3 訓練與超參數（Training & Hyperparameters） PyJuice

對于 PyJuice，我們遵循 [Liu et al., 2024] 的訓練設置，在隨機裁剪的 32x32 圖像塊上進行訓練。

• 對于 ImageNet32 數據集，這對應于整張圖像；

• 對于 CelebA 和 ImageNet，則對應圖像中的一個隨機 32x32 區域。

我們使用 EM 算法作為優化方法。

FCs 和 EiNets

對于 FCs 和 EiNets，我們按照 [Peharz et al., 2020a] 的設置，在整張圖像上進行訓練：

• ImageNet32：32x32 像素；

• CelebA / ImageNet：64x64 像素。

同樣地，我們使用 EM 算法進行優化。

我們使用了 5 個不同的隨機種子（0–4）運行所有實驗。

下表展示了每個數據集及聯邦學習設置下的所有相關超參數配置。

D.4 硬件配置（Hardware）

所有實驗均在配備 Nvidia A100 (40GB) GPU 的 Nvidia DGX 服務器上進行，處理器為 AMD EPYC 7742 64 核心處理器 ，內存為 2TB RAM 。

E 更多實驗結果（Further Results）

在此部分，我們提供了關于聯邦電路（FCs）的更多實驗細節。

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