これまで、Flux.jlの使い方を見てきましたが、この記事では応用例を見ていきます。
これまでの記事は、
- Juliaで機械学習：深層学習フレームワークFlux.jlを使ってみる その１：基本編
- Juliaで機械学習：深層学習フレームワークFlux.jlを使ってみる その２：線形回帰編
- Juliaで機械学習：深層学習フレームワークFlux.jlを使ってみる その3：ニューラルネットとバッチ正規化編
- Juliaで機械学習：Flux.jlで自由自在にオリジナルレイヤーを組んでみよう
- Flux.jlで量子力学
-
Juliaで機械学習：Flux.jlで自由自在にオリジナルレイヤーを組んでみよう 2023年版
などがあります。
今回は、量子スピン系の波動関数を、制限ボルツマンマシン(RBM)で表現し、エネルギーが最小となるような状態（基底状態）を求めることにします。
以下の説明は最終的には制限ボルツマンマシンによる波動関数について述べますが、それまでの流れは変分モンテカルロ法の説明となっています。
解くべき模型の説明：量子スピン系
今回は、N個の量子スピン系を考えます。ハミルトニアンは
H = J \sum_i S_i^z S_{i+1}^z + h_x \sum_i S_i^x
とします。これは横磁場イジング模型と呼ばれています。量子スピンが2個(N=2)の場合、波動関数は
|\psi \rangle = c_{11} |\uparrow \uparrow \rangle + c_{12} |\uparrow \downarrow \rangle + c_{21} |\downarrow \uparrow \rangle +c_{22} |\downarrow \downarrow \rangle
と書けます。ここで、$|s_1 s_2 \rangle$は1番目のスピンが$s_1$($= \uparrow,\downarrow$)、2番目のスピンが$s_2$($= \uparrow,\downarrow$)であることを意味しています。また、$S_i^z$や$S_i^x$は、$i$番目のスピンに対して、
\displaylines{
S_i^z |s_1,\cdots,s_i=\uparrow,\cdots,s_N \rangle &= &|s_1,\cdots,s_i=\uparrow,\cdots,s_N \rangle \\
S_i^z |s_1,\cdots,s_i=\downarrow,\cdots,s_N \rangle &=& -1|s_1,\cdots,s_i=\downarrow,\cdots,s_N \rangle \\
S_i^x |s_1,\cdots,s_i=\uparrow,\cdots,s_N \rangle &= &|s_1,\cdots,s_i=\downarrow,\cdots,s_N \rangle \\
S_i^x |s_1,\cdots,s_i=\downarrow,\cdots,s_N \rangle &=& |s_1,\cdots,s_i=\uparrow,\cdots,s_N \rangle
}
となります。ですので、2個のスピンの場合には、
\displaylines{
S_1^z |\psi \rangle = c_{11} |\uparrow \uparrow \rangle + c_{12} |\uparrow \downarrow \rangle - c_{21} |\downarrow \uparrow \rangle - c_{22} |\downarrow \downarrow \rangle \\
S_1^x |\psi \rangle = c_{11} |\downarrow \uparrow \rangle + c_{12} |\downarrow \downarrow \rangle + c_{21} |\uparrow \uparrow \rangle + c_{22} |\uparrow \downarrow \rangle
}
となります。これらの操作をベクトルで表すことにします。波動関数の係数をベクトルで表記すると、
\vec{\psi} = \left( \begin{matrix}
c_{11} \\
c_{12} \\
c_{21} \\
c_{22}
\end{matrix} \right)
となりますが、この時、$S_1^z$および$S_1^x$は、
\displaylines{
S_1^z = \sigma_z \otimes 1 \\
S_1^x = \sigma_x \otimes 1
}
と書けば、
\displaylines{
S_1^z \vec{\psi} = \left( \begin{matrix}
c_{11} \\
c_{12} \\
-c_{21} \\
-c_{22}
\end{matrix} \right) \\
S_1^x \vec{\psi} = \left( \begin{matrix}
c_{21} \\
c_{22} \\
c_{11} \\
c_{12}
\end{matrix} \right)
}
とかけます。これらをふまえると、ハミルトニアンは
H = J \sum_i 1 \otimes \cdots S_i^x \otimes S_{i+1}^z \otimes \cdots 1 + \sum_i h_x 1 \otimes \cdots S_i^x \otimes \cdots 1
と書けます。
ここで、行列のサイズは、各サイトでアップスピンかダウンスピンなので、$2^N$です。
Juliaでの実装：厳密対角化
上で定義したハミルトニアン行列を対角化し、固有値固有ベクトルを求めてみます。まず、パウリ行列を定義しておきます。
const σx = [
0 1
1 0
]
const σy = [
0 -im
im 0
]
const σz = [
1 0
0 -1
]
const σ0 = [
1 0
0 1
]
次に、$S_i^zS_j^z$を作成します。これは
function SziSzjmatrix(i,j,numspin)
hi = zeros(Float64,2,2)
hj = zeros(Float64,2,2)
for isite=1:numspin
if isite == i
hj .= σz
elseif isite == j
hj .= σz
else
hj .= σ0
end
if isite == 1
hi .= hj
else
hi = kron(hi,hj)
end
end
return hi
end
で作ることができます。定義通りにクロネッカー積を使用しています。
$S_i^x$は
function Sximatrix(i,numspin)
hi = zeros(Float64,2,2)
hj = zeros(Float64,2,2)
for isite=1:numspin
if isite == i
hj .= σx
else
hj .= σ0
end
if isite == 1
hi .= hj
else
hi = kron(hi,hj)
end
end
return hi
end
です。
この二つを使ってハミルトニアン行列を作るには、
function make_Hamiltonian(J,hx,numspin)
H = zeros(Float64,2^numspin,2^numspin)
for i=1:numspin
j = i+1
j += ifelse(j > numspin,-numspin,0)
H += SziSzjmatrix(i,j,numspin)*J
H += Sximatrix(i,numspin)*hx
end
return H
end
とします。
行列のサイズが$2^N$なのであまり大きな系は計算できませんが、固有値問題はLinearAlgebraパッケージのeigenを用いることで可能です。
このハミルトニアン行列を対角化して得られた最低固有値を基底状態のエネルギーとします。制限ボルツマンマシンを用いてこのエネルギーにどこまで近づけられるか、ということが、今回の記事の目的です。
基底状態の探索
シュレーディンガー方程式は
H |\psi \rangle = E |\psi \rangle
となります。固有値$E$の中で最小の固有値を求めたい、というのが今回の問題です。最小固有値を求めたい問題は、
\displaylines{
\langle \psi | H |\psi \rangle &=& E \langle \psi |\psi \rangle \\
E &=& \frac{\langle \psi | H |\psi \rangle}{\langle \psi |\psi \rangle }
}
となりますから、この$E$を最小化するような$|\psi \rangle|を求める問題に帰着されます。なお、一次元シュレーディンガー方程式の最小化問題に関しては、以前の記事で解説しています。
この式を基底$|s_1,s_2,\cdots,s_N \rangle$で表現するため、完全系を挟みます：
\displaylines{
E &=& \sum_{s_1,\cdots,s_N} \sum_{s_1',\cdots,s_N'} \frac{\langle \psi ||s_1,s_2,\cdots,s_N \rangle \langle s_1,s_2,\cdots,s_N |H||s_1',s_2',\cdots,s_N' \rangle \langle s_1',s_2',\cdots,s_N' |\psi \rangle}{\langle \psi |\psi \rangle }
}
ここで、$s_1,s_2,\cdots,s_N$の関数としての波動関数を
\psi(s_1,s_2,\cdots,s_N) = \langle s_1',s_2',\cdots,s_N' |\psi \rangle
とし、ハミルトニアンを基底$|s_1,s_2,\cdots,s_N \rangle$で表したものを
H(s_1,s_2,\cdots,s_N,s_1',s_2',\cdots,s_N') = \langle s_1,s_2,\cdots,s_N |H||s_1',s_2',\cdots,s_N' \rangle
とすると、
\displaylines{
E &=& \sum_{s_1,\cdots,s_N} \frac{\psi(s_1,s_2,\cdots,s_N)^{\ast} \sum_{s_1',\cdots,s_N'} H(s_1,s_2,\cdots,s_N,s_1',s_2',\cdots,s_N')\psi(s_1',s_2',\cdots,s_N')
}{\sum_{s_1,\cdots,s_N} |\psi(s_1,s_2,\cdots,s_N)|^2 }
}
となります。スピンが2個の場合であれば、$s_1,s_2$の二つですから、それぞれの和は4通りになります。もしスピンが増えてくると、和は$2^N$通りになり、計算が困難になっていきます。その場合には、モンテカルロ法で評価することになります。
モンテカルロ法の場合には、この式をさらに変形して、
\displaylines{
E &=& \sum_{s_1,\cdots,s_N} \frac{\psi(s_1,s_2,\cdots,s_N)^{\ast}\psi(s_1,s_2,\cdots,s_N)}{\sum_{s_1,\cdots,s_N} |\psi(s_1,s_2,\cdots,s_N)|^2} \frac{\sum_{s_1',\cdots,s_N'} H(s_1,s_2,\cdots,s_N,s_1',s_2',\cdots,s_N')\psi(s_1',s_2',\cdots,s_N')
}{ \psi(s_1,s_2,\cdots,s_N)} \\
&=& \sum_{s_1,\cdots,s_N} \frac{|\psi(s_1,s_2,\cdots,s_N)|^2}{\sum_{s_1,\cdots,s_N} |\psi(s_1,s_2,\cdots,s_N)|^2} \frac{\sum_{s_1',\cdots,s_N'} H(s_1,s_2,\cdots,s_N,s_1',s_2',\cdots,s_N')\psi(s_1',s_2',\cdots,s_N')
}{ \psi(s_1,s_2,\cdots,s_N)} \\
&=& \sum_{s_1,\cdots,s_N} P(s_1,s_2,\cdots,s_N) E_L(s_1,s_2,\cdots,s_N)
}
とします。ここで、
\displaylines{
P(s_1,s_2,\cdots,s_N) &\equiv& \frac{|\psi(s_1,s_2,\cdots,s_N)|^2}{\sum_{s_1,\cdots,s_N} |\psi(s_1,s_2,\cdots,s_N)|^2} \\
E_L(s_1,s_2,\cdots,s_N) &\equiv& \frac{\sum_{s_1',\cdots,s_N'} H(s_1,s_2,\cdots,s_N,s_1',s_2',\cdots,s_N')\psi(s_1',s_2',\cdots,s_N')
}{ \psi(s_1,s_2,\cdots,s_N)}
}
です。
変分試行波動関数
さて、エネルギーを最小化する波動関数$\psi(s_1,s_2,\cdots,s_N)$をどのように決めれば良いでしょうか。一般論としては、波動関数にあるパラメータ$\vec{\theta} = (\theta_1,\cdots,\theta_M)$を導入し、
\psi_{\vec{\theta}}(s_1,s_2,\cdots,s_N)
がエネルギーを最小になるようなパラメータ$\vec{\theta}$を見つければ良い、ということになります。この波動関数を変分試行波動関数と呼びます。
エネルギーの計算と最小値探索：少数スピン系の場合
次に、エネルギーの計算です。簡単のため、波動関数$\psi_{\vec{\theta}}(s_1,s_2,\cdots,s_N)$は実数であるとします。もしモンテカルロ法を使わなくても計算ができるようなサイズであれば、以下のように計算することができます。
まず、波動関数の各成分をベクトルにします。２個のスピンであれば、
\vec{\psi}_{\vec{\theta}} = \left( \begin{matrix}
\psi_{\vec{\theta}}(\uparrow,\uparrow) \\
\psi_{\vec{\theta}}(\uparrow,\downarrow) \\
\psi_{\vec{\theta}}(\downarrow,\uparrow) \\
\psi_{\vec{\theta}}(\uparrow,\uparrow)
\end{matrix}
\right)
となります。このベクトルを用いれば、エネルギーは
E =\frac{ \vec{\psi}_{\vec{\theta}}^{T} H \vec{\psi}_{\vec{\theta}} }{\vec{\psi}_{\vec{\theta}}^T \vec{\psi}_{\vec{\theta}} }
と簡単に求まります。
この計算を行うために、まずハミルトニアン行列のあるインデックスが決まった時、その時のスピンの状態を出力する関数を
function get_S(istate,numspin)
Sj = zeros(Int8,numspin)
k = lpad(string(istate-1,base=2),numspin,"0")
for ispin=1:numspin
Sj[ispin] = ifelse(k[ispin]=='0',-1,1)
end
return Sj
end
のように作ります。
これを用いれば、
Ss = [get_S(istate,numspin) for istate in 1:2^numspin]
vecψ = ψ.(Ss)
のようにすることで、$\vec{\psi}$が求まります。ここで、ψ.(Ss)はSsを引数とする関数です。
これらを組み合わせることで、エネルギーは
function calc_energy_full(H,ψ,Ss)
vecψ = ψ.(Ss)
Hψ = H*vecψ
E = dot(vecψ,Hψ)/dot(vecψ,vecψ)
return E
end
で求まります。この関数は、Fluxを使うことで、
θ,re = Flux.destructure(ψ)
E,dEdθ = Flux.withgradient(θ -> calc_energy_full(H,re(θ),Ss),θ)
のようにして値とその微分を計算できます。ここで、
θ,re = Flux.destructure(ψ)
は、Fluxのモデルをパラメータとその構造に分離するもので、re(θ)(x)のような形で使います。この形にすることでパラメータ$\vec{\theta}$に関する微分がやりやすくなります。
パラメータの訓練には、Optimisers.jlを使います。
function training_full!(H,ψ,numspin,nt)
θ,re = Flux.destructure(ψ)
state = Optimisers.setup(Optimisers.AdamW(), θ)
Ss = [get_S(istate,numspin) for istate in 1:2^numspin]
for it=1:nt
E,dEdθ = Flux.withgradient(θ -> calc_energy_full(H,re(θ),Ss),θ)
Optimisers.update!(state, θ, dEdθ[1])
println("$it energy: $E")
end
return θ,re
end
となります。これで、関数$\psi_{\vec{\theta}}(s_1,s_2,\cdots,s_N)$をFluxのモデルとして表すことができれば、パラメータを調節して最低エネルギーを持つ波動関数を求めることができます。
エネルギーの計算と最小値探索：一般の場合
スピン数が増えてくるとハミルトニアン行列のサイズが大きくなりすぎて計算ができません。そこで、モンテカルロ法を使ってエネルギーを評価することにします。
ハミルトニアンを行列で表現する代わりに、どのような演算子が含まれているかが記録されたハミルトニアンを独自structとして作ることにします。
まず、演算子の型としてTerm型を
struct Term{nspin}
term::Vector{Int8}
value::Ref{Float64}
end
function Term(nspin)
term = Array{Int8}(undef,nspin)
return Term{nspin}(term,0)
end
と定義します。この型はどのような演算子なのかを記録するものです。termは0,1,2,3のどれかをとります。0は単位行列、1は$\sigma_x$,2は$\sigma_y$,3は$\sigma_z$を表すこととします。
まず、$S_i^x$は
function Sxiterm(nspin,i::T,value) where T <: Integer
Sxi = Term(nspin)
Sxi.value[] = value
for isite=1:nspin
if isite == i
Sxi.term[isite] = 1
else
Sxi.term[isite] = 0
end
end
return Sxi
end
とします。Sxi.term[isite] において、isite=i番目の演算子が$\sigma_x$であるとしています。つまり、これは$1 \otimes \cdots \sigma_x \otimes 1$というクロネッカー積の情報を保持していることになります。
次に、$S_i^z S_j^z$は
function SziSzjterm(nspin,i,j,value)
SziSzj = Term(nspin)
SziSzj.value[] = value
for isite=1:nspin
if isite == i || isite == j
SziSzj.term[isite] = 3
else
SziSzj.term[isite] = 0
end
end
return SziSzj
end
です。今度は、$i$番目と$j$番目のサイトに$\sigma_z$が入ることになります。
ハミルトニアンは、これらの演算子の和ですから、
struct Hamiltonian{nspin}
terms::Vector{Term{nspin}}
end
function Base.length(h::Hamiltonian)
return length(h.terms)
end
function Hamiltonian(nspin)
terms = Array{Term{nspin}}(undef,0)
return Hamiltonian{nspin}(terms)
end
function Base.push!(h::Hamiltonian,term::Term)
push!(h.terms,term)
end
と定義します。これで、Term型をpush!でハミルトニアン型に追加できるようになりました。
次に、あるスピン状態$|s_1,s_2,\cdots,s_N \rangle$にハミルトニアンを作用させた時に現れる状態がどうなるかを調べる関数を作ります。例えば、2スピンの場合には$S_1^x$を作用させると
S_1^x | \uparrow \downarrow \rangle = | \downarrow \downarrow \rangle
となりますか、これは$| \uparrow \downarrow \rangle $という状態が$| \downarrow \downarrow \rangle $という状態に変わったことを意味します。$S_i^z$ならば、
S_1^z | \downarrow \downarrow \rangle = -| \downarrow \downarrow \rangle
のように状態は変化しませんが係数がつきます。このような状態を見つけることができれば、エネルギーのモンテカルロ計算に現れる
\sum_{s_1',\cdots,s_N'} H(s_1,s_2,\cdots,s_N,s_1',s_2',\cdots,s_N')\psi(s_1',s_2',\cdots,s_N')
という計算において、$H(s_1,s_2,\cdots,s_N,s_1',s_2',\cdots,s_N')$が非ゼロである部分だけを取り出すことができるようになります。
ということで、ある状態を指定した際に、指定した演算子をその状態に作用させた時に現れる状態を調べる関数を
function get_nonzero_index(term::Term{nspin},S,Sj) where nspin
Sj .= S
value = term.value[]
for isite=1:nspin
kind = term.term[isite]
if kind == 0
elseif kind == 1
Sj[isite] *= -1
elseif kind == 3
value *= ifelse(Sj[isite]==1,1,-1)
end
end
return value
end
とします。次に、ハミルトニアン全体において同じことをする関数を
function get_nonzero_indices!(h::Hamiltonian,Ss,values,S)
numterms = length(h)
for ikind = 1:numterms
term = h.terms[ikind]
values[ikind] = get_nonzero_index(term,S,Ss[ikind])
end
return numterms
end
とします。この関数は、状態を要素とした配列を得ることができます。
エネルギー期待値の微分
次は、エネルギーをモンテカルロ法で計算する関数と、その微分を計算する関数です。エネルギーの計算の式はすでに得られていますから、微分の式を導出します。あるパラメータ$a$で微分することを考えます。
\displaylines{
\frac{\partial E}{\partial a} &=& \left( \frac{\partial \langle \psi |}{\partial a} H |\psi \rangle + \langle \psi | H \frac{\partial |\psi \rangle}{\partial a}\right) \frac{1}{\langle \psi |\psi \rangle} -\frac{\langle \psi | H |\psi \rangle}{\langle \psi |\psi \rangle^2 } \left( \frac{\partial \langle \psi |}{\partial a} |\psi \rangle + \langle \psi |\frac{\partial |\psi \rangle}{\partial a} \right)
}
ここで、波動関数が実数であると仮定すると、
\displaylines{
\langle \psi | H \frac{\partial |\psi \rangle}{\partial a} = \left( \langle \psi | H \frac{\partial |\psi \rangle}{\partial a}\right)^{\ast} = \frac{\partial \langle \psi |}{\partial a} H |\psi \rangle \\
\langle \psi |\frac{\partial |\psi \rangle}{\partial a} = \left( \langle \psi |\frac{\partial |\psi \rangle}{\partial a}\right)^{\ast} = \frac{\partial \langle \psi |}{\partial a} |\psi \rangle
}
となりますから、
\displaylines{
\frac{\partial E}{\partial a} &=& 2 \frac{ \frac{\partial \langle \psi |}{\partial a} H |\psi \rangle }{\langle \psi |\psi \rangle} -2 \frac{\langle \psi | H |\psi \rangle}{\langle \psi |\psi \rangle } \frac{ \frac{\partial \langle \psi |}{\partial a} |\psi \rangle }{\langle \psi |\psi \rangle }
}
が得られます。これに、基底$|s_1,s_2,\cdots,s_N \rangle$の完全系を挟むと、
\displaylines{
\frac{\partial E}{\partial a} &= 2 \sum_{s_1,\cdots,s_N}\frac{ \frac{\partial \psi(s_1,\cdots,s_N) }{\partial a} }{\psi(s_1,\cdots,s_N)}\sum_{s_1',\cdots,s_N'}\psi(s_1,\cdots,s_N) H(s_1,\cdots,s_N,s_1',\cdots,s_N') \psi(s_1',\cdots,s_N') \frac{1}{\langle \psi |\psi \rangle} \nonumber \\ &-2 E \sum_{s_1,\cdots,s_N} P(s_1,\cdots,s_N) \frac{ \frac{\partial \psi(s_1,\cdots,s_N) }{\partial a} }{\psi(s_1,\cdots,s_N)} \\
&= 2 \sum_{s_1,\cdots,s_N}\frac{ \frac{\partial \psi(s_1,\cdots,s_N) }{\partial a} } {\psi(s_1,\cdots,s_N)}E_L(s_1,\cdots,s_N) P(s_1,\cdots,s_N) -2 E \sum_{s_1,\cdots,s_N} P(s_1,\cdots,s_N) \frac{ \frac{\partial \psi(s_1,\cdots,s_N) }{\partial a} }{\psi(s_1,\cdots,s_N)}
}
確率$P(s_1,\cdots,s_N)$による期待値計算を$\langle \rangle$と書くと、
\frac{\partial E}{\partial a} = 2 \langle \frac{ \frac{\partial \psi }{\partial a} }{\psi} E_L \rangle -2 E \langle \frac{ \frac{\partial \psi }{\partial a} }{\psi} \rangle
となります。ですので、この量をモンテカルロ法で計算すれば良いということになります。
モンテカルロ法
確率分布$P(s_1,\cdots,s_N)$によるモンテカルロ法ですが、マルコフ連鎖モンテカルロ法を用います。この時、$|\psi(s_1,\cdots,s_N)|^2$が分子にありますが、これを統計力学で出てくるようなボルツマン重みとみなせば、古典イジング模型のモンテカルロ法のように、スピンをひっくり返しながらその重みの変化を使ってアクセプト率を計算することができることがわかります。つまり、i番目のスピンをひっくり返した時のアクセプト率は
A(s_1,\cdots,s_i,\cdots,s_N \rightarrow s_1,\cdots,\bar{s_i},\cdots,s_N) = {\rm min} \left(1,|\frac{\psi(s_1,\cdots,\bar{s_i},\cdots,s_N)}{\psi(s_1,\cdots,s_i,\cdots,s_N)}|^2 \right)
となります。ここで$\bar{s_i}$は$s_i$と反対向きのスピン、という意味です。
モンテカルロ法で計算する際には、$E_L$という量が必要ですので、これを計算する関数を
function calc_EL!(ham,ψ,S,Sjs,values)
El = 0.0
ψi = ψ(S)
numS = get_nonzero_indices!(ham,Sjs,values,S)
for i=1:numS
Sj = Sjs[i]
value = values[i]
ψj = ψ(Sj)
El += value*ψj
end
El /= ψi
return El
end
とします。これを用いることで、
function calc_dEda(ham, re,θ,numspin,NMC=3000)
rbm = re(θ)
S = rand([-1,1],numspin)
Snew = deepcopy(S)
EL = 0.0
Eldlnψ = zero(θ)
dlnψ = zero(θ)
maxterm = 2^numspin
Sjs = Array{Vector{Int8}}(undef,maxterm)
for i=1:maxterm
Sjs[i] = zeros(Int8,numspin)
end
values = zeros(Float64,maxterm)
acceptance = 0
println("--------------------------------")
print_wavefunction(rbm,numspin)
for k=1:NMC
ψold = rbm(S)
Snew .= S
ispin = rand(1:numspin)
Snew[ispin] *= -1
ψnew = rbm(Snew)
ratio = abs(ψnew/ψold)^2
r = rand()
if r < ratio
#println("accepted")
#println("$k -th Sold $S Snew $Snew ratio $ratio")
S .= Snew
ψ = ψnew
acceptance += 1
else
ψ = ψold
#println("rejected")
end
Enew = calc_EL!(ham,rbm,S,Sjs,values)
dlnψnew = Flux.gradient(θ -> re(θ)(S),θ)[1]/ ψ
EL += Enew
Eldlnψ .+= Enew* dlnψnew
dlnψ .+= dlnψnew
end
println("acceptance ratio: $(acceptance/NMC)")
EL /= NMC
Eldlnψ /= NMC
dlnψ /= NMC
dEdθ = 2*(Eldlnψ .- EL*dlnψ)
return EL,dEdθ
end
によってモンテカルロ法で期待値とその微分を計算できます。
訓練
以上のコードを組み合わせることで、
function training!(ham,rbm,numspin,nt)
θ,re = Flux.destructure(rbm)
state = Optimisers.setup(Optimisers.Adam(), θ)
for it=1:nt
EL,dEdθ = calc_dEda(ham, re,θ,numspin)
Optimisers.update!(state, θ, dEdθ)
println("$it energy: $EL")
end
return θ,re
end
で訓練をすることができます。
制限ボルツマンマシン
以上の準備を済ませた上で、具体的な波動関数$\psi(s_1,\cdots,s_N)$の構成を考えます。ここでは、制限ボルツマンマシン(RBM)を使います。RBMは、
\psi(s_1,\cdots,s_N) = \sum_{h_1=\pm 1,\cdots,h_M=\pm 1} \exp \left[ \sum_{j=1}^N a_j s_j + \sum_{i=1}^M b_i h_h + \sum_{i,j} W_{ij} h_i s_j \right]
という関数です(Carleo et al., Science 355,602–606 (2017))。ここで、$h_1,\cdots,h_M$という$M$個の変数は隠れ層と呼ばれるものです。この模型の良いところは、$h_1,\cdots,h_M$という和を解析的に実行できる点でして、
\psi(s_1,\cdots,s_N) = e^{\sum_{j=1}^N a_j s_j} \prod_{i=1}^M 2 \cosh \left( b_i + \sum_j W_{ij} s_j \right)
という$h_1,\cdots,h_M$を含まない形に変形できます。この関数のパラメータは、$a_j,b_i,W_{ij}$です。
RBMで気をつけなければならない点としては、この関数は、パラメータが実数である限り、常に正であることです。もし波動関数が負になり得る場合には、パラメータを複素数にする必要があります。
Flux.jlによるモデリング
実数版
RBMをFlux.jlで実装しましょう。これはオリジナルのレイヤーを作ればよいので、
struct RBM{T1,T2,T3,numspin,numhidden}
a::T1
b::T2
W::T3
end
Flux.trainable(a::RBM) = (a=a.a,b=a.b,W=a.W)
function RBM(numspin,numhidden)
a = randn(numspin)
b = randn(numhidden)
W = randn(numhidden,numspin)
T1 = typeof(a)
T2 = typeof(b)
T3 = typeof(W)
return RBM{T1,T2,T3,numspin,numhidden}(a,b,W)
end
function RBM(a,b,W)
numhidden,numspin = size(W)
T1 = typeof(a)
T2 = typeof(b)
T3 = typeof(W)
return RBM{T1,T2,T3,numspin,numhidden}(a,b,W)
end
function (m::RBM{T1,T2,T3,numspin,numhidden})(x) where {T1,T2,T3,numspin,numhidden}
factor = exp(sum(x.*m.a))
Wx = m.W*x
for i=1:numhidden
factor *= 2*cosh(m.b[i] + Wx[i])
end
return factor
end
Flux.@functor RBM
となります。Flux.jlのモデルとして作りましたので、自動微分が可能であり、
function main()
numspin = 2
x = rand([-1,1],numspin)
rbm = RBM(numspin,6)
θ,re = Flux.destructure(rbm)
grad = Flux.gradient(θ -> re(θ)(x),θ)
println(grad)
end
main()
のように簡単に微分することができます。
複素数版
波動関数が負にもなるようにするためには、パラメータを複素数にする必要があります。上で構築したモデルのパラメータをそのまま複素数にした模型は
struct RBM_c{T1,T2,T3,numspin,numhidden}
a::T1
b::T2
W::T3
end
Flux.trainable(a::RBM_c) = (a=a.a,b=a.b,W=a.W)
function RBM_c(numspin,numhidden)
a = randn(numspin) + im*randn(numspin)
b = randn(numhidden) + im*randn(numhidden)
W = randn(numhidden,numspin) + im *randn(numhidden,numspin)
T1 = typeof(a)
T2 = typeof(b)
T3 = typeof(W)
return RBM_c{T1,T2,T3,numspin,numhidden}(a,b,W)
end
function RBM_c(a,b,W)
numhidden,numspin = size(W)
T1 = typeof(a)
T2 = typeof(b)
T3 = typeof(W)
return RBM_c{T1,T2,T3,numspin,numhidden}(a,b,W)
end
function (m::RBM_c{T1,T2,T3,numspin,numhidden})(x) where {T1,T2,T3,numspin,numhidden}
factor = exp(sum(x.*m.a))
Wx = m.W*x
for i=1:numhidden
factor *= 2*cosh(m.b[i] + Wx[i])
end
return real(factor)
end
Flux.@functor RBM_c
となります。上のモデルとの違いは、パラメータが複素数になっていることと、出力のところで実部を取っているところです。これによって、波動関数の値が負になることができます。
波動関数の値を表示する関数を
function print_wavefunction(rbm,numspin)
Sj = zeros(Int8,numspin)
ψjs = []
Sjs = []
for i=0:2^numspin-1
k = lpad(string(i,base=2),numspin,"0")
#println(k)
for ispin=1:numspin
Sj[ispin] = ifelse(k[ispin]=='0',-1,1)
end
ψj = rbm(Sj)
push!(ψjs,ψj )
push!(Sjs,copy(Sj))
#println("$Sj $ψj ")
end
ψjs /= norm(ψjs)
for i=0:2^numspin-1
Sj = Sjs[i+1]
ψj = ψjs[i+1]
println("$Sj $ψj ")
end
end
と作成しました。
学習
それでは、学習してみましょう。スピン数は2個という一番簡単な問題で試してみます。
using Random
function main()
Random.seed!(12)
numspin = 2
rbm = RBM_c(numspin,6)
value = 1
hx = 0.5
ham = Hamiltonian(numspin)
for i=1:numspin
j = i+1
j += ifelse(j > numspin,-numspin,0)
push!(ham,SziSzjterm(numspin,i,j,value))
push!(ham,Sxiterm(numspin,i,hx))
end
H = make_Hamiltonian(value,hx,numspin)
nt = 3000
θ,re = training_full!(H,rbm,numspin,nt) #小さい系
#θ,re = training!(ham,rbm,numspin,nt) #一般用
print_wavefunction(re(θ),numspin)
e,v = eigen(H)
println("energy $e")
return
end
main()
このコードを実行すると、後半で
2998 energy: -2.2360679774997894
2999 energy: -2.2360679774997894
3000 energy: -2.23606797749979
Int8[-1, -1] -0.16245984811645345
Int8[-1, 1] 0.6881909602355869
Int8[1, -1] 0.6881909602355866
Int8[1, 1] -0.1624598481164532
4×4 Matrix{Float64}:
0.16246 2.22045e-16 -0.707107 0.688191
-0.688191 0.707107 7.85046e-17 0.16246
-0.688191 -0.707107 7.85046e-17 0.16246
0.16246 2.22045e-16 0.707107 0.688191
energy [-2.236067977499788, -2.0000000000000004, 2.0000000000000004, 2.23606797749979]
のような出力が得られます。最後のenergyというのは、ハミルトニアンを対角化した結果です。ちゃんとよくあっています。固有ベクトルもよくあっています。
モンテカルロ法に切り替えてみると、出力は
2999 energy: -2.2360679774997583
--------------------------------
Int8[-1, -1] -0.16245984811645697
Int8[-1, 1] 0.6881909602356704
Int8[1, -1] 0.688190960235506
Int8[1, 1] -0.1624598481164374
acceptance ratio: 0.108
3000 energy: -2.2360679774997734
Int8[-1, -1] -0.16245984811644892
Int8[-1, 1] 0.6881909602356652
Int8[1, -1] 0.6881909602355125
Int8[1, 1] -0.16245984811644001
4×4 Matrix{Float64}:
0.16246 2.22045e-16 -0.707107 0.688191
-0.688191 0.707107 7.85046e-17 0.16246
-0.688191 -0.707107 7.85046e-17 0.16246
0.16246 2.22045e-16 0.707107 0.688191
energy [-2.236067977499788, -2.0000000000000004, 2.0000000000000004, 2.23606797749979]
となり、こちらもうまくいっています。どちらの方法も、乱数次第では最小値に落ちなかったりします。もしかしたらもう少し改善点があるのかもしれません。