テンソルネットワークをざっと確認

はじめに

量子コンピュータのアルゴリズムを今後作っていくにあたって、いろいろ覚えることが増えそうです。今回はざっくりテンソルネットワークを学んでいきたいと思います。

テンソルとは？

テンソル（英: tensor, 独: Tensor）とは、線形的な量または線形的な幾何概念を一般化したもので、基底を選べば、多次元の配列として表現できるようなものである。

テンソルネットワークとは？

大きなテンソルをより小さなテンソルに分解してネットワーク化したものです。

記法

足の数によって記法が変わります。*をノード、-をエッジとして見たときに、

・足が一つはベクトル

---*

v_j = \begin{bmatrix} v_1\\ v_2\\ \vdots\\ v_n \end{bmatrix}

・足が２つは行列

---*---

M_{ij} = \begin{bmatrix} M_{11} & \cdots & M_{1n}\\ \vdots\\ M_{m1} & \cdots & M_{mn} \end{bmatrix}

・足が３つはオーダー３のテンソル

---*---
   |

（表現の仕方がわからないので、なんとなく書いちゃいましたが、、、）

T_{ijk} = \begin{bmatrix} T_{111} & \cdots & T_{1n1}\\ \vdots\\ T_{m11} & \cdots & T_{mn1} \end{bmatrix} , \begin{bmatrix} T_{112} & \cdots & T_{1n2}\\ \vdots\\ T_{m12} & \cdots & T_{mn2} \end{bmatrix}, \cdots

計算

ノード同士はエッジで接続され、それぞれ計算されます。

・ベクトルと行列
ベクトルと行列は計算するとベクトルに、足の数は1になります。

--*--* = ----*

v'_i = \sum_jM_{ij}v_j

・行列同士
行列同士は行列になります。足の数も変わりません。

--*--*-- = --*--

M_{ik} = A_{ij}B_{jk}

参考

https://tensornetwork.org/
https://www.tensors.net/

波動関数

私たちは今後これを量子コンピュータの計算に応用したいので、波動関数・状態ベクトルなどというものと対応させてシミュレータやアプリケーションに応用をしていきたいと考えます。一般的に量子コンピュータとテンソルネットワークを直接結び付けて書いたようなものは多分少ないのと、全部を学ぶのは大変なので必要な量子コンピュータに特化したところを考えて見たいと思います。

テンソルネットワークと波動関数

量子コンピュータは通常量子ビットは|0>と|1>の重ね合わせで表現されます。

|\psi\rangle = \alpha|0\rangle + \beta|1\rangle

これをテンソルネットワークに対応させると、１ノードのテンソルに対応して、ベクトルでかけます。

-*

上記の波動関数を書き直すと、１量子ビットは、

|\psi \rangle = c_0|0\rangle + c_1|1\rangle = \sum c_i |i\rangle

また、複数量子ビットの場合には、（*同士が並んでる時はテンソル積を取ると考えて）、

-*
-*
-*
-*

|\psi \rangle = \sum_{i_0,i_1,\cdots,i_N} c_{i_0,i_1,\cdots,i_N}|i_0\rangle|i_1\rangle \cdots |i_N\rangle

これはパラメータの数が量子ビットの場合には $2^N$ になってしまうので、今後はこのパラメータを減らせるようにネットワークを構成して、パラメータ数を減らして波動関数を効率的に表現できるように目指していきます。

量子回路とテンソルネットワーク

上記波動関数が２次元のオーダーNのテンソルで表現でき、行列同士のテンソルは行列になるということなので、量子回路のようなユニタリ操作を状態ベクトルに対して適用するようなものはテンソルネットワークで表現できそうです。

量子回路の構造によって効率的なテンソルネットワークでかけて、効率的な波動関数の表現があるとすると、シミュレータやアプリケーションに応用できそうです。

