谷歌宣布開源張量計(jì)算庫TensorNetwork及其API

谷歌宣布開源張量計(jì)算庫TensorNetwork及其API，使用TensorFlow為后端，對(duì)GPU處理速度進(jìn)行優(yōu)化，與CPU相比，計(jì)算加速效果高達(dá)100倍。

現(xiàn)代科學(xué)領(lǐng)域中，有很多艱巨困難的科學(xué)任務(wù)，比如開發(fā)高溫超導(dǎo)體材料、了解空間和時(shí)間的本質(zhì)等，都涉及到處理量子系統(tǒng)的復(fù)雜性。這些問題之所以困難，是因?yàn)檫@些系統(tǒng)中的量子態(tài)數(shù)量呈指數(shù)級(jí)增長，使得暴力計(jì)算行不通了。

為了解決這個(gè)問題，人們利用名為“張量網(wǎng)絡(luò)”的數(shù)據(jù)結(jié)構(gòu)，可以專注于與現(xiàn)實(shí)問題最為相關(guān)的量子態(tài)——低能量狀態(tài)，而忽略其他不相關(guān)的狀態(tài)。張量網(wǎng)絡(luò)也越來越多地在機(jī)器學(xué)習(xí)中得到應(yīng)用。

然而，目前在機(jī)器學(xué)習(xí)中應(yīng)用張量還存在一些困難：比如用于加速硬件的生產(chǎn)級(jí)張量網(wǎng)絡(luò)庫尚未在大規(guī)模運(yùn)行張量網(wǎng)絡(luò)算法中部署，而且，大多數(shù)關(guān)于張量網(wǎng)絡(luò)的文獻(xiàn)是面向物理學(xué)科領(lǐng)域的應(yīng)用。這也讓人們產(chǎn)生一種錯(cuò)誤印象，認(rèn)為需要掌握量子力學(xué)的專業(yè)知識(shí)才能理解張量算法。

本次開源的TensorNetwork使用TensorFlow作為后端，并針對(duì)GPU處理進(jìn)行了優(yōu)化，與CPU相比，處理速度可以實(shí)現(xiàn)100倍的加速。此前已經(jīng)介紹了TensorNetwork，包括新的庫及其API，并針對(duì)非物理學(xué)背景的讀者對(duì)張量網(wǎng)絡(luò)進(jìn)行了概述，介紹了張量網(wǎng)絡(luò)在物理學(xué)中的特定應(yīng)用實(shí)例，展示了使用GPU帶來的處理速度的顯著提升。

為什么Tensor Networks有用？從張量的圖解表示說起

張量是一種多維數(shù)組，根據(jù)數(shù)組元素的順序按層級(jí)分類：例如，普通數(shù)是零階張量（也稱為標(biāo)量），向量可視為一階張量，矩陣可視為二階張量等等。低階張量可以很容易用一個(gè)明確的數(shù)字?jǐn)?shù)組或數(shù)學(xué)符號(hào)來表示。

不過涉及到高階時(shí)，這種符號(hào)法就變得非常麻煩。使用圖解符號(hào)對(duì)于解決這個(gè)問題很有用，一種方法是簡單地繪制一個(gè)圓（或其他形狀），引出多條線或者說“腿”，腿的數(shù)量與張量的階數(shù)相同。在這種表示法中，標(biāo)量表示為一個(gè)圓，矢量有一條腿，矩陣有兩條腿等。張量的每條腿也有一個(gè)尺寸，就是腿的長短。例如，表示物體通過空間的速度的矢量就是三維的一階張量。

張量的圖解表示法

以這種方式表示張量的好處是可以簡潔地對(duì)數(shù)學(xué)運(yùn)算進(jìn)行編碼，例如，用矩陣乘以向量，獲得另一個(gè)向量，或者將兩個(gè)向量相乘，得到一個(gè)標(biāo)量。這些都是所謂“張量收縮”的更一般的概念。

張量收縮的圖解表示法。圖中所示為矢量和矩陣乘法以及矩陣跡線（即矩陣的對(duì)角元素的總和）

以下是張量網(wǎng)絡(luò)的簡單示例。張量網(wǎng)絡(luò)是對(duì)幾個(gè)張量收縮，形成新張量的模式進(jìn)行編碼的圖形化表示。構(gòu)成新張量的每個(gè)張量具有各自的階數(shù)，圖上表示為腿的數(shù)量?；ハ噙B接的腿，在圖中形成邊，表示張量的收縮，而剩余的懸在外面的腿的數(shù)量就是生成的新張量的階數(shù)。

