直接拿來用！谷歌開源網路庫TensorNetwork，GPU處理提升100倍

編譯 | 琥珀

出品 | AI科技大本營（ID:rgznai100）

世界上許多最嚴峻的科學挑戰，如開發高溫超導體和理解時空的本質，都涉及處理量子系統的複雜性。然而，這些系統中量子態的數量程指數級增長，使得暴力計算並不可行。為了解決該問題，數據結構採用了張量網路的方式。張量網路讓人們關注與現實世界問題最相關的量子態，如低能態。張量網路已經越來越多地在機器學習中得到應用。

然而，目前仍存在相當多的一些挑戰阻礙了其在機器學習社區中的廣泛使用：1）用於加速硬體的生產級張量網路庫還不能用於大規模運行張量網路演算法；2）大多數張量網路文獻是面向物理應用的，併產生了一種錯誤的印象，即需要量子力學方面的專業知識來理解演算法。

為了解決這一問題，谷歌 X 實驗室與加拿大 Perimeter 理論物理研究所（Perimeter Institute for Theoretical Physics ）的研究人員合作開發了張量網路 TensorNetwork，以 TensorFlow 作為後端，針對 GPU 處理進行了優化。與在 CPU 上計算工作相比，可以實現高達 100 倍的加速。這是一個全新的開源庫，旨在提高張量計算的效率。

據悉，研究人員已經發布了一些列論文對張量網路的概念、特性以及應用特例等方面進行了詳細闡釋。

論文鏈接：

https://arxiv.org/abs/1905.01330

https://arxiv.org/abs/1905.01331

GitHub 鏈接：

https://github.com/google/tensornetwork

工作原理

張量是一種多維數組，按照順序進行分類。例如，一個普通數零階張量，也稱為標量，一個向量是一階張量，一個矩陣是二階張量。雖然低階張量可以很容易地用數字數組或像 Tijnklm 這樣的數學符號來表示，但一旦開始討論高階張量，這個符號就變得非常麻煩。

這一點上，使用圖解記數法是非常有用的，在這種記數法中，人們只需畫一個有許多條線或」 腿 「的圓（或其他形狀）。在這個符號中，標量只是一個圓，向量只有一條腿，矩陣有兩條腿等。張量的每條腿也有大小，也就是腿的長度。

張量的圖解符號

以這種方式表示張量的好處是簡潔地編碼數學運算，例如，將一個矩陣乘以一個向量得到另一個向量，或者兩個向量相乘得到標量。這個過程被稱為張量收縮。

張量收縮的圖解表示法

向量和矩陣乘法以及矩陣跡線（即矩陣對角元素的總和）。

除了這些案例之外，還有以張量收縮模式進行編碼以形成一個新張量的圖形方式。每個組成張量都有一個由自己的腿數決定的順序。連接的腿在圖中形成一條邊，代表收縮，而剩餘懸空腿的數量決定了合成張量的順序。

左：四個矩陣乘積的表示，即 tr（ABCD），它是一個標量。右：三個三階張量收縮，三條腿懸空，產生一個新的三階張量。

雖然這些例子非常簡單，但張量網路通常代表以各種方式收縮的數百個張量。用傳統的符號來描述這樣一件事情是非常模糊的，這也是 Roger Penrose 在 1971 年發明圖解符號（diagrammatic notation）的原因。

實踐過程

以一些黑白圖像為例，每個圖像可以被看做是 N 個像素值的列表。單個圖像中的單個像素可以被獨熱編碼（ one-hot-encoding）成二維向量，並且通過這些像素編碼組合在一起，我們可以得到 2N 個維獨熱編碼的結果。我們可以將高維向量轉化為 N 階張量，然後將圖像集合中所有張量相加，得到量 Ti1，i2，...，iN 的集合。

這聽起來像是一件非常浪費時間的事情。因為以這種方式對約 50 像素的圖像進行編碼已經佔用了數千兆位元組的內存。這正是張量網路的作用所在。研究人員沒有直接存儲或操縱張量 T，而是將張量 T 作為張量網路內的多個小的張量組合。結果證明這非常有效。例如，流行的矩陣乘積態（matrix product state，MPS）網路可以將 T 表示為 N 個較小的張量，從而總參數量為 N 的線性級，而不是指數級。

高階張量 T 用矩陣乘積態張量網路中的許多低階張量來表示。

但在有效構建或操縱大型張量網路的同時又能始終避免使用大量內容的需求，是不太明顯的。但事實證明，這在許多情況下是可能的，這也是張量網路廣泛用於量子物理和機器學習領域的原因。Stoudenmire 和 Schwab 使用這種編碼來構建圖像分類模型，展示了張量網路的新用途。而 TensorNetwork 庫旨在推進這類工作，研究人員在 TensorNetwork 相關論文中介紹了該庫如何應用於張量網路的操作。

在物理場景中的應用性能

TensorNetwork 是一個針對張量網路演算法的通用庫，因此它適用於物理學場景。逼近量子態是張量網路在物理學中的一個典型用例，可說明張量網路庫的能力。在另一篇論文《TensorNetwork on TensorFlow: A Spin Chain Application Using Tree Tensor Networks》中，研究人員提出了一種近似樹張量網路（tree tensor network，TTN），並使用張量網路庫實現了該演算法。此外，研究人員還對比了 CPU 和 GPU 的情況，發現在使用 GPU 和張量網路庫時，計算速度顯著提高了近 100 倍。

計算時間作為連接維度的函數 X。連接維度（ bond dimension ）決定了張量網路中張量的大小。連接維度越大意味著張量網路約強大，但這個過程也需要更多的計算資源。

總結及未來工作

本文是講述 TensorNetwork 實際應用案例的第一篇文章，後續的論文中，研究人員將使用 TensorNetwork 在 MNIST 和 Fashion-MNIST 上執行圖像分類，後面還將包括時序分析、量子電路模擬等。研究人員希望 TensorNetwork 能成為物理學家和機器學習從業者的寶貴工具。

相關鏈接：

https://ai.googleblog.com/2019/06/introducing-tensornetwork-open-source.html

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！