能「看到」的張量運算:因子圖可視化
- 2019 年 11 月 5 日
- 筆記
張量運算有時候並不容易直觀地理解:為什麼有時候改變計算順序不會影響結果,同時又能極大節省計算成本?使用因子圖來可視化或許能為人們提供簡潔直觀的理解方式。Rajat Vadiraj Dwaraknath 近日發佈了一篇文章,介紹了他在使用因子圖可視化張量運算方面的心得。
從格羅滕迪克那裡,我學習到不要以證明過程的難度為榮:困難意味着我們尚未理解。也就是說我們要能繪製出讓證明過程顯而易見的圖景。 ——著名數學家 Pierre Deligne
當維度超過 2 或 3 時,理解涉及多維數組的運算就會變得相當困難。但是,矩陣本身的特定性質可能讓你在初次與它們相遇時深感驚訝。軌跡運算的循環性質就是其中一例。

我最近遇到個能可視化這些所謂的張量運算的好工具——因子圖(factor graphs),它能得到視覺上很明顯(如循環軌跡)的結果。儘管我最初是在圖模型和消息傳遞的語境中遇到因子圖的,但我很快就意識到它們體現了一種更通用和更簡單的概念。在這篇文章中,我將主要在高層面介紹因子圖,而不會涉及圖模型或消息傳遞等算法的具體細節。
在深入因子圖之前,我們先簡單快速地介紹一下愛因斯坦求和符號。
愛因斯坦求和符號
愛因斯坦符號存在多種形式,尤其是在物理學領域,但我們要介紹的那種非常簡單,沒有任何物理學背景也能輕鬆掌握。
在矩陣乘法的定義中,

求和符號實際上是多餘的。我們可以直接捨棄它,並推斷出索引 k 必須被求和,因為它沒有出現在左側。

為什麼要這麼做?好吧,我們來看一個有一般張量的案例(將其看作是超過 2 維的 numpy 數組即可):

然後假設張量的形狀如下:

其中交織着複雜的「和」與「積」,而不斷寫求和符號是非常煩人的。同樣,我們不需要寫這些求和符號,因為我們可以通過查看僅出現在右側的索引來暗示所要求和的索引。用愛因斯坦表示法,寫起來就簡單多了:

但相比於公式 (4),這種表示方式確實丟失了一些信息——計算求和的順序。現在,求和的順序實際上不影響最終結果(福比尼定理),但事實表明某些順序的求和過程比另一些順序更高效(後面還會提到)。另外,你可以使用 numpy.einsum 在 Python 中輕鬆嘗試這些。這篇文章更詳細地介紹了 einsum,並給出了一些很好的示例:http://ajcr.net/Basic-guide-to-einsum/
因子圖
帶有多個不同大小的張量的和-積表達式也被稱為張量網絡。除了聽起來炫酷之外,這個名字也是合理的,因為你寫的任何有效的愛因斯坦求和實際上都可映射為一個張量圖。(「有效」是指同樣索引的不同張量的維度大小必須相等。)
我們可以很輕鬆地構建出這樣的圖。我們回頭看看這個例子,了解下發生了什麼:

- 這個圖有兩種節點——因子和變量
- 我們將用方框表示因子,用圓圈表示變量
- 因子對應張量 (A,B,C)
- 變量對應索引 (i,j,k)
- 邊僅出現在方框和圓圈之間
- 邊的規則很簡單——每個因子都連接其每個索引。在上面的例子中,A_{ijk} 表示 A 連接着 i、j、k
- 邊的厚度對應於因子中軸(即數組分量的長度)的大小
- 這使得圖成為了方框和圓圈之間的二部圖(bipartite graph)
- 僅出現在等式右側的索引(i 和 j)會被隱式地求和。我們通過加黑圖中對應的變量節點來表示它。
上面動畫的最後一部分給出了一個重要的直覺觀察:
每個因子圖都有一個完全收縮的狀態——愛因斯坦求和的右側(示例中的 2 維張量 D)。通過將圖中的所有因子組合成單個因子,然後將每個灰色的變量求和,可以得到這個狀態。現在剩下的是僅連接到未求和變量的單個因子——這就是收縮狀態。示例中轉變為因子 D 的綠色雲很好地展示了這種收縮。
在我們繼續探索這個奇特工具的能力之前,我們先談談它的來源。
名字從何而來?
這種圖被稱為因子圖的一大原因是右側看起來像是對左側張量的因子分解。在離散隨機變量的概率分佈語境中,這會更加具體。如果張量為正且總和為 1,則它們可以表示在不同隨機變量上的聯合分佈(這也是索引對應於變量的原因)。在這種設置中,因子圖是將許多變量的大型聯合分解成更小的互相獨立的變量集的聯合。
因子圖不僅會在概率圖模型中出現,而且也會現身於物理學當中,這時候它被稱為張量網絡或矩陣積狀態(matrix product states)。但是,這裡不會介紹這些應用的細節。
有一點需要注意,因子分解所需的內存實際上比整個聯合要少得多(存儲一個 10×10×10 張量對比存儲三個 10 維張量)。
可視化的 numpy 運算
為什麼這種表示方式有用?因為這能讓我們將複雜的因子分解轉換成更可視化的表示,從而更加輕鬆地處理。numpy 中的數值張量運算可以很好地適用於這個框架。下面給出了幾個無需過多解釋的示例:
矩陣-向量乘法