特にアプリケーションに関しても、完全にランダムな形式は難しそうですが、機械学習や量子化学計算や、社会問題もハミルトニアンの形式にある程度の傾向や偏りがある場合には効率的なテンソルネットワークの表現が作れそうです。

GoogleのTensornetworkツール

テンソルネットワークの計算は実際はツール経由で行います。ここではGoogleからでたTensornetworkをnumpyバックエンドとして使って見たいと思います。

https://github.com/google/TensorNetwork

Tensornetwork

Googleから提供されているオープンソースソフトウェアです。インストールはpip経由で、

!pip3 install tensornetwork

となります。blueqat cloudで使ってます。

基本の形

基本はテンソルネットワークの形を決めてそれぞれをつなぎます。参考は下記になります。

import numpy as np
import tensornetwork as tn

a = tn.Node(np.ones((10,))) 
b = tn.Node(np.ones((10,)))
edge = a[0] ^ b[0]
final_node = tn.contract(edge)
print(final_node.tensor)

10.0

こうすると

10.0

が得られます。

少しずつ見てみる

まず読み込みます

import numpy as np
import tensornetwork as tn

次にノードを指定します。今回はAとBというベクトルを指定します。

a = tn.Node(np.ones((10,))) 
b = tn.Node(np.ones((10,)))

aとbにエッジを指定します。

edge = a[0] ^ b[0]

そして、これらのテンソル縮約をとって表示をして見ます。

final_node = tn.contract(edge)
print(final_node.tensor)

今回は、AとB共に下記のような10この要素が全て1のベクトルなので、

\begin{bmatrix} 1\\1\\\vdots\\1 \end{bmatrix}

計算結果は、

1*1 + 1*1 + ... + 1*1 = 10

と、スカラー量となります。

いろいろ

計算して見ます。

ベクトルと行列

a = tn.Node(np.ones((5))) 
b = tn.Node(np.ones((5,5)))
edge = a[0] ^ b[0]
final_node = tn.contract(edge)
print(final_node.tensor)

[5. 5. 5. 5. 5.]

これはベクトルに、

[5. 5. 5. 5. 5.]

行列同士

行列の順番とかよくわからなかったですが、、、要素数を揃えたらいけました。

a = tn.Node(np.ones((5,3))) 
b = tn.Node(np.ones((5,2)))
edge = a[0] ^ b[0]
final_node = tn.contract(edge)
print(final_node.tensor)

[[5. 5.]
 [5. 5.]
 [5. 5.]]

[[5. 5.]
 [5. 5.]
 [5. 5.]]

GPUモード

バックエンドを変えることで、GPUモードが使えます。

!pip install tensornetwork jax jaxlib
import numpy as np
import jax
import tensornetwork as tn

例題があったので、それを実行して見ます。

def calculate_abc_trace(a, b, c):
  an = tn.Node(a)
  bn = tn.Node(b)
  cn = tn.Node(c)
  an[1] ^ bn[0]
  bn[1] ^ cn[0]
  cn[1] ^ an[0]
  return (an @ bn @ cn).tensor

a = np.ones((4096, 4096))
b = np.ones((4096, 4096))
c = np.ones((4096, 4096))

tn.set_default_backend("numpy")
print("Numpy Backend")
%timeit calculate_abc_trace(a, b, c)
tn.set_default_backend("jax")

print("JAX Backend")
%timeit np.array(calculate_abc_trace(a, b, c))

Numpy Backend
4.38 s ± 49.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
JAX Backend

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

2.1 s ± 61.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

高速化されているようです。blueqatの無料版ではGPU対応してないのえ、CPUになります。

Numpy Backend
4.38 s ± 49.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
JAX Backend
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
2.1 s ± 61.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

使ってみて

テンソルネットワーク理論を使って簡単にノードの計算や縮約の計算ができました。様々なテンソルネットワークの記述を駆使して効率的な計算が並列計算機を使ってできるというのでとても楽しいと思います。

量子コンピュータに活用するのが目的なので、シミュレータとして使って見たいと思います。

以上です。