テンソルネットワークへの関数の埋め込み

この記事は量子コンピューター Advent Calendar 2023の一日目の記事です。

お久しぶりです．wottoです．最近Twitterアカウントを新しく作りました．元気に生きてます．

振幅埋め込み

量子機械学習など，量子計算アルゴリズムにおけるサブルーチンとして，以下のように関数値を状態の振幅として準備したい（state preparation）状況はしばし発生します．
$$
\ket\psi=\frac{1}{\sqrt{\sum_{i=0}^{2^n-1}|f(x_i)|^2}}\sum_{j=0}^{2^n-1}f(x_j)\ket{j}
$$
ここで$n$は用いる量子ビット数とします．state preparationを実行するプロトコルはいくつか提案されていますが，その一つにテンソルネットワークを用いたものがあります．$\ket\psi$を表現するテンソルネットワークをまず準備し，そのテンソルネットワークを量子回路として埋め込むことでこれは達成されます¹．

テンソルネットワークを用いた関数の表現

ここで，$f$は定義域が$[0,1]$の関数であるとします．今やりたいことは，$[0,1]$を$2^n$個に分割し，離散化した関数を埋め込むことです：
$$
\ket\psi\propto\sum_{a=0}^{2^n-1}f\left(\frac{a}{2^n}\right)\ket{a}
$$
$n$が増えるごとに離散化のエラーは指数的に改善されます．$\ket\psi$は$2^n$の長さのベクトルで表現されますが，これを以下のようにMPS(Matrix Product States)と呼ばれるテンソルネットワークを用いてより安いコストで表現することを考えます．

赤く記したindexの次元をボンド次元と呼びます．サイズ$2^n$のベクトル（＝rankが$n$のテンソル）をMPSに直すには，以下のようにsequentialにSVD（Singular Value Decomposition, 特異値分解）を適用すればいいです．

SVDによる分解をする際に，近似を導入することでボンド次元を減らすことができます．具体的には．1）小さな特異値を切り捨てる　あるいは　2）ボンド次元の上限$\chi$を設定し，$\chi$個の大きな特異値のみを残し，その他は切り捨てる　という戦略を取ります．

理論的な保証

離散化した関数をMPSで常に効率的に表現できるとは限りません．実際，最悪のケースだと，十分な精度で近似するのに必要なボンド次元$\chi$は$n$に対して指数関数的に増加します．

ところが，[1]によると，関数$f$のフーリエ係数が波数に対して指数的に減衰するなら，$n$が十分大きいとき，MPS近似のエラーはボンド次元$\chi$に対して指数的に減衰し，また$n$に依存しないことが主張されています．これを信じるならば，多くの興味ある関数$f$に対し，十分に離散化すると，$f$の離散化はMPSで効率的に表現できることになります．

実装

ベンチマークとして正規分布の密度関数をMPSに埋め込むことを考えます．まず，以下のように離散化した正規分布のデータを準備します．今回は$n=10$としました．

mu = 0             # mean
sigma = 1         # standard deviation
x_min = -5        # truncation lower bound
x_max = 5         # truncation upper bound
n_bins = 2**10    # number of bins

x = np.linspace(x_min, x_max, n_bins)
pdf_values = norm.pdf(x, mu, sigma)
probabilities = np.array(pdf_values) / sum(pdf_values)

MPSを作成する関数と，MPSを状態ベクトルに戻す関数を作成します．今回は手で適当に書きましたが，本来は適切なライブラリを用いるべきであることは言うまでもないです．

