まとめ
行列の固有値を微分するとき、固有値が縮退していると発散します。適宜対処しましょう。
きっかけ
前回の記事
を書いて、固有値が微分できて最適化できるなら、固有ベクトルも同様にできて、欲しい固有ベクトルが得られる行列が得られるんじゃないかと思い、やってみました。
実際得られたのですが、ちょっとした問題があったので、対処しました。
数学
導出はこちらのフォーラムを参考に。
https://mathoverflow.net/questions/229425/derivative-of-eigenvectors-of-a-matrix-with-respect-to-its-components
ある行列 $A$ に対し、固有値 $\lambda_i$ があって、それに対応する固有ベクトル $\vec{n}_i$ を考えて、これを微分することを考えます。
固有値の微分
\frac{\partial \lambda_i}{\partial A_{kl}} = \left(\frac{\partial A}{\partial A_{kl}}\vec{n}_i\right)\cdot\vec{n}_i
ここで $\frac{\partial A}{\partial A_{kl}}\vec{n}_i$ は射影演算子のような感じになりますね。$\vec{n}_i$ の $k$ 番目の成分を $l$ 番目に持ってきて、あとは0というベクトルになります。するとこの微分は $\vec{n}_i$ の $k$ 番目と $l$ 番目を掛け算した値になりますね。思いのほか素直です。
固有ベクトルの微分
$i$ 番目の固有ベクトル $\vec{n}_i$ について、
\frac{\partial \vec{n}_i}{\partial A_{kl}} = \sum_{j\neq i}\left[\frac{1}{\lambda_i - \lambda_j}\left(\frac{\partial A}{\partial A_{kl}}\vec{n}_i\right)\cdot\vec{n}_j\right]\vec{n}_j
と書けます。要は、他の固有ベクトルに適当な重みを付けて足し合わせていく感じです。
ここで注意したいのは、$\frac{1}{\lambda_i - \lambda_j}$ という項です。これにより、 同じ固有値が複数あるとき（物理的には縮退に対応）、この微分は発散します。
検算
準備
適当な行列 X を決めて、固有値と固有ベクトルを確認します。
import tensorflow as tf
X = tf.Variable(initial_value=[[1.0, 0.0, 0.12975], [0.0, 1.0, 0.0], [0.12975, 0.0, 1.1373545]])
eigval, eigvec = tf.linalg.eigh(X)
print(eigval)
print(eigvec)
tf.Tensor([0.9218725 1. 1.2154819], shape=(3,), dtype=float32)
tf.Tensor(
[[-0.8566836 0. 0.51584214]
[-0. -1. 0. ]
[ 0.51584214 0. 0.8566836 ]], shape=(3, 3), dtype=float32)
固有値を微分
GradientTapeを使って、最小の固有値を計算してみます。
with tf.GradientTape() as g:
g.watch(X)
eigval, eigvec = tf.linalg.eigh(X)
Y = eigval[0]
dYdX = g.gradient(Y, X)
print(dYdX)
tf.Tensor(
[[ 0.7339068 0. 0. ]
[ 0. 0. 0. ]
[-0.88382703 0. 0.2660931 ]], shape=(3, 3), dtype=float32)
- $0.7339068 \approx 0.8566836^2$
- $0.2660931 \approx 0.51584214^2$
- $0.88382703 \approx 2 \times 0.8566836 \times 0.51584214$
ということから、まあ妥当な結果ですね。$2\times$ となっているのは対称行列で下半分しか使っていないためと思われます。
固有ベクトルを微分
eigvec[i, j] は $j$ 番目の固有値に対する固有ベクトルの $i$ 番目の成分です。
with tf.GradientTape() as g:
g.watch(X)
eigval, eigvec = tf.linalg.eigh(X)
Y = eigvec[0, 1]
dYdX = g.gradient(Y, X)
print(dYdX)
tf.Tensor(
[[ 0. 0. 0. ]
[-8.158832 0. 0. ]
[ 0. 7.707127 0. ]], shape=(3, 3), dtype=float32)
こっちはめんどくさいので検算はさぼります。
ここまでは普通です。
縮退させる
以下のようにXの値を変更すると、2つの固有値が1になります。
X = tf.Variable(initial_value=[[1.1225665, 0.0, 0.12975], [0.0, 1.0, 0.0], [0.12975, 0.0, 1.1373545]])
なお、前回の記事のコードを使って求めました。
tf.Tensor([1. 1. 1.2599211], shape=(3,), dtype=float32)
tf.Tensor(
[[-0.7269436 0. 0.6866972]
[-0. -1. 0. ]
[ 0.6866972 0. 0.7269436]], shape=(3, 3), dtype=float32)
これを微分すると、両方とも
tf.Tensor(
[[nan 0. 0.]
[nan nan 0.]
[nan nan nan]], shape=(3, 3), dtype=float32)
となりました。
固有値の方が nan になるのは解せませんが、固有ベクトルの微分の方は理論式通り nan になりました。
なお、微分する固有ベクトルを3番目にして、つまり $\lambda_i - \lambda_j = 0$となるような $j$ が存在しない $i$ での固有ベクトルの微分も nan になりました。つまりTensorflowの実装では縮退があるような行列では容赦なくすべての微分が nan になるようです。
対策
いくつか回避策が考えられます。
- 問題における同値の範囲で行列を変形させる
- 多少値が変わることを覚悟で値を動かす
- そもそも固有値・固有ベクトルの微分を使わない方向でアルゴリズムを書き直す
ここではランダムに摂動を与えてみます。というのも、この微分は勾配法のために使いますので、一種のアニーリングのような感じで摂動させるだけで、最終的な結果には影響しないとできるからです。
当然、最終的な結果で固有値が縮退している場合は特別に考えなければなりません。
縮退を見つける
ともあれ、微分する前に同じ固有値があるかどうかを見つけることを考えます。
eigval[1:] - eigval[:-1]
これで一つ隣同士の固有値の差が格納された配列が得られます。tf.linalg.eigh が返す固有値はすでに昇順で並べられていることを利用すれば、絶対値を使わずとも $0$ 以上であることが分かります。そしてこれのうち１つでもほとんど０の成分があると縮退しているとします。
tf.math.reduce_any((eigval[1:] - eigval[:-1]) < 1e-20)
あとはこれを条件にして、満たさなくなるまでループさせます。 A は N行N列の行列とします。
eigval, eigvec = tf.linalg.eigh(A)
while tf.math.reduce_any((eigval[1:] - eigval[:-1]) < 1e-20):
Ap = A + tf.linalg.diag(tf.random.uniform(shape=[N])*0.001)
eigval, eigvec = tf.linalg.eigh(Ap)
判定基準の 1e-20 や摂動の大きさの 0.001 は問題によって変わってくると思います。
ひとまず僕が今やりたかったことはこれで解決しました。
おまけ
1次元の量子井戸を計算してみます。物理的な説明は
など、他のページに譲ります。
$x\in \left[-\pi, \pi\right]$ の範囲が有限のポテンシャル、それ以外は $+\infty$ であるようなポテンシャル $U(x)$ を考えたとき、どのような $U(x)$ を設定すれば基底状態の波動関数が
\psi(x) = A\exp\left(-2x^2\right)
となるかを数値的に求めます。やり方としては、ハミルトニアン $H$ を行列で書き、この固有値・固有ベクトルを求め、最小の固有値に対する固有ベクトルがこの $\psi(x)$ になるように勾配法で近づけていきます。 ただ数値計算において、関数を関数として扱えないので、範囲 $\left[-\pi, \pi\right]$ を $N$ 個の点で分割します。 $i$ 番目の点の座標を $x_i$ とすると、波動関数は $\vec{\psi}$ の$i$番目の成分 $ = \psi(x_i)$ というベクトル $\vec{\psi}$ で書けます。
気を付けるところとして、固有ベクトルの符号の自由度がありますので、損失関数は
L_+ = \sum_i^N(n_i - \psi(x_i))^2
L_- = \sum_i^N(n_i + \psi(x_i))^2
の両方が考えられることが挙げられます。イテレーションを回していると反転するということも普通にあります。今回はこのうち最小の方
L = \min(L_+, L_-)
としたらうまくいきました。
これらを踏まえ、以下のようなコードを書きました。
import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
def main():
max_x = np.pi
N = 150 + 1
dx = max_x*2/N
x = np.arange(-max_x, max_x, dx)
D = np.eye(N, k=1)
D += -1 * np.eye(N)
D += D.T
D = D/(dx**2)
m = 1.0
D_tf = tf.constant(D/(2.0*m), dtype=tf.float32)
V0_np = np.exp( - x**2/0.5)
V0_np = V0_np/np.linalg.norm(V0_np)
V0_target = tf.constant(V0_np, dtype=tf.float32)
U0 = np.zeros(shape=[N])
U = tf.Variable(initial_value=U0, dtype=tf.float32)
def calc_V(n):
H = - D_tf + tf.linalg.diag(U)
eigval, eigvec = tf.linalg.eigh(H)
while tf.math.reduce_any((eigval[1:] - eigval[:-1]) < 1e-20):
H = - D_tf + tf.linalg.diag(U) + tf.linalg.diag(tf.random.uniform(shape=[N])*0.001)
eigval, eigvec = tf.linalg.eigh(H)
print("found lambda_i+1 - lambda_i = 0")
v_raw = eigvec[:, n]
return v_raw
def calc_loss():
v0 = calc_V(0)
dplus = tf.reduce_sum((v0 - V0_target)**2)
dminus = tf.reduce_sum((v0 + V0_target)**2)
return tf.minimum(dplus, dminus)
opt = tf.keras.optimizers.Adam(learning_rate=0.001)
L = calc_loss()
v0_current = calc_V(0)
print(L)
while L > 1e-11:
opt.minimize(calc_loss, var_list=[U])
v0_current = calc_V(0)
L = (tf.abs(tf.reduce_sum(v0_current*V0_target)) - 1)**2
print(L)
plt.plot(x, U.numpy())
plt.show()
if __name__ == "__main__":
main()
これを実行し、 Ryzen 5とGTX 1060 を積んだマシンで数十分放置すると、下図のようなグラフが得られました。
つまりこのようなポテンシャルを設定すると、量子井戸中の基底状態の波動関数がGaussian波束になるだろうと分かりました。
残念ながらこれは数値解ですので、解析的に確認ができません。なにかよい関数でフィットして解いてみたいところ。人間にできるかどうかわかりませんが。