左：四個(gè)矩陣乘積的表示，即tr（ABCD），它是一個(gè)標(biāo)量，圖中沒有腿。右：三個(gè)三階張量收縮，最終有三條腿懸在外面，即產(chǎn)生一個(gè)新的三階張量。

雖然這些例子非常簡單，但張量網(wǎng)絡(luò)通常代表以各種方式收縮的數(shù)百個(gè)張量。用傳統(tǒng)的符號(hào)來描述這樣一件事是很難理解的，這就是Roger Penrose在1971年發(fā)明圖解符號(hào)（diagrammatic notation）的原因。

張量網(wǎng)絡(luò)在實(shí)踐中的應(yīng)用

想象一組黑白圖像，每個(gè)圖像都可以看作是一個(gè)n個(gè)像素值的列表。單個(gè)圖像的單個(gè)像素可以被一個(gè)one-hot編碼為二維矢量，通過將這些像素編碼結(jié)合在一起，我們可以對(duì)整個(gè)圖像進(jìn)行2N維的one-hot編碼。我們可以將這個(gè)高維向量重塑成一個(gè)order-N張量，然后將圖像集中的所有張量相加，得到一個(gè)總張量Ti1,i2,...,iN集合。

這聽起來是一件非常浪費(fèi)的事：用這種方式編碼大約50像素的圖像將占內(nèi)存許多PB的空間。這就該用到張量網(wǎng)絡(luò)了。與其直接存儲(chǔ)或操縱張量T，不如將T表示為張量網(wǎng)絡(luò)形狀中許多較小組分張量的收縮。結(jié)果證明效率更高。例如，流行的矩陣積態(tài)（MPS）網(wǎng)絡(luò)將把T寫成N個(gè)更小的張量，這樣參數(shù)的總數(shù)在N中只是線性的，而不是指數(shù)的。

在矩陣積態(tài)張量網(wǎng)絡(luò)中，高階張量T用許多低階張量表示。

不明顯的是，大張量網(wǎng)絡(luò)可以被有效地創(chuàng)建或操作，同時(shí)始終避免占用大量內(nèi)存。但事實(shí)證明，這在許多情況下是可能的，這就是為什么張量網(wǎng)絡(luò)在量子物理學(xué)和現(xiàn)在的機(jī)器學(xué)習(xí)中被廣泛使用的原因。

谷歌AI的研究人員Stoudenmire和Schwab使用剛才描述的編碼來建立一個(gè)圖像分類模型，展示了張量網(wǎng)絡(luò)的新用途。TensorNetwork庫的設(shè)計(jì)就是為了方便這種工作，我們第一篇論文（https://arxiv.org/pdf/1905.01330.pdf）就描述了該庫如何用于一般的張量網(wǎng)絡(luò)操作。

性能實(shí)例分析：計(jì)算速度提升100倍

張量網(wǎng)絡(luò)是張量網(wǎng)絡(luò)算法的通用庫，對(duì)物理學(xué)家也有一定的幫助。量子態(tài)的近似是物理中張量網(wǎng)絡(luò)的一個(gè)典型用例，非常適合用來說明張量網(wǎng)絡(luò)庫的功能。在第二篇論文（https://arxiv.org/pdf/1905.01331.pdf）中，我們描述了一種tree tensor network（TTN）算法，用于估算周期性量子自旋鏈（1D）或薄環(huán)面上的晶格模型（2D）的基態(tài)，并用張量網(wǎng)絡(luò)實(shí)現(xiàn)了該算法。在使用GPU和TensorNetwork庫時(shí)，我們比較了CPU和GPU的使用情況，并觀察到計(jì)算速度顯著提高，高達(dá)100倍。

計(jì)算時(shí)間作為鍵維數(shù)的函數(shù)，χ。 鍵合維度決定了張量網(wǎng)絡(luò)的組成張量的大小。 更大的鍵維度意味著更強(qiáng)大的張量網(wǎng)絡(luò)，但需要更多的計(jì)算資源來操縱。

未來方向：時(shí)間序列分析和量子電路模擬

我們計(jì)劃用一系列的論文來說明張量網(wǎng)絡(luò)在實(shí)際應(yīng)用中的強(qiáng)大之處，這是第一篇。在下一篇論文中，我們將使用TensorNetwork對(duì)MNIST和Fashion-MNIST數(shù)據(jù)集中的圖像進(jìn)行分類。

未來的計(jì)劃包括機(jī)器學(xué)習(xí)方面的時(shí)間序列分析和物理方面的量子電路模擬。通過開源社區(qū)，我們會(huì)經(jīng)常為TensorNetwork添加新功能。我們希望TensorNetwork將成為物理學(xué)家和機(jī)器學(xué)習(xí)實(shí)踐者的寶貴工具。