def create_MPS(psi, n, max_singular_values=None, max_truncation_err=0.0):
    M = []
    tmp = psi
    fidelity = 1.0
    bdim = 1
    bond_dims = []
    for i in range(n-1):
        U, s, Vh = np.linalg.svd(tmp.reshape(bdim*2, -1), full_matrices=False)
        cum_error = np.sqrt(1.0 + 1e-15 - np.cumsum(s**2)) # truncation error: \sqrt{s_(i+1)^2+...+s_n^2}
        num_sing_vals_keep = np.count_nonzero((cum_error > max_truncation_err).astype(np.int32)) + 1
        num_sing_vals_keep = min(num_sing_vals_keep, len(s))
        if max_singular_values is not None:
            num_sing_vals_keep = min(max_singular_values, num_sing_vals_keep)
        s = s[:num_sing_vals_keep]
        U = U[:,:num_sing_vals_keep]
        Vh = Vh[:num_sing_vals_keep]
        M.append(U.reshape(bdim, 2, -1).transpose(1,0,2))
        tmp = np.dot(np.diag(s), Vh)
        bdim = num_sing_vals_keep
        fidelity *= (np.sum(s**2))
        bond_dims.append(bdim)
    M.append(tmp.reshape(bdim,2,1).transpose(1,0,2))
    return M, fidelity, bond_dims

def contract_MPS(M):
    tmp = M[0]
    for i in range(1, len(M)):
        tmp = np.einsum("iab,jbc->ijac",tmp,M[i]).reshape(2**(i+1),1,-1)
    return tmp

上で作った関数を用いて離散化した関数値データをMPSに変換します．まずはmax_truncation_error=1e-5として，小さな特異値を切り取る近似を用いて実行します．

tensor = np.sqrt(np.array(probabilities)).reshape(2,2,2,2,2,2,2,2,2,2)
M, fidelity, bond_dims = create_MPS(tensor, 10, None, 1e-5)
print("bond dimensions:", bond_dims)
print(f"estimated fidelity: {fidelity}")

tensor_mps = contract_MPS(M)
print(f"l2-norm between |mps> and |true>: {np.linalg.norm(tensor_mps.flatten() - tensor.flatten())}")

def kl_divergence(p, q):
    return np.sum(p * np.log(p / q))

def js_divergence(p, q):
    m = 0.5 * (p + q)
    return 0.5 * kl_divergence(p, m) + 0.5 * kl_divergence(q, m)

mps_probability = np.power(tensor_mps, 2).flatten()
mps_probability /= np.sum(mps_probability)

print(f"KL-divergence between p_mps and p_true: {kl_divergence(probabilities, mps_probability)}")
print(f"JS-divergence between p_mps and p_true: {js_divergence(probabilities, mps_probability)}")

結果は以下のようになりました．

bond dimensions: [2, 4, 6, 4, 4, 3, 3, 3, 2]
estimated fidelity: 0.9999999994423893
l2-norm between |mps> and |true>: 9.851397877556107e-06
KL-divergence between p_mps and p_true: 1.9410098748939853e-10
JS-divergence between p_mps and p_true: 4.8525018738518734e-11

bond dimensionsの欄に載っているものがボンド次元のリスト（今回は$n=10$なので$9$個）です．本来，MPSで$2^n$のデータを表現するには$2^5=32$程度のボンド次元が必要ですが，今回はそれよりずっと小さなボンド次元で十分精度良く近似できていることがわかります．

次に，最大のボンド次元を4として再び実行してみます．

max_bond_dimension = 4
M, fidelity, bond_dims = create_MPS(tensor, 10, max_bond_dimension, 1e-5)

結果は以下のようになりました．

bond dimensions: [2, 4, 4, 4, 4, 4, 4, 4, 2]
estimated fidelity: 0.9999998592080019
l2-norm between |mps> and |true>: 0.00014182501983998135
KL-divergence between p_mps and p_true: 4.025722071084061e-08
JS-divergence between p_mps and p_true: 1.0057021666133458e-08

先ほどよりは精度は少し劣るものの，より効率的に表現できていることがわかります．プロットして確認してみます．

目で見てほとんど違いがわからないほど，正確に元の確率分布を再現できていることがわかります．

参考文献

[1] Efficient MPS representations and quantum circuits from the Fourier modes of classical image data, arXiv:2311.07666

1. テンソルネットワークで表現された波動関数を量子回路に埋め込む方法については色々提案されていますが，これがbestだみたいなものはまだ定まっていないと思います． ↩