What is the Relationship between Tensor Factorizationsand Circuits (and How Can We Exploit it)?
張量分解和電路之間有什么關系(以及我們如何利用它)?
https://arxiv.org/pdf/2409.07953
https://github.com/april-tools/uni-circ-le
摘要
本文在電路表示(circuit representations)與張量分解(tensor factorizations)這兩個看似不同但本質上密切相關的領域之間建立了嚴格的聯系。通過連接這兩個領域,我們指出了能夠使兩個社區都受益的一系列機會。
我們的工作在電路語言(circuit language)中推廣了流行的張量分解方法,并在一個統一的廣義分層分解框架下整合了多種電路學習算法。
具體而言,我們引入了一種模塊化的“樂高積木”方法來構建張量化(tensorized)的電路架構。這種方法使我們能夠系統地構造和探索各種電路與張量分解模型,同時保持其可處理性。
這一聯系不僅澄清了現有模型之間的相似性與差異,還為構建和優化新的電路/張量分解架構提供了一個全面的流程。
我們通過廣泛的實證評估展示了該框架的有效性,并強調了張量分解在概率建模中的新研究機會。
1 引言
本文旨在連接兩個看似遙遠、實則緊密相關的領域:電路表示(circuit representations)與張量分解(tensor factorizations)。具體來說,我們在兩種表示之間建立了正式聯系,并展示了后者如何為前者的學習算法提供統一視角,同時也為兩個研究社區帶來新的研究機會。
張量是矩陣的多維推廣,廣泛用于表示高維數據(Kroonenberg, 2007)。張量分解是一種數學上已被深入理解的方法,它通過低維張量上的簡單操作來緊湊地表示高維張量(Kolda, 2006)。
它們在機器學習和人工智能中被廣泛應用,例如在計算機視覺(Vasilescu & Terzopoulos, 2002;Savas & Eldén, 2007;Panagakis 等人, 2021)、圖分析(Kolda 等人, 2005)、計算神經科學(Vos 等人, 2007;Tresp 等人, 2021)、神經符號 AI(Nickel 等人, 2015;Balazevic 等人, 2019;Gema 等人, 2023;Loconte 等人, 2023)、語言建模(Ma 等人, 2019;Hu 等人, 2022;Xu 等人, 2023),以及作為編碼概率分布的方式(Jaini 等人, 2018b;Novikov 等人, 2021;Amiridi 等人, 2022;Hood & Schein, 2024)。雖然通常以淺層分解的形式定義,但張量分解也可以表達為分層結構(Grasedyck, 2010),有時用張量網絡(Orús, 2013;Biamonte & Bergholm, 2017;Glasser 等人, 2019)的圖形形式表示。
另一方面,電路表示(Darwiche & Marquis, 2002;Choi 等人, 2020;Vergari 等人, 2021)是在邏輯推理和概率建模背景下提出的結構化計算圖(Darwiche, 2003;Poon & Domingos, 2011;Kisa 等人, 2014)。其中,概率電路(PCs)(Vergari 等人, 2019b;Choi 等人, 2020)是能編碼可處理概率分布的電路。它們支持多種需要精確高效推理的應用,如無損壓縮(Liu 等人, 2022)、生物醫學生成建模(Dang 等人, 2022b)、可靠的神經符號 AI(Ahmed 等人, 2022;Loconte 等人, 2023)以及受約束的文本生成(Zhang 等人, 2023)。
過去已提出許多從數據中學習 PC 的算法(參見 Sidheekh & Natarajan (2024) 的綜述),其中一種范式逐漸顯現:構建過參數化的電路,包含數百萬甚至數十億參數(Liu 等人, 2023a;Gala 等人, 2024a),并通過梯度上升、期望最大化(EM)或正則化方法訓練這些參數(Peharz 等人, 2016;2020c;Dang 等人, 2022a)。
層次化張量分解和概率電路都曾被作為概率圖模型的替代表示(Song 等人, 2013;Robeva & Seigal, 2017;Glasser 等人, 2020;Bonnevie & Schmidt, 2021),一些工作也暗示了某些電路與張量分解之間的聯系(Jaini 等人, 2018b;Glasser 等人, 2019)。然而,它們在應用方式上有所不同:張量分解通常用于存在一個待逼近的真實張量或可以表述為降維問題(即張量草圖)的任務,而概率電路則是以類似于生成模型的方式從數據中學習得到的。
不過,現代的概率電路表示同樣是過參數化的,并通常以張量集合的形式編碼,以便利用并行性和現代深度學習框架(Vergari 等人, 2019a;Peharz 等人, 2020c;Mari 等人, 2023)。這引發了一個問題:電路與張量分解之間是否存在正式且系統的聯系?
我們的回答是肯定的。我們證明了電路可以被表示為一種廣義稀疏的層次化張量分解,其中其參數編碼了該分解中的低維張量。或者換句話說,層次化張量分解是一種具有特定張量化架構的深層電路的特例。
對于概率電路而言,這意味著將表示為非負張量的概率分布進行分解(Cichocki & Phan, 2009)。同時,經典張量分解也可以被準確地編碼為(淺層)電路。通過確認張量分解與電路之間的對偶性,我們系統整理了文獻中的已有結果,為電路的表示與學習提供了新視角,并指出了構建新型或擴展已有(概率)分解的可能路徑。
具體而言,本文首先推導出一種簡潔的方式來表示多種張量化的電路架構,并使用“樂高積木”方法將其表示為計算圖——通過堆疊(局部)密集的張量分解,同時保留電路所需的可處理結構性質。這使我們能夠以即插即用的方式使用新穎的“模塊”。
接著,我們將文獻中提出的眾多不同學習概率電路的算法(Peharz 等人, 2020c; a;Liu & Van den Broeck, 2021b)統一起來。這些算法來源不同,所構造的電路也被視為不同的模型。我們表明,它們的差異實際上只是其張量參數的分解方式和語法轉換,因為它們都可以基于Tucker 張量分解(Tucker, 1966)及其特例(Kolda & Bader, 2009)在一個統一的廣義(層次化)分解框架下理解。因此我們認為,文獻中經常報告的不同性能表現,更多源于不同的超參數和學習方法,而非不同的歸納偏置(Liu 等人, 2023b)。
進一步地,在建立這種聯系之后,我們利用張量分解對已經以張量格式表示的現代概率電路架構進行進一步壓縮。由此,我們引入了比以往更參數高效的概率電路,并表明為特定任務尋找最佳電路架構的問題遠未解決。
最后,我們強調這種與電路的聯系為張量分解研究社區帶來了有趣的研究機會——這些機會在本文中以方框標注的方式呈現,包括從數據中學習張量分解、將張量分解解釋為潛在變量概率模型,以及通過背景知識指定誘導稀疏性等方向。
貢獻:
- 我們將流行的張量分解方法及其層次化形式推廣到電路語言中
(第2節)。
- 我們將概率電路與非負張量分解聯系起來,并指出后者可被解釋為潛在變量模型,從而可用于生成建模和神經符號 AI
(第3節)。
- 在我們的框架下,我們抽象出構建和學習現代過參數化架構的各種選項,提出了一個通用的算法流程
(第4節),用于表示和學習張量化的層次化張量分解。
- 這使我們能夠利用張量分解分析現有不同電路參數化之間的關系,并提出更具參數效率的建模選擇,同時保留一定的表達能力
(第5節)。
- 我們在廣泛的分布估計任務中評估了我們框架中的多種算法選擇,突出了時間與空間復雜度之間的主要權衡及最終性能表現
(第6節)。
2 從張量分解到電路
2.1 淺層張量分解即淺層電路 Tucker 張量分解 張量分解通過一組低維張量來近似高維張量。形式上,給定一個張量T ∈ ??1??2?...???,其大小隨著維度呈指數增長,我們希望為其尋找一個低秩分解(Kroonenberg, 2007)。
許多流行的張量分解方法,例如典范多元分解(CP)(Carroll & Chang, 1970)、RESCAL(Nickel 等人, 2011)以及高階奇異值分解(HOSVD)(De Lathauwer 等人, 2000),都是Tucker 分解(Tucker, 1964;1966)的特例。
因此,我們對張量分解的討論將首先聚焦于 Tucker 分解,隨后介紹其層次化形式(Grasedyck, 2010)。我們的結論也將推廣到其特例,如 CP、RESCAL 和 HOSVD。
電路可以被理解為具有指數多項的多線性多項式,但它們以一個大小為多項式級的深層計算圖進行緊湊編碼(Darwiche, 2003;Zhao 等人, 2016;Choi 等人, 2020)。
從這個角度來看,我們可以直觀地理解電路與張量分解之間的聯系及其差異。事實上,正如后者也編碼了緊湊的多線性算子(見公式(2)),電路多項式的不定元(indeterminates)可以不僅僅是定義2中提到的矩陣條目,還可以是潛在的非線性輸入函數。
例如,一個電路可以編碼一組連續隨機變量上的聯合密度,而輸入函數f?可以表示高斯密度(見圖1)。關于如何在電路中編碼輸入單元的多種方式,詳見“機會4”部分的討論。
電路中函數c的求值是通過通常的前饋方式遍歷其計算圖完成的——先輸入后輸出,如圖1所示。
此外,我們所提供的電路定義可以比張量分解更具一般性,因為它能夠表示稀疏的計算圖,即其中單元之間是非規則連接的。但我們將在后文論證,這并非必需情況。
實際上,電路也可以被設計成局部稠密(locally-dense)的形式,這在許多現代實現中是常見的做法(見第4節)。
正如我們在下面針對一般 Tucker 分解(定義1)所提出的構造性命題中所展示的那樣,當把張量分解轉化為電路時,它們看起來也正是這種局部稠密結構。
附錄 A.1 詳細描述了我們的證明構造,圖2 則以一個三維張量的 Tucker 分解為例進行了說明。
簡而言之,我們構建了一個作用于相同變量上的淺層電路 c,當對其進行求值時,它輸出一組坐標對應的重建張量元素,即它編碼了公式 (2)。事實上,它的輸入函數f?將變量狀態映射為嵌入向量,也就是從 Tucker 分解中得到的矩陣所包含的實數值,見圖1。
請注意,我們可以很容易地對我們的構造進行特化,從而得到對應于其他分解形式(如 CP、RESCAL 和 HOSVD)的電路。
作為我們構造方法的一個具體示例,請考慮以下情況。設T ∈ ?3?3?3是一個三維張量,其定義如下:
然后,我們可以構建一個結構與圖2相同的電路c,為其輸入單元配備來自V?1?、V?2? 或 V?3?的嵌入向量,具體取決于它們的作用域,并將加法單元的參數設置為通過對張量W進行向量化得到的向量w ∈ ??,其值為(0.5, ..., 0.5)。
現在,為了計算張量T中條目t???的近似值,我們可以以前饋方式對電路c進行求值——即先計算輸入,再計算輸出——從而得到c(1, 2, 2)的值。該求值過程將執行如下計算:
請注意,括號內顏色編碼的塊對應于電路中輸入函數的輸出(見圖2),而向量外積(?)實現了電路中的乘法單元,最后與w的點積則由加法單元進行編碼。
我們邀請讀者嘗試這個例子,并試著恢復張量中的其他條目,直到你對如何將張量分解轉換為電路形式有更清晰的理解。
此外,由于電路可以表示張量分解,它們也繼承了張量分解方法中常見的“非唯一性”問題(例如 Tucker 分解)。也就是說,由電路所編碼的張量分解不是唯一的:在保持函數不變的前提下,我們可以改變電路的參數。
最后,我們要指出的是,分解的多線性秩現在轉化為電路表示中輸入單元的數量。稍后,在將層次化分解轉化為深層電路(第2.2節)時,秩也將體現為不同深度層級上單元的數量。
將張量分解表示為這種類型的計算圖,為我們擴展原有模型類別提供了許多機會。這些機會將在本文中以方框標注的形式加以強調。
同時,我們也可以更好地理解為什么這些分解已經能夠支持某些感興趣量的可處理計算,例如積分、信息論度量或最大化運算(Vergari 等人, 2021)。這些計算可以在電路框架下系統地完成,該框架將這些計算能力映射到計算圖的某些結構性質的存在上,并精確地定義了可處理性的充分條件(有時也是必要條件)。
我們首先定義兩個結構性質:平滑性(smoothness)和可分解性(decomposability),它們使得我們可以高效地計算關于指數多個變量賦值的求和運算,而這類運算在其他模型中通常是難以處理的。
定義3(逐單元平滑性與可分解性)(Darwiche & Marquis, 2002)
一個電路是平滑的(smooth),如果對于每一個加法單元n,它的所有輸入單元都依賴于相同的變量集合,即:
對任意 i, j ∈ inp(n) ,都有 scp(i) = scp(j) 。
一個電路是可分解的(decomposable),如果對于每一個乘法單元n,它的所有輸入單元分別依賴于互不相交的變量集合,即:
對任意 i ≠ j ∈ inp(n) ,都有 scp(i) ∩ scp(j) = ? 。
對于一個既平滑又可分解的電路,我們可以精確地計算如下形式的求和:
其中Z ? X,Y = X \ Z,這被稱為邊緣分布(marginals),只需一次前饋遍歷其計算圖即可完成計算(Choi 等人, 2020)。
請參見我們在第3節中對平滑性與可分解性更多應用場景的討論。
很容易驗證,將 Tucker 張量分解表示為電路(如圖2所示)時,它既是平滑的又是可分解的,因此繼承了可處理的邊緣化能力。
此外,從這一視角出發,我們還可以理解這些分解的表達能力。因為多線性多項式的表達能力通常通過具有這些結構性質的電路來刻畫(Shpilka & Yehudayoff, 2010;Martens & Medabalimi, 2014;de Colnet & Mengel, 2021)。
張量分解與電路表示究竟從何而來?
在我們已經建立了張量分解與電路之間的初步聯系之后(即前者可以被重寫為具有特定結構性質的計算圖,用后者語言來表達),我們也指出這兩個研究社區在如何獲得和處理這些對象方面的一個首要區別。
張量分解起源于對壓縮給定高維張量的需求,而這個張量通常是顯式表示的(即使不在內存中,也在磁盤上)。接著,通過一個優化問題來獲得其分解形式,例如:找到使某種重建損失最小化的因子(Sidiropoulos 等人, 2017;Cichocki 等人, 2007)。
相比之下,現代的電路是從數據中學習得到的。雖然這種學習既可以在有監督的方式下進行,也可以在無監督的方式下進行,但后者更為常見,因為電路通常被用來編碼一個概率分布。我們可以將這樣的分布看作是一個從未被觀測到的“隱式張量”,我們僅能從中采樣出數據點。
第3節將對此進行形式化,并介紹電路學習問題。
盡管張量重建通常與從數據中學習電路的方法不同,但一旦獲得了某個張量分解,將其視為一個電路后,我們就打開了許多新的使用和挖掘它的機會。我們在后續章節中以方框標注的形式突出這些機會。
接下來,我們將討論電路框架如何也推廣了層次化(或更深的)張量分解,這也將為我們提出的學習電路與張量分解的統一流程提供切入點(見第4節)。
2.2 層次化張量分解即深層電路
張量分解可以通過堆疊組合形成一個深度或層次化的分解結構,這種結構在空間效率上可以遠優于其淺層形式(即具有更低的秩)。例如,Grasedyck (2010)提出了層次化 Tucker 分解(hierarchical Tucker),它根據對張量維度的一個固定層次劃分方式,將多個低秩的 Tucker 分解進行堆疊。
Cohen 等人 (2015)指出,在大多數情況下,若使用等效甚至近似的淺層分解,所需的秩將隨維度數量呈指數增長。對于電路而言,也有類似的理論結果:深層電路可以比淺層電路小指數級,其中電路的大小指的是單元之間的連接數量(Delalleau & Bengio, 2011;Martens & Medabalimi, 2014;Jaini 等人, 2018b)。
在本節中,我們首先介紹層次化 Tucker 分解,并展示它可以被視為一種深層電路。隨后,我們將利用這一聯系來描述現代張量化的電路表示(見第4節)。為此,我們借鑒了電路領域中的一個工具:電路作用域的層次劃分(Vergari 等人, 2021),也稱為區域圖(Region Graph, RG)(Dennis & Ventura, 2012)。
正如我們接下來將形式化定義的那樣,區域圖是一個二分圖,其節點要么是變量集合(即張量的維度),要么表示這些變量集合是如何被劃分的。
定義4(區域圖)(Dennis & Ventura, 2012) 給定一個變量集合X,一個區域圖 R是一個二分的、有根的有向無環圖(DAG),其節點要么是區域節點,表示X的子集,要么是劃分節點,表示某個區域是如何被劃分為其他區域的。該圖的根節點是區域節點X
不失一般性,我們假設使用的是二元區域圖(binary RG),即每個區域都被劃分為兩個子區域,如圖3所示。
與我們對電路的圖形表示法(定義2)類似,在圖中我們省略了節點連接的方向性,并假設邊是從包含更多變量的區域節點指向包含更少變量的區域節點。
接下來,我們定義 Tucker 分解的一個層次化變體。
附錄 A.2 展示了該構造過程,并在圖4a中以基于圖3所示區域圖的層次化 Tucker 分解為例進行了說明。
以完全相同的方式,我們可以將任何張量分解擴展為層次化的形式,并將其表示為一個電路。
然而,在電路相關的文獻中,我們發現許多架構并不局限于樹狀區域圖(tree RGs),也不局限于具有單變量輸入區域的結構。
(Def. 4 允許使用任意的 DAG 和任意范圍的區域,而分層張量因子化通常以具有樹結構的區域圖 (RG) 表示,其中每個葉子區域包含恰好一個變量(例如,圖 3 中的 RG),這些有時被稱為“維度樹”或“模式簇樹”(Grasedyck, 2010),在張量因子化社區中,或者在電路文獻中稱為“vtree”(Pipatsrisawat & Darwiche, 2008; Kisa et al., 2014; Wedenig et al., 2024a)。更復雜的 RG 結構可以增加表達能力和靈活性,從而在構建深度電路/分層因子化時更具優勢。直觀上,我們可以在多個因子化中共享因子矩陣,從而減少模型參數的數量,使得更高效的實現成為可能。更多細節請參見 Peharz 等人(2020c;a)。我們在左側提供了這樣一個 RG 的例子,分別見圖 10 和圖 5。
讀者可以檢查圖 1 中的這個 RG 編碼了可分解電路的層次作用域劃分。圖 9 然后展示了這個 RG 的一部分,并說明了如何根據它構建作為電路的張量因子化。電路文獻提供了幾種適合某些數據模態(例如,圖像、序列、表格數據)的 RG 構建方法,這些也可以從數據中學習。第 4.1 節提供了這些技術的概述。)
通過利用一個區域圖(RG)來施加特定的分解結構,并在其中為每個區域選擇特定的參數化方式(如第4節將詳細討論),我們能夠編碼出與現有方法不同的新型層次化分解結構。圖6展示了一些示例。
在圖中,我們以“按層展開”的形式表示電路,這種表示方式將在第2.3節進一步說明。
請注意,從上述定義的 RG 出發構建張量分解,可以保持可分解性;而文獻中基于 RG 構建的電路通常也滿足平滑性(定義3)。
層次化 Tucker 及其變體也同時具有平滑性與可分解性,因此支持許多(概率)推理任務的可處理計算(見第3節)。
這些遵循“樹狀區域圖且葉節點為單變量”的層次化分解(及其對應的深層電路)還滿足一個額外的結構性質,稱為結構化可分解性(structured decomposability)。
結構化可分解性使我們能夠高效地執行一些僅靠平滑性與可分解性無法完成的更復雜的運算。例如,在張量網絡的圖形語言中,對某些特定張量分解進行平方操作——這在物理學中被稱為Born 規則(Feynman, 1987;Glasser 等人, 2019)(也見第2.4節)。
下面我們給出結構化可分解性的定義:
定義6(結構化可分解性)(Pipatsrisawat & Darwiche, 2008)
一個電路是結構化可分解的,如果它滿足以下兩個條件:
它是平滑的且可分解的;
對于任意兩個具有相同作用域的乘法單元n和m,它們在其輸入單元上以相同的方式對該作用域進行分解。
我們可以很容易驗證,層次化 Tucker 分解所生成的電路是結構化可分解的:這是因為它基于一棵區域圖樹堆疊 Tucker 分解(由可分解電路計算得到),從而使得所有乘法單元都以相同方式進行分解。
我們強調,明確那些能夠解釋多種感興趣量可處理計算的少數結構性質,有助于避免重復發現和重新設計特定層次化分解算法的努力。
2.3 在張量化形式中表示電路
將(層次化)張量分解表示為(深層)電路,突顯了我們可以自然地按類型和作用域對電路單元進行分組,并組織成“層”(layer),如圖2所示。
這一視角帶來了新的機會:將某些電路結構定義并表示為張量化的計算圖。盡管文獻中電路通常是以標量級別的計算單元(加法、乘法、輸入)和單個連接來定義的(定義2),但如今許多成功的電路實現已經將單元以張量的形式進行組織(Vergari 等人, 2019a;Peharz 等人, 2020c; a;Liu & Van den Broeck, 2021b;Loconte 等人, 2024),目的是利用 GPU 提供的加速能力來提升計算效率。
基于這些思想,我們現在給出一個通用的張量化電路定義,它提供了一種模塊化方式來構建過參數化的電路架構。這使我們能夠設計一個統一的學習流程,涵蓋許多現有架構(見第4節),并提出一種通過混合與復用小型“積木塊”來創建新型架構的方法。
定義7(張量化電路)
一個張量化電路 c是由三種類型的層組成的計算圖:
- 輸入層
- 乘法層
- 加法層
每一層?都由具有相同作用域scp(?)的計算單元組成。
每一個非輸入層接收來自其他層的輸出向量作為輸入,記作集合inp(?)。
這三種層定義如下:
圖7展示了張量化電路中的層類型,并與逐單元表示(定義2)進行了對照。
此外,當我們設置每層大小K = 1時,就恢復了之前的標量單元表示。
上述四種類型的層構成了我們將用來構建更復雜層的基本“樂高積木”(參見第4.3節和第5節),并可用于再現所有現代電路架構(見表1)。
表 1:將去結構化的電路和張量因子化架構及其實現,拆解為更簡單的設計選擇,以符合我們的流水線設計:應使用哪些區域圖(第 4.1 節)和和積層(第 4.3 節),以及是否應用折疊(第 4.4 節)。通過混合搭配這些已有的基本組件,可以創造出新的設計。此外,我們提出了新的區域圖,以實現更高效的張量化電路:QG、QT-2 和 QT-4。通過利用折疊電路權重的張量因子化,我們提出了兩種新的和積層:CP、CPS 和 CPXS。打勾標記 ? 表示盡管 HCLTs 的原始實現并未如我們在此描述的那樣實現折疊,但它們通過自定義的 CUDA 內核實現了類似的并行性。在附錄 B 中,我們詳細討論了在每種概率電路(PC)架構中隱式做出的、與我們流水線相關的設計選擇。
示例說明
為了說明這種定義如何幫助我們抽象出電路架構中的細節,請參見圖4。
在圖中,加法層與 Kronecker 乘法層被用于堆疊兩個 Tucker 張量分解,從而表示一個層次化的 Tucker 分解。
我們在第4節中系統性地介紹如何堆疊不同類型的層,并以此方式構建深層電路。
現在,我們可以輕松地將定義3中關于結構性質的逐單元定義擴展到本節提出的按層表示,只需定義每一層的作用域即可。
請注意,通過假設每一層都由具有相同作用域的單元組成,并使用定義7中所定義的三種層結構,我們所構建的張量化電路在設計上就具有平滑性(smoothness)和可分解性(decomposability)。
此外,如果一個深層電路所對應的區域圖(RG)是一棵樹,那么該張量化電路也將是結構化可分解的(structured-decomposable)(定義6)。
我們可以從圖4b中層次化 Tucker 分解所對應的張量化電路的圖形表示中快速識別出這些性質。
接下來,我們將利用這種分層抽象來連接當前流行的張量網絡(tensor networks),并展示它們如何被自然地編碼為深層電路。
2.4 張量網絡作為深層電路
在物理學和量子計算等領域,張量網絡(Tensor Networks, TNs)通常是表示層次化張量分解的首選方式(Markov & Shi, 2008;Schollwoeck, 2010;Biamonte & Bergholm, 2017)。
張量網絡擁有一套圖形化語言——Penrose 圖形記號——用于以緊湊的圖示形式編碼張量點積(也稱為張量收縮)(Orús, 2013 提供了綜述)。
其中,最流行的張量網絡分解之一是矩陣乘積態(Matrix-Product State,MPS),它也被稱為張量列分解(Tensor-Train factorization,TT)(Pérez-García 等人, 2007;Oseledets, 2011;Glasser 等人, 2019;Novikov 等人, 2021)。
例如,給定一個張量T ∈ ??1?...???,其秩為 R 的 MPS/TT 分解可逐元素表示為:
在圖8中,我們展示了一個表示 MPS/TT 的張量化電路,其變量為X = {X?, X?, X?}。正如 Loconte 等人 (2024) 中命題3的證明所詳述的那樣,該電路的輸入層和密集層的參數是通過對 MPS/TT 中的張量 進行分解得到的。
與層次化 Tucker 分解的張量化電路表示(命題2)類似,命題3也導出了一個具有結構化可分解性(structured decomposability)的張量化電路(定義6)。
結構化可分解性是 MPS/TT 中的一項關鍵性質,它使得某些操作可以高效執行。例如,對 MPS/TT 進行平方運算從而構建一個Born 機器(Born machine)——這是一種用于模擬物理學中量子多體系統的概率模型(Orús, 2013;Glasser 等人, 2019)。
理解這一點后,實踐者就可以設計替代性的 Born 機器架構,這些架構不再局限于“線性”區域圖所編碼的一系列張量操作,也無需從頭開始證明這些新架構上平方運算的可處理性(Shi 等人, 2005)。
這是我們強調的、當將層次化張量分解表示為電路時所帶來的研究機會之一(見機會1與機會2)。下一節我們將介紹更多此類機會,并且這些機會同樣適用于張量網絡(TNs)以及傳統的張量分解。
下一步 到目前為止,我們討論的是實值張量的一般分解形式。然而,還有一些專門針對非負數據(如圖像)的張量分解方法,稱為非負張量分解(non-negative tensor factorizations),它們將張量分解為非負因子,因此更容易解釋(Cichocki & Phan, 2009)。
在第3節中,我們將非負張量分解與概率建模電路的文獻聯系起來,從而將其解釋為一種深層潛在變量模型(deep latent-variable models)。
此外,通過建立非負張量分解與其作為(深層)電路表示之間的橋梁,我們將展示與以下兩個方向相關的未來研究機會:
如何對張量分解進行參數化;
如何使用這些分解進行概率推理。
3 從非負分解到用于概率建模的電路
在機器學習領域,可處理的概率建模(tractable probabilistic modeling)中的電路表示受到了廣泛關注,即:對那些支持高效推理的概率分布進行建模。以這一目標構建的電路通常被稱為概率電路(Probabilistic Circuits, PCs)(Vergari 等人, 2019b;Choi 等人, 2020)。
在本節中,我們將非負張量分解與概率電路聯系起來,并展示張量分解社區在概率機器學習領域中的一系列研究機會。
首先,我們將非負(層次化)張量分解與(深層)概率電路中關于離散潛在變量的解釋聯系起來,并舉例說明一些現有的、利用這種解釋實現線性時間概率推理的算法(不僅包括前一節討論的邊緣分布,也包括采樣)。
其次,我們展示了概率電路文獻中豐富的研究成果提供了多種緊湊參數化技術,這些技術可以產生非線性張量分解。
同時,我們也借鑒非負張量分解領域的優化技巧,用于學習概率電路。
最后,我們將非負張量分解與無限維張量分解的研究聯系起來,揭示其與編碼概率密度函數的概率電路之間的關系,以及與具有無限維加法單元的概率電路之間的聯系。
我們首先描述如何將有限離散隨機變量上的概率分布表示為一個張量分解。
確保一個電路是一個概率電路(PC)的一個充分條件是:限制加法單元的參數和輸入單元的輸出均為非負值,這樣的電路被稱為單調電路(monotonic circuit)(Shpilka & Yehudayoff, 2010)3。
例如,我們前面提到的編碼非負層次化 Tucker 分解的電路就是一個單調的概率電路(PC),因為它的加法單元權重(即核心張量W的元素)以及輸入單元的輸出(即因子矩陣{V???}?=1?的元素)都被限制為非負值。
電路中的平滑性(smoothness)與可分解性(decomposability)使得我們可以高效地計算求和與積分運算(見第2.1節),這在概率建模中意味著:對于具有這些結構性質的概率電路(PC),我們可以精確地計算任意邊緣分布或條件分布(Vergari 等人, 2019b)。
然而,這些概率電路(PC)不僅僅是可處理的概率模型,它們同時也是生成模型,我們可以從中進行精確采樣。
3.1 非負張量分解作為生成模型
由于非負分解(如非負層次化 Tucker 分解)是平滑的且具有(結構化)可分解性的概率電路(PC)(見定義3和定義6),它們繼承了概率電路在進行可處理推理和生成新數據點(即其所定義變量的某些配置)方面的能力。
據我們所知,將張量分解作為生成模型來處理的研究目前尚未受到廣泛關注。在下文中,我們將對此進行討論,并展示如何為這些表示設計(更高效的)采樣算法。
3.2 如何參數化概率張量分解?
電路(circuits)與張量分解(tensor factorizations)是兩種不同優化問題的輸出結果,但它們在優化過程中也面臨一些共同的挑戰。理解這些挑戰可以為兩個研究社區帶來新的研究機會。
在(非負)張量分解的應用場景中,主要任務是對一個給定的張量進行壓縮或重建,而該張量通常在內存中是顯式表示的。因此,張量分解的參數是通過最小化某種重建損失函數來優化的(Cichocki 等人, 2007)。
相比之下,現代的概率電路(PCs)是從數據中學習得到的。也就是說,我們通常會獲得一個包含N個樣本的數據集{x???}?=1?,并假設這些樣本獨立同分布地來自某個未知分布p(X)(Bishop & Nasrabadi, 2006)。編碼該分布的概率張量因此是隱式的:一方面因為分布本身未知,另一方面由于其可能具有指數級甚至無限大的規模(見第3.4節),所以無法完整地顯式構造出來。
由于概率電路的學習通常歸結為一個優化問題——即最大化數據的(對數)似然函數(Peharz 等人, 2016),因此為了保證電路的非負性,通常采用一種或多種重參數化方法(reparameterization),即將實值參數映射為正值的加法單元權重。
這是必要的,因為單調概率電路(monotonic PC)的加法權重必須構成一個凸組合,才能生成一個合法的概率分布(如公式 (8) 所示)。
例如,我們可以使用softmax 函數對具有 K 個輸入的加法單元的 K 個參數θ ∈ ??進行“壓縮”處理:
將這種重參數化方式與用于編碼概率分布的輸入函數結合使用,可以確保所構建的概率電路滿足歸一化條件(即所有變量賦值的概率之和為1)。這是因為每個加法單元的權重之和等于1。
對于張量化電路而言,這種重參數化會在每一個加法層的參數矩陣上按行進行。
幸運的是,如果電路是平滑且可分解的(定義3),即使加法權重未被歸一化,我們仍然可以高效且精確地計算其歸一化常數(Peharz 等人, 2015)。這使我們能夠使用其他方式來參數化單調概率電路c,即使它輸出的是一個未歸一化的分布(即總概率不等于1的分布)。
事實上,我們仍可以通過歸一化操作來高效地恢復一個合法的概率分布:
其中ε是一個接近零的正閾值。每種重參數化方式都可能導致不同的損失函數景觀(loss landscape),從而在優化過程中引導出不同的解。
在我們的實驗中(第6節),我們發現這種第三種重參數化方法在學習概率電路(PCs)時最為有效。
對于單調概率電路(monotonic PC)中的輸入單元,它們需要建模合法的概率分布。常見的參數化方式可以包括簡單的概率質量函數(PMF)(或密度函數,見第3.4節),例如伯努利分布(Bernoulli)或分類分布(Categorical),甚至也可以使用其他能夠高效進行邊緣化的概率模型。
這產生了一組可能的參數化方式,遠遠超出了張量分解中通常使用的那種“從索引到矩陣條目的簡單映射”方式(見命題1與圖2)。
3.3 可靠的神經符號整合
概率電路(PCs)在可處理推理方面的一個重要應用場景是安全關鍵型系統(safety-critical applications),其中需要對神經分類器的預測結果施加硬性約束(Ahmed 等人, 2022;van Krieken 等人, 2024)。
這類約束可以通過感知組件(如分類器)提取出的符號,以邏輯公式的形式表達出來。例如,在自動駕駛場景中,“汽車必須在行人或紅燈前停車”這一安全規則(Marconato 等人, 2024b;a)可以寫成一個命題邏輯公式:
?:(P∨R)?S
其中P、R、S是布爾變量,分別表示在車載視頻流中檢測到了行人(Pedestrian)和紅燈(Red-light),以及必須執行停車(Stop)操作。
電路特別適合用于這種神經符號整合(neuro-symbolic integration)(De Raedt 等人, 2019),因為它們既可以表示概率分布,也可以表示邏輯公式。
這兩種表示可以在同一個分類器中結合使用,以確保任何違反給定約束的預測結果的概率恒為零。
形式上,我們可以實現這樣一個分類器:它將輸入x映射到輸出y,且輸出必須滿足某個約束條件?,其定義如下(Ahmed 等人, 2022):
其中:
- q(y | x)
是一個由電路編碼的條件分布,它可以由神經網絡進行參數化(見“機會4”);
- ?{y ? ?}
是一個指示函數,當預測結果y滿足約束?時取值為 1(即邏輯上成立,記作 ?)。
例如,在我們自動駕駛的例子中,y是對變量P、R、S的布爾賦值,而?{y ? ?}在將y代入公式?后結果為“真”(J)時取值為 1。
這個指示函數可以通過一種稱為知識編譯(knowledge compilation)的過程,被緊湊地表示為由加法和乘法單元構成的電路形式(Darwiche & Marquis, 2002;Chavira & Darwiche, 2008;Choi 等人, 2013)?。
如果概率分布q與約束?的指示函數所對應的電路是兼容的電路結構(見“機會2”),那么我們可以高效地將它們相乘,并通過計算配分函數(partition function)來進行歸一化。
該配分函數等于在給定輸入x的條件下,硬性約束?成立的概率,即:
這也被稱為加權模型計數(Weighted Model Count,WMC)(Chavira & Darwiche, 2008;van Krieken 等人, 2024),它是進行邏輯推理與概率推理結合時所需計算的核心量(Darwiche, 2009;Zeng 等人, 2020)。
據我們所知,這種可能的整合方式目前尚未進入張量分解研究社區的關注視野。
3.4 無限維概率張量與連續分解
到目前為止,我們討論的電路所表示的是具有有限維度的張量(包括層次化)分解,也就是說,每個維度上的元素數量是有限的。這類電路定義在一組離散變量之上,每個變量也只有有限個狀態。
在本節中,我們將關注那些維度可能具有無限個(甚至可能是不可數的)元素的張量,或稱為準張量(quasi-tensors)(Townsend & Trefethen, 2015)。
類比于(層次化)張量分解與電路之間的對稱性(見第2節),我們展示準張量也可以被表示為定義在至少一個具有無限(甚至不可數)域的變量上的電路。
此外,通過與一類配備了積分單元(integral units)的最新電路模型建立聯系,我們指出了一種新的研究機會:關于無限秩(hierarchical)張量分解的參數化問題,即那些秩不必為有限的分解形式。
我們將這些思想扎根于一個具體任務:建模概率密度函數(PDF)。
類似地,我們也可以構建這種連續張量分解的層次化版本,并將其應用于概率建模(Gala 等人, 2024b)。
如果公式 (13) 中的積分在計算上是難以處理的,可以使用求積規則(quadrature rules)對其進行近似。詳見 Gala 等人 (2024a) 的相關討論。
在接下來的第4節中,我們將提出一個通用的流程框架,可用于構建有限維與無限維的層次化概率張量分解,并將它們表示為深層張量化的概率電路(PCs)(定義7)。
在此之前,在下面的“機會”框中,我們強調:電路還可以作為某些概率分布的替代表示形式,而這些分布并不對應于概率張量分解。
4 如何構建和擴展電路:一種張量化的視角
我們現在已具備所有必要的背景知識,可以開始探索(分層)張量分解與(深度)電路之間的聯系。特別是,在本節中,我們將展示如何通過利用張量分解作為模塊化抽象工具,理解和統一許多——表面上看似不同——構建電路(以及其他分解方式)的方法。通過這種方式,我們可以“解耦”構建和有效學習過參數化電路的關鍵要素,即:構建具有大量參數的電路(見表1)。
圖9總結了我們的流水線結構:
i)首先,構建一個RG結構,以強制實現所需的結構性質(第4.1節);
ii)然后,通過引入基本單元并將它們分組為層來填充該模板(第4.2節),這一過程遵循多種可能的張量分解抽象方式(第4.3節);
可選地,iii)這些層可以被“折疊”,即將它們堆疊在一起以利用GPU的并行計算能力(第4.4節)。
最后,可以通過梯度下降或期望最大化(EM)算法來優化電路的參數(Peharz et al., 2016;Zhao et al., 2016)。
4.1 構建與學習區域圖(Region Graphs)
我們流水線的第一步是構建一個 RG(定義4)。它根據輸入變量的層次劃分,用于構建深度電路架構。具體來說,基于滿足某些關鍵結構性質(如光滑性和平分解性)的 RG 所構建的概率電路(PC),可以保證許多查詢任務的高效推理(第2節)。如果 RG 是一棵樹,并且葉子節點是單變量的,則還滿足結構分解性(見第2.2節)。雖然一些論文中明確使用 RG 來構建 PC(Peharz 等,2020c;a),但正如我們將展示的那樣,RG 的結構也隱含地存在于許多其他 PC 和張量分解架構中。此外,我們還介紹了一種新的方法,可以快速為圖像數據構建 RG,這種方法不依賴于具體數據集,但利用了像素的結構信息。
線性樹形 RG(LT)
構建 RG 的一種簡單方式是一次只分解一個變量。也就是說,給定變量 X 上的一個排序 π,每個第 i 個分區節點將其作用域 {Xπ(1), ..., Xπ(i)} 分解為兩個區域:{Xπ(1), ..., Xπ(i?1)} 和 {Xπ(i)}。我們稱這種 RG 為線性樹(Linear Tree, LT),并在圖3中展示了三個變量情況下的示例。變量排序可以是字典序,也可以根據額外信息(例如在建模序列數據時的時間順序)來決定。這種順序型 RG 被鏈狀張量網絡分解方法采用,如 MPS、TT 或 BM(Pérez-García 等,2007;Oseledets,2011),以及以 PC 形式表示的隱馬爾可夫模型(HMM)(Rabiner & Juang,1986;Liu 等,2023a)。
隨機樹形 RG(RND)
一種稍微復雜一點的 RG 構建方法是構造一棵平衡樹。這可以通過遞歸隨機劃分變量的方式實現,而無需依賴特定數據集。即,根區域 X 被遞歸地劃分為大致相等的子集,直到無法再進一步劃分為止。我們將這種方法稱為 RND,最初用于構建隨機化和張量化化的和積網絡(RAT-SPN)(Peharz 等,2020c)。Di Mauro 等人(2017;2021)也提出了類似的方法,不同之處在于他們在參數化電路時還會考慮部分隨機選擇的數據子集,從而將 RG 構建與電路參數化耦合在一起。
Poon-Domingos 構造法(PD)
我們可以設計專門針對特定數據模態的 RG 構造方法,同時仍保持對數據集的無關性。對于圖像數據而言,當變量對應像素值時,Poon & Domingos(2011)提出通過遞歸進行水平和垂直切割的方式,形成一個由圖像塊組成的深層結構。然而,這種方法(標記為 PD)的主要缺點是通常會生成非常深的電路結構,難以優化(第6節),因為它考慮了所有可能的遞歸劃分方式,導致圖像塊數量隨圖像大小快速增長。PD RG 在電路文獻中被廣泛使用,例如 EiNets 架構(Peharz 等,2020a)。
面向圖像數據的新 RG:四分圖(QG)與四叉樹(QT)
我們希望設計出既不依賴具體數據集又能像 PD 一樣感知像素結構的 RG,同時避免其優化問題。因此,我們提出了一種更簡單的面向圖像的 RG 構建方法,能夠生成更小的電路,并在性能上優于從數據中學得的 RG(見第6節)。附錄中的算法 D.1 詳細描述了我們的構造方法。與 PD 類似,它通過遞歸分割大小相近的圖像塊來構建 RG,但與 PD 不同的是,它每次僅將圖像塊劃分為四個部分(一次水平切分和一次垂直切分),并共享新生成的圖像塊。我們稱這種 RG 為四分圖(Quad-Graph, QG)。圖10 展示了一個 3×3 圖像的 QG 示例。
另外,我們也可以通過在水平和垂直方向上分割圖像塊(但不共享圖像塊)來構建樹狀 RG,我們稱之為四叉樹(Quad-Tree, QT)。由于這些 RG 中的區域對應圖像塊,我們可以選擇不同的劃分方式。特別地,我們用 QT-2 表示將區域劃分為兩部分(圖像塊的上下部分)的 QT,用 QT-4 表示將區域劃分為四部分(按象限劃分)的 QT。使用 QT-2 可以還原先前工作中用于圖像數據的張量分解方法(Cheng 等,2019)。
從數據中學習 RG
到目前為止討論的方法都不依賴訓練數據。為了在 RG 構建過程中利用數據信息,可以在區域節點 Y ? X 內部測試特征子集的統計獨立性。這是早期 LearnSPN 算法所采用的方法(Gens & Domingos,2013),后來被許多研究擴展(Molina 等,2018;Di Mauro 等,2019)。盡管這些變體從未明確提到 RG,但實際上它們通過執行這些統計檢驗并引入與聚類得到的不同“數據塊”相關聯的區域,隱式地構建了一個 RG(Vergari 等,2015)。另一種方法是根據數據上的啟發式規則進行區域劃分,使得區域節點可以共享(Jaini 等,2018a)。這一思想也是 Chow-Liu 算法的基礎,用于學習最佳逼近數據似然性的樹狀概率圖模型(PGM)(Chow & Liu,1968b)。Chow-Liu 算法(CL)同樣可以隱式地構建 RG,這在許多結構學習方法中都有應用(Vergari 等,2015;Rahman 等,2014;Choi 等,2011)。最近的一種方法在此基礎上表現優異:首先學習 Chow-Liu 樹,然后將其視為潛在樹模型(Choi 等,2011),最后編譯成 PC(Liu & Van den Broeck,2021b)。
一旦我們將 RG 的作用與其他部分解耦,這種隱藏的 Chow-Liu 樹(HCLT)的構建過程恰好遵循我們流水線中的步驟。
到目前為止提到的其他 PC 和張量分解架構(例如 RAT-SPN、EiNets、MPSs、BMs 等)的構建方式也遵循相同的模式,并且可以很容易地歸類到我們的流水線中(見表1)。它們不僅在所使用的 RG 類型上有所不同,還在所選擇的和層(sum layers)與積層(product layers)的類型上存在差異。在下一節中,我們將提供一個通用算法,該算法可以根據給定的 RG 以及選定的、用于編碼張量分解的和層與積層,構建一個張量化的電路架構。
4.2 Overparameterize & Tensorize Circuits
4.2 過參數化與張量化電路
給定一個區域圖(RG),構建電路最簡單的方法是為每個葉區域關聯一個輸入分布單元,為每個內部區域關聯一個求和單元,并為每個劃分關聯一個乘積單元,然后按照 RG 的結構將它們連接起來。這種方法可以生成一個平滑且(結構可分解)的電路,并且它是稀疏連接的。事實上,前一節中討論的許多結構學習算法都是隱式地采用這種策略來構建電路的(Gens & Domingos, 2013;Vergari 等,2015;Molina 等,2018)。
我們可以將這一策略適配到“深度學習范式”中,轉而輸出一個局部密集連接的過參數化電路。所謂過參數化,是指在一個 RG 中不是只放置一個,而是放置多個具有相同作用域的求和、乘積和輸入單元的過程。這樣得到的張量化的計算圖(定義7)包含更多可學習參數,并適合在 GPU 上進行并行化處理,因為我們可以將具有相同作用域的計算單元向量化,形成密集層。
算法1詳細描述了過參數化和張量化的過程。該算法的輸入包括:一個區域圖 R、輸入函數的類型 F(例如高斯分布),以及控制電路表達能力的求和單元數量 K(等價于因子分解的秩)。此外,我們還可以自定義輸入層的選擇方式,以及如何堆疊求和層和乘積層,從而以不同的效率和表達能力來構建電路。
構建輸入層
算法1的第一步是為葉區域關聯輸入單元,即那些不再被進一步分解的區域。葉區域通常是單變量的,即形式為 Y = {Xj},其中 Xj ∈ X 是某個變量。對于每個作用于變量 Xj 的葉區域,我們引入 K 個輸入單元,每個都計算一個函數 fi: dom(Xj) → ?。為了保證單調概率電路(monotonic PCs)的輸出非負性,fi 通常被選為非負函數,例如概率質量或密度函數(Choi 等,2020)。然而,我們也可以從更廣泛的表達能力強的函數族中選擇 fi,例如多項式樣條(de Boor, 1971;Loconte 等,2024)、神經網絡(Shao 等,2020;Correia 等,2023;Gala 等,2024a;b)以及歸一化流(Sidheekh 等,2023)。參見機會4。
接著,可以通過張量化輸入單元來實現高效計算:用一個輸入層 ?: dom(Xj) → ?^K 來替代這些輸入單元,使得 ?(Xj)i = fi(Xj),其中 i ∈ [K] 可以并行計算(對應算法1第11行)。接下來,根據給定 RG 中對變量的劃分方式,構建并連接相應的求和層和乘積層。
4.3 將求和層與乘積層抽象為模塊
除了輸入層之外,我們在定義7中還介紹了張量化電路中的其他基本“樂高積木”:求和層、哈達瑪積層(Hadamard product layer)和克羅內克積層(Kronecker product layer)。在接下來的內容中,我們將使用這些基本模塊來構建復合層(composite layers),這些復合層將作為更高層次的抽象模塊,無縫地插入到算法1中。
這些復合層包括:Tucker 層(圖11)、CP 層(圖15)和CPJ 層(圖16)。每種層都編碼了一種局部因子分解方式,并以不同的方式堆疊和連接內部的求和單元與乘積單元,從而提升模型的表達能力或計算效率。
請注意,基于我們對張量化層的語義定義,通過在區域圖(RG)上應用算法1來堆疊這些復合抽象模塊,最終輸出的電路始終是一個張量化的電路,并且該電路是平滑的且(結構-)可分解的(定義8)。
我們首先考慮采用Tucker 分解中計算單元連接方式的復合層,如圖2所示。這種結構最早出現在諸如RAT-SPNs(Peharz 等,2020c)和EiNets(Peharz 等,2020a)這類架構中。在這些模型中,一個作用于變量集合 Y?X 上的區域節點,若其被劃分為 (Z?, Z?),則會被參數化為一個層 ?,該層是由一個克羅內克積層(Kronecker product layer)后接一個求和層(sum layer)所組成的復合層,即它計算:
4.4 折疊以進一步加速學習與推理
我們所提出流程的最后一步(也是可選步驟)(見圖9)是將具有相同函數形式的層堆疊在一起,以提升 GPU 的并行計算能力。我們將這一步驟稱為“折疊(folding)”。
請注意,折疊僅僅是一種對電路的語法變換,也就是說,它不會改變電路所表示的函數,因此也保持了其表達能力。然而,這種簡單的語法“重寫”卻可能對學習和推理性能產生顯著影響。
事實上,“折疊”是 EiNets(Peharz 等,2020a)相較于未折疊的電路架構(如 RAT-SPNs(Peharz 等,2020c))在速度上取得額外提升的核心機制。而 EiNets 和 RAT-SPNs 在其他架構細節(例如使用 Tucker 層)方面是相同的(參見表1)。
因此,當將 RAT-SPNs 和 EiNets 視為兩類不同的概率電路模型時,通常報告的性能差異(例如 Liu 等(2023a)中的結果),實際上應歸因于其他因素,例如區域圖(RG)的選擇,或用于訓練這些模型的其他超參數之間的差異(例如優化器的選擇)。
通過在我們的流程中將這些因素分離出來,我們可以設計出更清晰的實驗,真正揭示哪些因素對性能提升起到了關鍵作用(見第6節)。
折疊層 為了獲得 Tucker 層(公式(Tucker-layer))的折疊表示,我們需要沿著一個新引入的維度(我們稱之為“折疊維度”)來堆疊參數矩陣,并根據這個新增的維度進行相應的乘法運算。
其中,??(相應地,??)表示一個折疊層,用于計算 ???? 的 F 個左側輸入(相應地,右側輸入),每個輸入分別定義在變量 Z????(相應地,Z????)上;而每個 Wn::∈RK×K2 是第 n 個 Tucker 層 ???? 的參數矩陣。換句話說,Wn:: 是張量 W∈RF×K×K2 在第一個維度上的第 n 個切片(slice),該張量是通過將每個 Tucker 層的參數矩陣堆疊在一起得到的。
由于同一個區域節點可能參與其他區域節點的多個劃分方式中(例如見圖9i),我們可能會有折疊后的輸入 ?? 和 ?? 計算相同輸出的情況。我們在圖9iii中展示了這樣一個示例:兩個共享一個輸入的 Tucker 求和-乘積層被折疊在一起。在附錄F中,我們提供了一個使用einsum 操作實現折疊 Tucker 層的 PyTorch 代碼片段。
因此,雖然折疊在評估張量化電路時可以帶來顯著的速度提升,但它也可能帶來內存使用的增加,這取決于所選擇的區域圖(RG)結構。
如何選擇要折疊的層?
接下來的問題是:我們應該如何選擇哪些層進行折疊?
最簡單的方法是從上到下(即從輸出到輸入方向)遍歷整個張量化電路,并對計算圖中處于同一深度的層進行折疊。但請注意,我們也可以對不同深度的層進行折疊。
例如,如果所有輸入層都為所有變量編碼相同的輸入函數形式,則它們可以被一起折疊。這正是 EiNets(Peharz 等,2020a)中采用的方式,也是我們在所有實驗和基準測試中所使用的方法(見第6節)。
然而,請注意,這并不一定是最優的折疊方式,針對特定架構定制不同的折疊策略可能會帶來額外的速度提升和內存節省。
盡管我們沒有探索除上述方法之外的其他折疊方式,但在我們提出的流程中,將“折疊”與“過參數化”步驟進行了明確區分(見第4.2節),這一設計將有助于未來的研究借鑒大量關于并行化通用計算圖的文獻(Shah 等,2023),從而進一步優化折疊策略。
5 通過張量分解壓縮電路與參數共享
在本節中,我們將再次借助張量分解(tensor factorizations)領域的研究成果,以改進電路架構的設計與學習過程。
我們首先觀察到,在我們提出的流程中,電路各層的參數是以大型張量的形式存儲的(例如見公式 (Tucker-layer) 和 (Tucker-folded)),因此原則上可以對這些張量進行進一步的分解。而由于這些分解本身也可以被表示為電路(命題1),最終我們可以得到多種電路架構和層的變體:其中一些是新的設計,能夠在速度與精度之間取得良好的平衡(第5.2節);而另一些則已經被隱式地應用于現有電路與張量分解結構的構建中(見表1)。
我們仍然從Tucker 層出發,目標是使用它們來壓縮一個深層電路,即用更少的參數對其進行近似。
5.1 壓縮 Tucker 層
5.2 通過張量分解實現參數共享
我們現在將注意力轉向在張量化概率電路(PC)中跨層共享參數的問題。我們再次利用張量分解技術來完成這一任務。
考慮一個按照我們提出的流程(第4.2節)基于區域圖(RG)構建的張量化 PC。可以合理地假設:位于相同深度的層,其參數張量中可能存儲著相似的結構。
例如,兩個作用域分別為相同大小的相鄰像素塊的層,可能會對其各自的輸入執行相似的變換,因為我們通常可以假設這兩個像素塊所對應的分布是高度相似的。
如果 RG 是一棵完美平衡的二叉樹,那么對整個電路進行折疊操作,實際上就是在對同一深度的層進行折疊,而這些層在參數空間中很可能具有相似的結構。
這啟發我們將參數共享機制實現為一種跨折疊層的因子分解。
具體來說,我們首先從壓縮一個折疊后的 Tucker 層(見公式 (Tucker-folded))開始;為了獲得一個能夠實現上述參數共享的新層,我們再次對該層的參數張量應用CP 分解(定義10)。
這一次,我們需要分解的是張量 W∈RF×K×K×K,它是一個四維張量,通過對折疊后的 Tucker 層 ? 的參數張量進行重塑得到,其中 F 表示折疊維度。
如果我們應用一個秩為 R 的 CP 分解,并且令 R 遠小于 K(即 R?K),則我們可以得到:
6 實驗評估:應該選擇哪種區域圖(RG)和層?
將現代概率電路(PC)架構(以及張量分解方法)解構為我們的統一流程(圖9),使我們能夠通過簡單的“混合搭配”方式構建新的張量化架構(見表1)。同時,這也有助于我們從表達能力、推理速度和優化難易度等角度,理解不同模型類別之間真正關鍵的差異。
現在我們可以輕松地分離出幾個關鍵因素,例如 RG 的作用,以及在現代電路架構中復合層的選擇,并明確指出哪一部分對性能提升起到了關鍵作用。
例如,HCLTs 曾在最近的基準測試中被認為是表現最好的電路架構之一(Liu 等,2022;2023a),但直到現在,人們還不清楚它為何能優于 RAT-SPNs 和 EiNets 等其他架構。
在我們的框架下,我們可以嘗試回答更具體的問題來解釋這一現象:
是它們所使用的、從數據中學習得到的區域圖(RG)(第4.1節)帶來的優勢嗎?
是它們采用的復合求和-乘積層參數化方式(第5.1節)嗎?
還是其他超參數的選擇起了作用?
(劇透:最終答案將是使用了CP 層。)
具體來說,在本節中我們將通過嚴格的實證研究來回答以下三個研究問題:
- RQ1
:我們目前可以構建的多種張量化架構,在測試和訓練階段所需的計算資源(時間與GPU內存)分別是多少?
- RQ2
:在將張量化電路作為分布估計器進行訓練時,區域圖(RG)與復合求和-乘積層的選擇對性能有何影響?
- RQ3
:如果我們像圖14a→圖14b所示那樣,將預訓練的 Tucker 層因子分解為 CP 層后,是否仍能保留(大部分)原始性能?
請注意,我們不關心折疊(folding)的影響(第4.4節),因為我們已經知道答案:對于大規模的張量化架構來說,折疊是至關重要的。因此,在所有實驗中我們都使用了折疊后的張量化電路。
我們要強調的是,實驗的目的不是為了達到最先進的分布估計效果,而是為了理解張量化電路架構中各個組件的作用。
所有實驗均在一個配備單塊 NVIDIA RTX A6000 GPU(48GB 顯存)的設備上運行。我們的代碼可在github.com/april-tools/uni-circ-le獲取。
一種新的電路命名方式
需要說明的是,表1中的 HCLT、EiNets、RAT-SPNs 等縮寫并不代表不同的模型類別,而只是不同的架構設計。它們都屬于同一類模型:平滑且(結構)可分解的概率電路。
在下文中,我們將使用如下格式表示一個張量化架構:[RG]-[sum-product layer]
,必要時加上 K 表示如算法1中所述用于過參數化的單元數量。
當 RAT-SPNs 和 EiNets 使用隨機區域圖(random RG)構建時,它們將被統一命名為
RND-Tucker
;若使用 Poon&Domingos 區域圖,則稱為
PD-Tucker
;而 HCLTs 則對應
CL-CP
架構。
我們主要通過對圖像數據集進行分布估計來評估所提出的架構。
我們使用的Mnist家族數據集包括6個灰度圖像數據集(28×28):
Mnist(LeCun 等,2010)
FashionMnist(Xiao 等,2017)
EMNIST 及其4個劃分(Cohen 等,2017)
此外,我們還使用了 CelebA 數據集,將其縮放為 64×64 像素大小,并提供兩個版本:
RGB 像素版本
經過無損 YCoCg 顏色編碼預處理的像素版本(Malvar & Sullivan, 2003),因為近期研究表明這種變換可以顯著降低 bpd(bits per dimension)
除此之外,我們還在具有連續變量的表格數據上進行了實驗。特別地,我們在5個 UCI 數據集上執行密度估計任務,這些數據集通常用于評估歸一化流模型(Papamakarios 等,2017)。UCI 數據集的統計信息詳見附錄表 E.5。
參數優化
我們訓練電路以估計生成圖像的概率分布,將每個像素視為一個隨機變量。因此,電路中的輸入單元代表具有256個取值的分類分布(Categorical distributions)。對于 RGB 圖像,每個像素關聯三個分類分布單元(每個顏色通道一個)。
而對于5個 UCI 數據集(見表 E.5),我們則使用代表單變量高斯分布的輸入單元,并同時學習其均值和標準差。
我們通過隨機梯度上升法進行最大似然估計,即最大化如下目標函數:
其中,是概率電路 c 的配分函數,而 B 表示一個訓練批次的數據。
在進行了一些初步實驗后,我們發現使用 Adam 優化器(Kingma & Ba, 2015),并將學習率設為 ,平均而言可以為我們所考慮的數據集帶來表現最好的模型。
此外,我們還確定在每次優化步驟之后,通過對電路中的求和參數進行clamp 操作并設(見公式(9)),以保持其非負性,這種方式在所有可能的重參數化方法中帶來了最佳的學習動態(參見第3.2節)。
在下文中,我們將總結針對 RQ1–RQ3 所得的結論,并提煉出一些面向實踐者的建議,用于指導如何構建和訓練概率電路。
RQ1)不同張量化架構的時間...
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.