矩陣-矩陣乘法

逐元素求積

外積

軌跡

注意,沒有邊的因子是 0 維張量,其實就是單個數值(軌跡就該是這樣)。
可視化證明
使用簡潔的可視化證明,我們不僅可以理解 numpy 運算,還能迅速搞定數學定理。下面就簡潔地證明了我一開始提到的軌跡恆等:
軌跡是循環的

張量收縮的計算成本
現在,我們已經將因子圖中的那些綠色雲壓縮成了一個大因子,而沒有探討這種變換究竟是怎樣計算的。將許多因子組合成單個因子並求灰色變量的和的過程涉及到兩個基本的計算操作:
- 求和:移除僅有一條邊的灰色節點
- 求積:將兩個因子合併成一個因子
可以很容易看到,這樣的操作能保留網絡最終的收縮狀態,所以如果我們不斷應用它們直到只剩僅連接到未求和變量的單個因子,那麼我們就實現了網絡收縮。
求和
求和是不言自明的。基本上就是將 numpy.sum 運算應用於對應的軸。這涉及到對大小等於所有其它軸大小的積的張量求和,而且這些張量的數量就是被求和的軸的大小。因此,加法的總數量就是所有軸大小的積。我們也能從可視化表示中看出這一點:

求積
求積運算本質上就是兩個張量的外積泛化為一般張量。用愛因斯坦表示法,組合兩個因子就等同於通過兩個因子的項相乘而將兩個因子當成一個,從而得到一個更大的因子:

這種求積是用一個因子中的每個元素與另一個因子的整體相乘。因此最終結果的大小是各個因子的總大小的積,這會大很多。最終積的每個元素都只是兩個數值相乘的結果,所以乘法總數量就是最終積的項總數。這也很容易可視化:

另外,如果兩個因子共享一個變量,則兩條邊會結合成單條邊——在效果上是執行類似於軌跡動畫中的對角運算。
當收縮一個網絡時,對變量求和並以不同的順序組合因子會導致不同的計算成本。研究表明,尋找實現成本最小化的收縮一般因子圖的最優順序實際上是 NP-hard 問題。作為一個有趣的練習,你可以試試解讀矩陣鏈乘法(matrix chain multiplication)過程,並使用因子圖理解尋找一個鏈矩陣積的總計算成本是如何受乘法順序影響的。這個特定案例有非常簡潔的動態編程算法,可在平方時間內獲得最優順序。
希望我現在已經使你信服因子圖是非常強大的可視化工具。我們不僅能用它證明一些簡單的恆等關係,而且還能進一步將其用於理解一些複雜概念,比如用於概率圖模型推理的有效方法。
一些細節
因子圖其實不能精確地表示愛因斯坦求和。你們可能已經注意到我們丟失了張量的哪個軸對應於圖中哪條邊的信息。但是,只要將源自每個因子的邊加上軸標籤,就能輕鬆解決這個問題。但這會使可視化無必要地雜亂和醜陋,所以我決定不包含它們。
另外,我們還可將這種可視化擴展用於不只是求和變量節點的網絡。將一個變量節點變為灰色在效果上就是將對應的軸約簡為單個數值,因此我們可以用任何執行這種約簡的運算替代求和。舉個例子,我們不用求和,而是取該軸中所有元素的最大值,或者就簡單地索引該軸上一個特定位置。這在 MAP 估計和最大積信念傳播方面是相關的。我們甚至可以將這種解讀方式擴展到連續域,並用積分替代求和,其中的因子也不再是離散矩陣,而是連續域上的多變量函數。