Zenn
🐷

テンソルネットワークによる連続関数の最適化

2023/12/16に公開
テンソル
tech

この記事について

量子コンピューター Advent Calendar 2023にて投稿させて頂いた記事になります。

TTOpt

本日はTTOptというある連続関数の最小値/最大値を求めるソフトウェアを少し紹介したいと思います。

原論文はこちらになりますTTOpt: A Maximum Volume Quantized Tensor Train-based Optimization and its Application to Reinforcement Learning

TTOptの背景にある技術として、quantized (quantics) tensor train / tensor cross interpolationをまずは簡単に説明したいと思います。

はじめに

本記事でのテンソルの定義は、ただの多次元配列であることに注意してください。

Quantized (Quantics) tensor train (QTT)

簡単のため0 \leq x<1の区間で定義された一次元の関数f(x)を例にとり、QTTの概念を説明する。

はじめに連続変数xの離散化するために、等間隔にメッシュを用いる。そのように離散化されたxの値は、q進数(q=2,3,,,)を用いて以下のように表現される。ここではq=2(binary codingの場合を考える。

\begin{equation} x=0 . b_1 b_2 \cdots b_R=b_1 / 2+b_2 / 2^2+\ldots+b_R / 2^R \end{equation}

ここで、b_i=0,1である。ビット列の一番左側b_1がスケールが一番大きく、b_Rがスケールが最も小さいものに対応する。ビットの数Rを増やすことで、指数関数的に幅1/2^Rが小さくなるので、関数の解像度も指数関数的に向上する。

メッシュ上の関数値F(x)は

\begin{equation} f(x=0 . b_1 b_2 \cdots b_R)=f(b_1 / 2+b_2 / 2^2+\ldots+b_R / 2^R)= F^{b_1 b_2 \cdots b_R} \end{equation}

として表され、これは、R階のテンソルとして見なすことができる。

ここでQTTでは、次の仮定をおくことで圧縮を行う。それは、同程度の長さスケール間の相関が強く、 指数関数的に異なる長さスケール間の相関が弱いという仮定です。この相関構造を自然に満たす(テンソルネットワークの1種が行列積状態/Matrix Product State(MPS) になります。また、多変数関数への拡張も同様に可能です。

Tensor cross interpolation (TCI)

tensor cross interpolation (TCI)は、多次元空間上の点をサンプルしてきて、元々の関数を行列積状態へと近似を行う技術です。ここで、関数に関連するテンソルが低ランク構造を持つことを仮定しています。

まずは簡単のため行列の圧縮について見ていきたいと思います。これをMatrix Cross Interpolation（MCI）と呼びます。
MCIのプロセスは、元々の行列\boldsymbol{J}の特定の行\boldsymbol{J}_{\boldsymbol{R}}と列\boldsymbol{J}_{\boldsymbol{C}}を選択し、これらの行と列が交差する行列\hat{\boldsymbol{J}}を利用して元の行列を\tilde{\boldsymbol{J}}で近似するというものです。

\begin{equation} \boldsymbol{J} \simeq \tilde{\boldsymbol{J}}, \quad \tilde{\boldsymbol{J}}=\boldsymbol{J}_{\boldsymbol{C}} \hat{\boldsymbol{J}}^{-1} \boldsymbol{J}_{\boldsymbol{R}} \end{equation}

ここで、\boldsymbol{J}_{\boldsymbol{C}}は選んだ列のセットからなる行列で、\boldsymbol{J}_{\boldsymbol{R}}は選んだ行のセットからなる行列、\hat{\boldsymbol{J}}は行と列が交差する要素からなる行列である。
（以下の図は、こちらの論文のFigure.2を参考に作りました。)

この近似には次の二つの重要な性質を持つことが知られています：

  1. 内挿公式である。選んだ行と列の交差点上で、元の行列と近似行列の要素が等しくなる。
  2. もし、元の行列のランクと同じ数の行と列を選ぶと、近似行列は元の行列と等しくなる。

ここで、どのように行と列を選ぶのかという問題が起こる。愚直に最適な行と列の組みを探そうとすると、「組合せ爆発」を起こしてします。それを回避するために、いくつかのヒューリスティックな方法が提案されています。その中の一つにmaxvolアルゴリズムがあります。これはある行列の行と列を選択し、その行列式の絶対値（|det(\hat{\boldsymbol{J}})|）を最大化する要素を探すというプロセスです。つまり、選ばれた要素の絶対値が最大になるように行と列を探索するもので、これにより行列の近似が得られます。この精度についてはここでは割愛する。ただし、誤差の評価についてはフロベニウスノルムの意味で特異値分解よりも（当然）悪くなります。

このアルゴリズムはテンソルに拡張することができ、それが「tensor cross interpolation」と呼ばれています。詳細については原論文を参照。

しかし通常のmaxvolでは、絶対値の大きな要素を探すので符号までは考慮しません。これでは、連続関数の最小値最大値問題を解きづらくなってしまいます。そこで論文では、TCIの内部で(\hat{\boldsymbol{J}})）を評価するのではなく、以下の関数を代わりに評価しています。
J_{\min }はこれまでTCIの内部で評価してきた関数地で最も低い値になります。

\begin{equation} \mathbf{g}(\boldsymbol{x})=\frac{\pi}{2}-\operatorname{atan}\left(\mathrm{J}(\boldsymbol{x})-J_{\min }\right), \end{equation}

実際に使ってみる。

以下のソースコードは、TTOptのgithubのnotebookのqtt.pyから持ってきました。

ここでは答えが厳密にわかる関数(X-1)^2 * ((X+1)^2 + 0.01)の最小値問題を解いてみることにします。

TTOptでは、最初に最大のボンド次元(=rank)を指定する必要があります。
ここでは、rank=2としています（どのような値を設定すれば事前にわからないのが難点かもしれません）

import numpy as np


from ttopt import TTOpt
from ttopt import ttopt_init


np.random.seed(42)


d = 1                      # Number of function dimensions:
rank = 2                   # Maximum TT-rank while cross-like iterations
def f(X):                   # Target function
    #return np.sin(X) -10
    return (X-1)**2 * ((X+1)**2 + 0.01)
    #return np.sum(np.abs(X * np.sin(X) + 0.1 * X), axis=1)


# We initialize the TTOpt class instance with the correct parameters:
tto = TTOpt(
    f=f,                    # Function for minimization. X is [samples, dim]
    d=d,                    # Number of function dimensions
    a=-10.,                 # Grid lower bound (number or list of len d)
    b=+10.,                 # Grid upper bound (number or list of len d)
    p=2,                    # The grid size factor (there will n=p^q points)
    q=10,                   # The grid size factor (there will n=p^q points)
    evals=1.E+5, )           # Number of function evaluations
    #name='Alpine',          # Function name for log (this is optional)
    #x_opt_real=np.ones(d),  # Real value of x-minima (x; this is for test)
    #y_opt_real=0.,          # Real value of y-minima (y=f(x); this is for test)
    #with_log=True)


# And now we launching the minimizer:
tto.optimize(rank)


# We can extract the results of the computation:
x = tto.x_opt          # The found value of the minimum of the function (x)
y = tto.y_opt          # The found value of the minimum of the function (y=f(x))
k_c = tto.k_cache      # Total number of cache usage (should be 0 in this demo)
k_e = tto.k_evals      # Total number of requests to func (is always = evals)
k_t = tto.k_total      # Total number of requests (k_cache + k_evals)
t_f = tto.t_evals_mean # Average time spent to real function call for 1 point
                       # ... (see "ttopt.py" and docs for more details)
print(x)
print(y)
# We log the final state:
#print('-' * 70 + '\n' + tto.info() +'\n\n')

実行すると正しく解が得られていることがわかりました。今回は1変数関数の自明な例でしたが、多変数関数の例や人工的な関数などに適用するなど色々と遊んでみると面白いかもしれません。

TToptを使った応用研究

量子アニーリングとTTOptを比較している論文

https://inspirehep.net/literature/2079821

VQEの最適化にTTOptを使用
https://arxiv.org/pdf/2306.02024.pdf

Discussion