プログラムの高速化・並列化サービスの事例

概要

本レポートは、HECToR上の2つの著名な密度汎関数法コード、CASTEP[4]およびONESTEP[7]の座標最適化におけるスケーラビリティ改善を目的としたdCSEプロジェクトの結果を報告します。座標最適化はイオン座標に関する原子系エネルギーの最小化です。これは両コードにおいて2番目に良く用いられる機能で、物理的に意味のある結果を出すためには、電子密度分布だけでなくイオン座標に関しても収束した状態でなくてはならないためです。Broyden-Fletcher-Goldfarb-Shanno(BFGS)法は、このような非線形最適化問題を解く手法であり、コードのデフォルト手法です。ここでメモリーのスケーリング性が特に重要で、現状のBFGS法はシステムサイズ$N$に対して$O(N^2)$でスケールしており、座標最適化対象の原子数はHECToRフェーズ2bの32ノードでは2,000に抑えられています。BFGSは、調和的な制限内の超一次線形収束を示す計算過程でヘッシアン行列を構築します。このヘッシアン行列が2次のメモリースケールの要因です。しかしながらBFGSは、ヘッシアン行列を一から作成せず、座標とグラジェント更新ベクトルのみ保存することで計算が可能です。この記憶制限BFGS(LBFGS)法はメモリーを$O(N)$でスケールし、ある条件下での性能改善を示すことが示されます。

目次

イントロダクション

BFGS法

 ニュートン法

 準ニュートン法

 BFGS

 直線探索

 Wolfe条件

記憶制限BFGS

 履歴の制限

 ヘッシアン解析

 弾性特性とフォノンモードの計算

測定

実装

結果と比較

結論

イントロダクション

計算ハードウェアの進歩に伴い、計算科学は利用可能な資源と歩調を合せ、より複雑かつより物理的に正確な系へ探索を拡大しています。物理や化学の領域では、特にDFTに対して当てはまり、その計算は最高速のスーパーコンピューターの限界を常に追求しています；80年代のベクトル型スーパーコンピューターから、90年代の特注かつ様々な分散メモリーマシンを通して、現在の比較的ホモジニアスなアーキテクチャを用いた共有/分散メモリー型の高並列スーパーコンピューターに渡って。DFT計算は、8原子のシリコン単位セルから、数千原子のバルク金属系や万単位の原子数の有機分子に至るまで、幅広く計算対象とします。

これらの計算は配位空間の一点で求められますが、真の収束性のためには、最小化すべき物理量は固定された座標の電子状態エネルギーだけでなく、その座標に対しても最適化が必要です。計算量から言えばこれは極めて膨大な作業量です；作業を座標最適化と電子状態計算に分ければ、電子状態計算の方がより負荷が高く、座標最適化の繰り返しステップでは電子状態の最適化を実行しなければならず、そうした時点でようやく力を計算することが出来ます。

現状のCASTEPおよびONESTEPはBFGSコードを共有しており、それが開発された時点では、DFTを用いて1,000原子より多くの系を扱うことは非現実的でした。BFGSは、イオン座標やエネルギー・グラジェントを用いたランク2の更新による近似ヘッシアンを用いる準ニュートン法です。ヘッシアンは自由度のべき乗の量なので、座標最適化で必要なメモリー量は最大$(6+3N)^2$のオーダーです。これは3次元空間の各イオンの3自由度とセル空間の6自由度(エネルギー不変な回転を除きます)を含みますが、様々な条件を加味して削減されます。また、ONETEPは現在、可変セルの計算を行うことはできません。

大規模計算では、この2次のスケール性を持つヘッシアンは各MPIプロセスにコピーされ、分散配置が不可能なため、すぐにノード上のメモリーを使い切ってしまいます。問題はさらに複合的で、ヘッシアンは良好な収束性を確保するために頻繁に初期化され、初期化プロセスにはPfrommerらが示したようにさらに6つのヘッシアンサイズの行列が必要になります。

このプロジェクトの目標は以下の通りです

1.	BFGSアルゴリズムと実装の理解
2.	CASTEPの現状のBFGSコードの最適化(BLASの有効利用など)
3.	CASTEPへベースとなるL-BFGSアルゴリズムを実装する
4.	CASTEPのSVDコード内へL-BFGSを実装する
5.	ONESTEPへ新しいアルゴリズムを実装する
6.	大規模ベンチマークで新しいアルゴリズムの実行時間とメモリー性能を調査する
7.	最終レポート作成

BFGS法

Broyden,Goldfarb,Fletcher,Shanno(BFGS)法は、非線形最適化問題を解くためのニュートン法に対する近似です。これは、Davidon,Fletcher,Powell(DFP)法のように、初期の準ニュートン法をベースに構築されましたが、より優れた収束特性を持ちます[3]。さらに、最急降下法や共役勾配法(CG)と比較して、準ニュートン法は安定点への直接的なパスを示すことが多く、曲率情報からより良い関数的座標モデルを作ることが出来ます[5]。しかしながら準ニュートン法は、CG法に比べて明らかに不利な性質を持っています。それは、$O(N^2)$オーダーのヘッシアンが必要とするメモリー要求量です。

ニュートン法

二階微分可能な関数$f$は点$\mathbf{x}$において以下のテーラー展開が可能です

\begin{eqnarray} f(\mathbf{x}+\Delta \mathbf{x})= f(\mathbf{x})+J[f(\mathbf{x})]\Delta \mathbf{x} +\frac{1}{2}\Delta \mathbf{x}^T H[f(\mathbf{x})]\Delta \mathbf{x}… \tag{1} \end{eqnarray}

ここで$J$はヤコビアン、$H$はヘッシアンです。
ここで、この関数の$\Delta \mathbf{x}$に関する微分をゼロとすれば、この近似の局所安定点が見つかり、この点において以下の式が成り立ちます

\begin{eqnarray} 0=J[f(\mathbf{x})]+H[f(\mathbf{x})]\Delta \mathbf{x}. \tag{2} \end{eqnarray}

このプロセスは、$f$の安定点が見つかるまで繰り返され、関数は最小値に十分近くでテーラー展開されて、以下の式が成り立ちます、

\begin{eqnarray} \Delta \mathbf{x}= -H[f(\mathbf{x}_n)]^{-1} J[f(\mathbf{x}_n)]. \tag{3} \end{eqnarray}

Wolfe条件を満たすように、ステップサイズは$\alpha$でスケールされなければなりません、

\begin{eqnarray} \mathbf{x}_{n+1}=\mathbf{x}_n-\alpha \Delta \mathbf{x}_n, \tag{4} \end{eqnarray}

そしてこのプロセスは数値的に収束するまで繰り返されます。

準ニュートン法

安定点近傍においては、ニュートン法は良好な収束性と精度を持ちます。その欠点の一つは全ヘッシアンの計算が必要になることです。これは計算負荷が大きく、メモリーは$N^2$オーダーの量が必要です。別のやり方として、ヘッシアンの幾つかの重要な性質を再現しつつ、各繰り返しにおいて低いランクの更新情報からヘッシアンを近似的に構築する方法があります。これは準ニュートン法と呼ばれ、多くのバリエーションがありますが、ヘッシアンの更新のやり方によっては優れた収束特性を持つことが知られています。

BFGS法

ヘッシアンが持つべき最も重要な性質は、セカント法に従う、あるいは座標変化とグラジェント変化間の変換行列として用いることが出来ることです、

\begin{eqnarray} H_{n+1}s_n=y_n, \tag{5} \end{eqnarray}

ここで$s_n,y_n$はそれぞれ座標およびグラジェントの変化量です。
最も成功し、最も広く利用されている更新の式の一つがBFGSです。これは、その安定性およびinexact直線探索を用いた良好な収束性に優れ、近似ヘッシアンは対称性と正定値性を有しています。
更新式は以下のものです：

\begin{eqnarray} H_{n+1}^{-1} =(I-\rho_ky_ks_k^T) H_n^{-1}(I-\rho_ky_ks_k^T)+ \rho_ks_ks_k^T, \tag{6} \end{eqnarray}

ここで、$\rho_k=1/(y_k-T s_k)$です。式(6)から解る通り、BFGSでは逆ヘッシアン行列(以降、これを$B$と記します)を更新します。これが便利であるのは、各繰り返しにおいて、ヘッシアンを最小化の探索方向を見出して、次のポイントを見つけるのに用いる場合です。これは以下のようにします：

\begin{eqnarray} p_n=-B_n \nabla f(x_n), \tag{7} \end{eqnarray}

従って、DFPのような方法が、ヘッシアンを更新してその逆行列を必要とするのに対し、これは逆ヘッシアン行列更新のみで良いことになります。ここで、Sherman-Morrison-Woodbury(SMW)式によれば、行列のランクn更新の逆行列は、行列の逆行列のランクn更新として計算可能です

\begin{eqnarray} (A + BCD)^{-1} = A^{-1}-A^{-1}B(C^{-1}+DA^{-1}B) ^{-1}DA^{-1}. \tag{8} \end{eqnarray}

この式は、BFGSとDFP両方で、それぞれヘッシアンおよび逆ヘッシアン行列へ適用可能です。

直線探索

準ニュートン法はそれほど正確なものでないため、特に更新数が少ない場合、あるいは非調和制限の場合には、$\alpha$(式4を参照してください)を決めるために直線探索がよく用いられます。つまりこれが意味するのは、準ニュートン法は放物型モデルを探索方向の決定に用い、正確な関数がその直線に沿って最小化されることです。直線探索を厳密に実行するのは、最小化の繰返し数が増えてしまうだけなので実用的ではありません。効率性を考えれば、直線探索の精度は準ニュートン法のステップ数とバランスさせて、関数(力)の計算を減らすべきでしょう。

\begin{eqnarray} x_{n+1}=x_n+\alpha_n p_n \tag{9} \end{eqnarray}

Wolfe条件

直線探索を実行する際の考え方は、探索方向での目的関数の最小値を見つけることです。ですがこれを実行するには正確なステップ長$\alpha$を計算しなくてはなりません。つまり探索方向上の数点を見つけて、その点で関数を計算しなくてはなりません。Wolfe条件はステップ長に対する2つの条件です。これは、大域的最小値を保証するものでなく、特定の$\alpha$を用いて関数計算を明確に削減する(小さくない$\alpha$を用いてプロセスを加速する)ことを保証するものです。Wolfe条件の一つは、目的関数を減少させる以下の式です：

\begin{eqnarray} f(x_n+\alpha_n p_n) \leq f(x_n)+w_1 \alpha_n \nabla f_n^T p_n; \tag{10} \end{eqnarray}

また二つ目の条件は、大きなステップサイズの選択です：

\begin{eqnarray} \nabla f(x_n+\alpha_n p_n)^T p_n \geq w_2 \nabla f_n^T p_n. \tag{11} \end{eqnarray}

スカラー係数$w_1,w_2$は任意ですが、計算を通して最初のものは小さく、二番目のものは1より少し小さく(0.8など)すべきです。CASTEPの直線探索では二番目の条件は不要で、ある決まった値から開始するバックトラック直線探索を用いています。

記憶制限BFGS法

BFGS(より一般的には準ニュートン法)のメモリー・スケール性を改善するためには、その更新のやり方を注意深く検証する必要があります。文献[2]に見られるように、再帰方程式[6]は以下の式と同等です：

\begin{eqnarray} B_n=B_0+[S \quad B_0Y]W \left[ \begin{array}{c} S^T \\ (B_0Y)^T \end{array} \right], \tag{12} \\ W_n= \left[ \begin{array}{cc} R_n^{-T}(D_n+Y_n^TB_0Y_n)R_n^{-1} & -R_n^T \\ -R_n^{-1} & 0 \end{array} \right]. \tag{13} \end{eqnarray}

$B_0$は初期逆ヘッシアン行列で、

\begin{eqnarray} R_{n+1}^{-1}= \left[ \begin{array}{cc} R_n^{-1} & -\rho R_n^{-1}S_n^T y_n \\ 0 & \rho_n \end{array} \right], \tag{14} \end{eqnarray}

\begin{eqnarray} D_{n+1}= \left[ \begin{array}{cc} D_n & 0 \\ 0 & y_n^T s_n \end{array} \right], \tag{15} \end{eqnarray}

\begin{eqnarray} \rho_n=\frac{1}{ y_n^T s_n }, \tag{16} \end{eqnarray}

\begin{eqnarray} S_n=[s_0,…,s_{n-1}], Y_n=[y_0,…,y_{n-1}]. \tag{17} \end{eqnarray}

ここで、$s_n,y_n$はそれぞれn番目の座標およびグラジェントの更新ベクトルです。これは式(6)よりかなり複雑な表現となっていますが、その利点として、逆ヘッシアン行列に施す演算はベクトルを一度掛けるのみで良く、また、座標とグラジェントの更新ベクトルの保存のみで可能になります。n番目の計算には以前のステップと重複が多いため、$R_n^{-1}$は保存されて、各ステップにおいて式(14)と同じやり方で更新されます。注意深く操作すれば、$W_n$を保存せずに$m^2$サイズの行列(ここで、mは保存する更新ペア数、あるいはBFGS繰り返し数)、$R_n$と$(T_n=D_n+Y_n^TB_0Y_n)$、およびm*nサイズの$S_n,Y_n$の保存のみで行うことが出来ます。

BFGSにおいては、ヘッシアンは力のベクトルの掛けることによって探索方向を与えるものですが、LBFGSではヘッシアンは一切、計算あるいは保存されません。このやり方をメモリーに対して線形にスケールさせるためには、その表現を展開せずに計算を実行すべきです。こうしたアルゴリズムを以下に示します。また計算に必要となるBLASライブラリールーチン名を示しています。$v$は逆ヘッシアン行列で乗算されるベクトル、$p$は計算結果のベクトルです：

Function $f:\mathbb{R}^n \rightarrow \mathbb{R}^n,p \mapsto B_nv$
\begin{eqnarray} p & = & B_0v & DSBMV & \\ w_{1:m} & = & Y_n^Tp & DGEMV & \\ w_{m+1:2m} & = & S_n^Tv & DGEMV & \\ w_{1:m} & \mapsto & R_n^{-1}Y_n^TB_0v & DTRSM & \\ w_{m+1:2m} & \mapsto & R_n^{-1}S_n^Tv & DTRSM & \\ \Xi & = & T_n & & \\ \Xi & \mapsto & R_n^{-1} \Xi & DTRSM & \\ w_{1:m} & \mapsto & w_{1:m}-\Xi w_{m+1:2m} & DGEMV & \\ p & \mapsto & p+S_n^T w_{1:m} & DGEMV & \\ t & = & -Y_n^T w_{m+1:2m} & DGEMV & \\ p & \mapsto & p+B_0t & DSBMV & \end{eqnarray} End Function

ここで$w$は2mの長さの作業用ベクトルで、$\Xi$はm*mサイズの作業用行列、$t$は長さnの作業用ベクトルです。この関数を用いて、BFGSの代わりにLBFGSを実行することが出来ます。この他に幾つかの小さな行列も必要です。

履歴の制限

LBFGSがBFGSと最も異なる性質は、更新ペアの構築でなく、一つの更新ペアを移動することで逆ヘッシアン行列近似を更新する点です。実際には、k個の更新ベクトルを用いて実行しており、これがLBFGSのメモリーに対する線形なスケール性を意味しています。もし更新ベクトルの数に制限がなければ、BFGSのようなメモリーに対する要求が生じるでしょう。
その履歴長を制限するためには、最も古い更新ベクトルを削除して、効果的に保存領域をシフトする機能が必要です。

図1

図2：行列Rに対して更新ベクトルを削除して追加する。削除は決まった履歴長と履歴ベクトルが飽和した場合に必要となる。

原子座標最適化においてLBFGSを用いる場合、さらに複雑な事情があります。

ヘッシアン解析

LBFGSを用いる理由はメモリー依存性を線形にすることですが、再初期化や最終的な解析(リスタートで用いられる)はこのようにすることが困難です。そこで弾性特性とフォノン周波数(ブリユアン領域の中心での)解析へ応用する方法を開発し、以下のセクションで詳細を示します。プロジェクトの最初では、特異値分解(SVD)は、更新ベクトル空間の基底を見出すのに用いられるため、その項目として述べる予定でしたが、LBFGS実装においても組み込まれたため、ここで示すこととします。

これはPfrommerら[6]による逆ヘッシアン行列の初期化に全てを負っています。これは以下のように初期逆ヘッシアン行列がセットアップされるという前提のもとに計算されます：

これはブロック対角ですが、上左が9*9ブロックで、残りは3*3ブロックの対角形です。$\Omega$はセル体積、$B_0$はブロックモジュール、$\tilde{M}$は平均イオン質量、$\omega_0^{-2}$はブリユアン領域中心での平均フォノン質量、行列$g_0$は3*3メトリックテンソルです。

上左ブロックはセル間相互作用、下右ブロックはイオン間相互作用を表します。これは原子間相互作用やセル・イオン相互作用による音響モードは含みません。 対角上の3*3ブロックは対称で、行列全体を以下のように対称な5重対角帯行列として保存可能です。

ここで、黒四角は初期逆ヘッシアン行列のブロック要素、白四角はゼロ、＊は無関係な要素を表します。

つまり初期逆ヘッシアン行列の保存領域は線形スケールします。初期ヘッシアン行列で必要な変数はこれだけで、後はその構成と、逆ヘッシアン行列の乗算ルーチンの呼出に関連するスカラー値のみです。しかしながら。将来的により不均一な初期化が必要になる事から、この手法は用いません。

初期ヘッシアン行列にベクトルを掛ける場合は常に、BlasルーチンDGBMVを用いますが、一般的な行列に$B_0$を掛ける場合にはBLASやLAPACKに適切なルーチンが存在しないため、新たに専用ルーチンを作成しました。5重対角行列では上の対角バンド幅は2となり、このルーチンにはこの制限が課されます。

現状のCASTEPにおけるフォノンモードと弾性特性計算は、Pfrommerのやり方に従っています。この計算を開始する前に、Pfrommerのプロセスが力のベクトルを増強するやり方で、$B_n$は有限の歪が入った行列で補正されます：

\begin{eqnarray} B_n'=D^{-T}B_nD^{-1}, \tag{20} \end{eqnarray}

\begin{eqnarray} D= \left[ \begin{array}{cccccc} (1+\epsilon) & & & & & 0 \\ & (1+\epsilon) & & & & \\ & & (1+\epsilon) & & & \\ & & & 1 \\ & & & & … & \\ 0 & & & & & 1 \end{array} \right], \tag{21} \end{eqnarray}

CASTEPにおいては、この行列は逆行列計算前に、サイズ$ndim^2$分すべてがアロケーションされます。この行列は、上左部でブロック対角で残りは1の対角形です。3*3ブロックには$\epsilon_f=(1+\epsilon)$が保存されており、$B$はそのセル間相互作用部に、3*3ブロックおよびその逆行列が掛けられて、以下のようになります：

このルーチンでは次に、ヘッシアンのサンプリング・サブ空間の基底を得るために、更新空間のヘッシアンに対してPfrommerが提案したSVDを用います。初期逆ヘッシアンは最新のものであることから、これが変更の全てになります。これはCASTEPのBFGSが実施する方法です：

\begin{eqnarray} B'=V \sigma W \tag{23} \end{eqnarray}

ここで$B'$は逆ヘッシアン行列に対する更新ベクトルで、$B_0$の最新の初期化となります。$V$は$B'$の特異ベクトルの直交行列、$\sigma$は特異値の対角行列、$W$は零空間の$B'$の特異ベクトルの直交行列です。LAPACKは$N^2$行列Vの計算に用いられています。特異値はN要素のベクトルへ纏められ、Wは計算されません。

式(12)の展開をせずにSVDを実行して、少なくともこれと他の$N^2$の量を含めて、$B'$のLBFGS表現の構造が詳細に検証されました。

\begin{eqnarray} B'=[S \quad B_0Y]W \left[ \begin{array}{c} S^T \\ (B_0Y)^T \end{array} \right], \tag{24} \\ \end{eqnarray}

更新行列は直交行列である必要はないですが、QR分解を用いれば可能です。この行列演算は展開する必要はありませんが、原理的には以下のようになります：

\begin{eqnarray} B'=QRWR^TQ^T, \tag{25} \\ \end{eqnarray}

$RWR^T$はスペクトル分解できて、

\begin{eqnarray} B'=QP \Lambda P^TQ^T, \tag{26} \\ \end{eqnarray}

$\Lambda$は固有値行列、$P$は固有ベクトルの直交行列です。$P,Q$は共に直交行列なので、その積も直交行列となり、結局、直接的な計算なしで逆ヘッシアン行列の更新空間の基底を抽出したことになります。ここで、$\Lambda$内の非零特異値に対応するQPの要素のみを保存すれば良く、この最大N*Mの行列を$Z$を呼ぶことにします。

Pfrommerの論文において述べられているのは、基底$Z$は純粋なセルのみとイオンのみのベクトルから構成されるという点です。元々のCASTEPでの解釈の論理は、$B'$へSVDを実施して、ハウスホルダー変換によりブロック対角へ変換するというものです。

\begin{eqnarray} B'= \left( \begin{array}{c|c} Cell-cell & 0 \\ \hline 0 & ion-ion \end{array} \right), \tag{27} \\ \end{eqnarray}

しかしこの作業には$B'$のコピーが必要になります。同じことですがより効果的な方法は、論文が意図しているかどうか不明ですが(ハウスホルダー変換は明確に言及されていませんが、暗示されています)、cell-cellおよびion-ion部分を別々に基底ベクトルを抽出することです。その後これらベクトルは一つの基底行列へ纏められ、形式的な手法から得られるベクトルとは異なるものですが、依然として完全なサンプル空間を表現します。

図3：基底ベクトル$Z$を線形メモリー量で抽出する方法。番号付き矢印はこの方法の5つのステージを表す。1.更新行列をQR分解する。2.行列$RWR^T$を構成する。3.この行列を固有値分解する。4.非零特異値に対応する固有ベクトルを抽出する、5.これらのベクトルを更新行列へ掛けて$Z$を生成する。

通常はセルベクトルは更新ベクトルよりも少ないため、$B'$のcell-cellブロック行列を構成して、この9*9行列のSVDを、残りのion-ionブロックにQR/SVDを行う前に実行する方が理にかなっています。

次にルーチンから消去すべき$N^2$行列は、質量スケーリング行列で、以下の対角行列です：

\begin{eqnarray} \mu= \left[ \begin{array}{ccccccc} I_6 & & & & & 0 \\ 　　　& M_1　& & & & \\ 　　　& & M_1　& & & \\ 　　　& & & M_1　& & \\ 　　　& & & & M_2　& \\ 　　　& & & &　… & \\ 0　　 & & & & & M_N \end{array} \right], \tag{28} \\ \end{eqnarray}

これは単純にベクトルとして保存でき、各要素は$Z$の対応する列が掛けられ、Pfrommerに名づけられた縮小された座標でのヘッシアン行列は次の形式を取ります：

\begin{eqnarray} \bar{A}=(Z^T \mu^{1/2} B_{\mu}^{1/2} \mu^{1/2} Z)^{-1}. \tag{29} \\ \end{eqnarray}

ただし、PfrommerはLBFGSのグラジェント更新において、$Z$の代わりに$Y$を用いていました。

$Z$はN*N空間からより少ないM*M次元の縮小された座標空間への線形写像として働きますが(したがって、$\bar{A}$およびその構築に含まれる乗算は線形スケールのメモリー量で可能です)、行列$B_0$に対するランクnの補正に際して、以前に指摘したWoodbury式を用いることが出来ます。

\begin{eqnarray} \bar{A}=\left( \bar{Z}^T \left( B_0+[S \quad B_0Y]W \left[ \begin{array}{c} S^T \\ (B_0Y)^T \end{array} \right] \right) \bar{Z} \right)^{-1}. \tag{30} \\ \end{eqnarray}

ここで$\bar{Z}=\mu^{1/2}Z$は、質量でスケールされた基底です。

\begin{eqnarray} \left( \begin{array}{c|c} Cell-cell & cell-ion \\ \hline Ion-cell & ion-ion \end{array} \right), \tag{31} \\ \end{eqnarray}

弾性特性とフォノンモードの計算

この時点では全ての$O(N^2)$オーダーの行列は削除されていますが、この場合はそうではありません。元々のコードは、オーバーフローを避けるために$N^2$の行列で$\bar{A}$を保持しています。固定の履歴長を用いればこの問題を避けることが出来て、階数・退化次数の定理(rank-nullity theorem)により$\bar{A} \leq m$が保証されます。

もう一つの問題は、スケール性には関係ありませんが、LBFGSのPfrommerの手法の実装の観点から重要な事です。それはPfrommerのやり方が$B_n$のcell-cellセグメントをVoigt空間へ入れていることです。これは、セル系の対称性は課されますが、行列の対称性は失われます。$B_n$のサイズを3倍小さくするためには、LBFGSの対称積とVoigtセルの間にはなんら関係が無いため、これは良いアイデアとは言えませんでした。

6*6セルから9*9セルへの入れ替えにおいて、ストライドの課題があります。9*9はランク4テンソルを表すため、体積弾性率に総和される要素はVoigt記法で、ランク2の対称歪テンソルの1,2,3ではなく1,5,9です。
また、フォノン振動数計算においてコレスキー分解を効率性改善に利用できることが明らかになりました。

測定

新旧のコードのメモリーサイズ測定には、/proc配下の仮想ファイルから、座標最適化コードの呼出し前後でピークのメモリー使用量を観測しました。 ファイル/proc/self/statusのVmPeakとVmExeで始まる行を読み込み、計算中の特定の時点でのピークメモリーサイズを計測しました：利用する全メモリー=VmPeak-VmExe。 ここでVmExeは実行形式ファイルのサイズで、この測定とは関係ありません(LBFGS版のバイナリーはそれがない場合よりも大きくなるため、多少の誤差は混入します)。

このようにして、全コードのメモリー使用を把握できるようになり、他の全てのコードの座標最適化コードのメモリーを測定するには、 親の座標最適化ルーチン(geometry_optimise)の最初と最後で測定すればよいことになります：座標最適化メモリー=最後のメモリー-最初のメモリー。

このコードに掛かる時間の測定も同じやり方です。ただし仮想ファイルでなくMPI経過時間ルーチンを用います(MPI_Wtime)。 この関数の評価に掛かる時間は、座標最適化モジュールに掛かる時間には含まれないように差し引かれます。 この関数の実行は多数発生します。よってLBFGSが多く実行されれば、BFGS実行に掛かる総時間へ影響を与えます。

測定における人為性は、系が少し揺らぐと、配位空間内の明確な最小値に異なるルートを取ることです。 このため、直接比較できるのは何かについて注意が必要です。 異なるLBFGS履歴長は必然的に異なるパスへ導き、ヘッシアンは空間の各点において異なる曲率をもたらし、異なる探索方向へ導きます。
メモリーの比較は理にかなったことですが、固定の履歴長に対しては、収束へ至るパスに独立です。

実装

実装において本プロジェクトの目標であるメモリースケーリングの改善に注力されましたが、性能を維持するためにBLASやLINPACKを可能な限り利用しました。実際に、これらライブラリーを用いない重要な操作は"geom_init_H_mul"のみで、これは対称かつ5重対角バンド行列に一般行列を掛けるという特殊なルーチンです。このルーチンはコード中で3回呼ばれるのみです。

最初の標的は、Fortran90組込み行列演算関数(matmul)で、これらは同等のBLASで置き換えられました。対称行列であるヘッシアン(およびその逆行列)はDSYMMとDSYMVで置き換えられました。これら初期の置き換えは小さいですが性能に大きな影響を与えるものでした。

図4：BFGSへBLASを組み込んだ結果、性能は向上したがメモリーは増加した。

次に行ったのは、CASTEPに基本LBFGSを実装することです。これはByrdらのやり方に従って、迅速に行われました。 元々のコードは、例えば制約条件のセクションで、式(5)のヘッシアンを使って座標から力へ変換する、 あるいは逆ヘッシアン行列を使って力から座標へ変換するために、組込み関数matmulを使用していました。

さらに元のコードの重大な特徴が、ヘッシアン行列の構築に在りました。メインループ内の各繰り返しにおいて座標及び力のベクトルの変化は、 式(6)のランク2更新行列生成時に用いられますが、これはBFGS更新として知られるもので、保存されたヘッシアン行列に適用されます。これは全て一つのルーチン内で処理されます。

基本LBFGSの実装において、ヘッシアン行列にベクトルを掛ける形式的なルーチンを一つ開発することとしました。 このルーチンは、必要となれば常にmatmulの代わりに呼び出すことが出来て、BFGSが呼び出されればいつでもBLASのDSYMVのラッパーとして実行することが出来ます。

このルーチンでどちらの手法を用いるかを指定する、do_bfgsおよびdo_lbfgsという2つの相互排他的な論理変数をグローバル宣言し、さらにBFGSとLBFGSの機能を指定するための他の変数も追加されました。 式(2)で用いる行列・ベクトル乗算ルーチンを実装し、これは式(6)で示される演算を実行するもので、最も効率的になるように中間行列$W$の2*2ブロック対角(および、1零ブロック)構造を考慮したものです。

すでにBFGSルーチンは存在しているため、このルーチンを書き直すのでなく、 何時完全なヘッシアンが出来るかと何時更新ベクトル用のメモリーが利用可能になるかを指定するグローバルフラグを用意して、 LBFGS用に拡張しました。これは、n*m次元の更新履歴行列に、最新の位置と力の更新ベクトルを追加することを意味します。

効率性のために、式(14)の行列$R$も保存されますが、この行列はm*m次元しかなく、m*n次元の座標更新行列$S^T$による$R$へのレベル3BLASの乗算、 さらにレベル2BLASの、最新のグラジェント更新ベクトルによる乗算で、行列の最右列を更新します。最後にこの列の対角成分をレベル1BLASで計算して、$R$は新しい行を持つ右三角構造になります。

BFGS更新ルーチンの計算効率性は$O(N^2)$であり、LBFGSは$O(M^2N)$です。 以上により各々に対していくらかの性能向上がありましたが、BFGS更新はスカラー演算で処理されており、新しいベクトル化されたLBFGSはさらに改善すべきでしょう。

この手法のスケーリング性能全体に対して考慮すべき点が少なからず存在します。それはBFGSでなく、(上で述べたような)ベクトルに対する行列の適用です。 なぜなら多数の演算が更新計算からそこに移されているためです(BFGSでは$O(N^2)$の行列・ベクトル演算一つのみです)。アルゴリズムからこの演算数は、$2NM^2+2MN+3M^3+2M^2+6N$となります。 これは保存されたベクトルの数に大きく依存しています。BFGSの演算数プリファクターは、元のコードでのループ内のスカラー演算数を数え上げたところ、10でした。

よってLBFGS演算数に対して10$N^2$をプロットすれば、理論的な性能スケーリング比較が出来ます。LBFGSはメモリーはリニアですが、そのプレファクターはより大きなものです。 結果として繰り返し毎の性能の観点では、大規模系で、比較的短い履歴長の場合に有利です(図5を参照してください)。

勿論固定履歴長でBFGSとLBFGS間に見られる収束回数の違いもく考慮しなくてはなりません。これは総実行時間に大きくかかわります。 何故なら各ステップでエネルギーと力の計算が必要になり、これらが座標最適化と比べて膨大な計算だからです。

図5：自由度数の変化に対する、オリジナルBFGSとLBFGSの性能予測。この自由度は原子数でなく、位置および力のベクトル長である。これは演算数とコードから導いたプレファクターに基礎をおいている。y軸のスケールに注意：これらは現代的なCPUで負担になる数ではない。

ヘッシアン解析コードには、大きな改善余地があります。やるべきことは、6$N^2$の行列保存を線形なメモリースケールにすることですが、実装の仕方によって演算数の削減による大きな性能利得を得ることが可能です。
まず、基本的にはヘッシアンの更新を図3のように計算しなくてはならず、このアルゴリズム詳細を以下に示します：

Function $f:\mathbb{R}^n \rightarrow \mathbb{R}^n,Z \mapsto basis(B')$
\begin{eqnarray} U_{1:n \times 1:2m} & = & [S_n \ \ B_0Y_n] & & \\ [Q_{1:n \times 1:2m} \ \ R_{1:2m \times 1:2m}] & = & qr(U) & DGEQRF & \\ C_{1:2m \times 1:2m} & = & RW & DTRMM & \\ C & \mapsto & CR & DTRMM & \\ [P_{1:2m \times 1:2m} \ \ \Lambda_{1:2m}] & = & spectral(C) & DSYEVD & \\ v & = & pivot(\Lambda) & algor \underline{ } sort & \\ \forall i \leq 2m:\Lambda(i) & \mapsto & \Lambda(v(i)) & & \\ \forall i \leq 2m \land \Lambda(i) \gt 0:P_{1:n}(i) & \mapsto & P_{1:n}(v(i)) & & \\ Z_{1:n \times 1:l} & = & U_{1:n \times 1:2m} P_{1:2m \times 1:l} & DGEMM & \tag{32} \end{eqnarray} End Function

ここで$qr$はQR分解を、$spectral$は固有値分解を示し、$pivot$は、ベクトル要素の大きさの降順で並べ替えて、整数ピボットベクトルを生成することを示しています。並べ替えとピボットベクトル生成は、CASTEPとONESTEPでは別々に実行されます。CASTEPではalgoモジュールのalgo_sortルーチンを用いますが、ONESTEPでは同じ機能のutilsモジュールのutils_heapsortを用います。

更新ベクトル空間の基底を展開した後、完全なヘッシアンに適用して、縮小された座標でのヘッシアン$\bar{A}$を生成します。ここで階数・退化次数の定理より$i \leq 2m$です。このやり方には2つあります。単純なやり方は式(30)を展開して逆行列計算することですが、うまいやり方はWoodbury式(8)を使って、一連の行列のランク更新の逆行列を得ることです。いずれにせよ、初期逆ヘッシアン行列($B_0$)およびその更新行列($B'$)の基底変換という、2つの部分に分けられます。これをWoodbury式を用いて表現しますが、最初の部分では、一般の密な$n \times l$行列と対称5重対角バンド行列の乗算が存在します。このような演算のためのBLASやLAPACKルーチンは、図6で述べたように存在しません。特に保存されたバンド行列の構造は(19)で示したものであることに注意が必要です。

全表式に対して、Woodbury式は直截的に、一般的な行列演算(DGEMM)と、上記のルーチンおよびLAPACKの一般逆行列ルーチン(DGETRI)で実装可能です。

コードには更にリスタートのための変更が含まれます。これは全ての正しい配列が割り当てられて、初期化されたことを確認するものです。これは固定長の履歴を用いる通常の場合に使うことはありませんが、古いBFGSとの互換性を確かめるための履歴長を制限しない場合に用います。これは、昔のサイズでテンポラリー配列を割当て、データをこれにコピーしてその配列を削除し、新しいサイズで配列を割当てて、テンポラリー配列からデータをコピーしてからテンポラリー配列を削除するという単純な操作です。この作業は、引数によって異なる大きさを持つ各配列に対して行います。

ONESTEPはCASTEPのコードをベースにしているため、多少の追加修正がありますが、LBFGSとヘッシアン解析は全て共有します。唯一の違いはチェックポイント・リスタートで、そこではONESTEPは座標最適化のCASTEPモデルのモックを用いており、例えば両者の間で逆行列やソートに異なるルーチンを用いていることから、多少の修正が必要となります。

図6：一般の行列Aに、5重対角行列Bを乗算して、行列Cへその結果を移す場合に、このようなステンシルを用いる。ここで、a,b,c,d,eはBの要素である。対称の要素は乗算され、加算されて、Cの対応する要素へ加算される。この操作はiとj全てに対して行われる。ただしエッジの場合には注意が必要である。ドットはAとCの現在のiとjを示し、Bの現在の列はAに対応して並べられている。

結果と比較

CASTEPとONESTEP共に、様々な原子数に対して、元のBFGSと新しいLBFGSコードを実行しました。CASTEPとONESTEPは全く同じ座標最適化コードを持ち、エネルギーおよび力の計算のみ異なります。イオン座標にノイズを入れた水晶の様々なスーパーセルとセル自由度についてテストをしました。これはCASTEPとONESTEPに対する座標最適化のレグレッションテストで、LBFGS前後の結果が同じであるかを確かめるものです。レポートの採用に述べたようなメモリーサイズ履歴も測定しました。また、幾つかのLBFGSの履歴長に対して、3回以上実行しました。また、物理情報(体積弾性率やフォノン周波数)計算での性能に焦点を当てて、その計算の有り無しでも実行しました。

図7に示すように、LBFGSはBFGSの2次のスケーリングを比較して、良好なスケール性を示しました。履歴長100のLBFGS(一番上のライン)は、予想通りにBFGSのラインと100原子近傍で交わりました。

グラフ上の最大原子数において、100履歴長のLBFGSはほぼ10MBのメモリーを必要としますが、BFGSは300MB以上です。

図7：ヘッシアンを用いた物理情報計算無しの場合。+記号のBFGSは2次スケール性を示し、×記号のLBFGSとは対照的である(線の下から上に、履歴長15,30,45,60,100を示す)。LBFGSは直線フィッティングされ、BFGSは2次のフィッティングが当てはまる。対数―対数グラフであることに注意。

同様の特性は、弾性特性やフォノンモード計算を行なった場合も示され、BFGSとLBFGS間で2次からリニアーに改善することが示されました。メモリー使用率はさらに差が広がりました。例えば最大の2,300原子の場合、同じく100履歴長で、LBFGSは60MB以下でしたが、BFGSは2.5GBに達しました。

図8：ヘッシアンを用いた物理情報計算有りの場合。+記号のBFGSの2次スケール性は、×記号のLBFGS(線の下から上に、履歴長15,30,45,60,100を示す)と比較してかなりの性能劣化を示した。LBFGSは直線フィッティングされ、BFGSは2次のフィッティングが当てはまる。対数―対数グラフであることに注意。

元々のコードでこの点に着目すると、$N^2$配列はMPIプロセスにコピーされ、HECToRフェーズ2BはCPU当り24コアを持つことから、各ノードが飽和した最悪の場合は、座標最適化計算には2.5*24=60GBが必要となります。これはノード当たりのメモリー32GBを超えてしまいます。またこの他にも座標最適化モジュール以外に、波動関数計算やOSからのオーバーヘッドもあります。こうしてHECToR上での数千原子に対する、ヘッシアン解析による座標最適化は、ノード飽和率が大きく減少してノード時間を浪費します。

CASTEPによる新しい計算が可能になり、数千原子の系の座標最適化が現実的なものとなりました。15履歴長のLBFGSは600,000原子に対して2.5GBで計算可能ですが、こうした短い履歴長の場合は、BFGSの収束速度が遅いことが解りました(ですが収束はします)。

履歴長30以上は良好な結果を示しました。特に30は良い選択です。

収束時間の比較については、更新ベクトルの保存数を変えるとヘッシアンは探索方向と収束パスを変更するので、公平な比較が出来ません。その代わりにヘッシアン計算に固定履歴長で測定を実施しました。これは、ヘッシアン解析までは同じパスを描くことを意味し、ヘッシアン計算を公平に比較することが出来ます。

図9：固定履歴長のヘッシアン解析では、LBFGSは小さな系に対してはBFGSと同程度の性能を示し、系のサイズが大きくなるにしたがって改善する。最も良くフィットする$y=ax^b$の形式の直線を描いてある。これは、ヘッシアン解析が主要を占める大きな系の振る舞いを良く近似している。小さな系では、更新時のオーバーヘッドによりデータは直線からずれている。

図10：ヘッシアン解析部の前(あるいはそれをスキップするように指示した場合)では、新しい手法はBFGSにとてもよく似ているが、全体的に僅かに速い。ただし、これはBLAS化したBFGSであり、元々のコードはこれよりも遅いことに注意(図4を参照のこと)。

文献

[1]	Aarons, J., Hasnip, P., and Probert, M. An improved hessian initialisation in the lbfgs formalism. In preparation (2011).
[2]	Byrd, R., Nocedal, J., and Schnabel, R. Representations of quasi-newton matrices and their　use in limited memory methods. Mathematical Programming 63, 1 (1994), 129-156.
[3]	Byrd, R., Nocedal, J., and Yuan, Y. Global convergence of a class of quasi-newton methods on　convex problems. SIAM Journal on Numerical Analysis 24, 5 (1987), 1171-1190.
[4]	Clark, S. J., Segall, M. D., Pickard, C. J., Hasnip, P. J., Probert, M. I. J., Refson, K., and Payne, M. C. First principles methods using castep. Zeitschrift für Kristallographie 220, 5-6-2005 (May 2005), 567-570.
[5]	Nazareth, L. A relationship between the bfgs and conjugate gradient algorithms and its implications for new algorithms. SIAM Journal on Numerical Analysis 16, 5 (1979), 794-800.
[6]	Pfrommer, B. G., Côté, M., Louie, S. G., and Cohen, M. L. Relaxation of crystals with the quasi-newton method. Journal of Computational Physics 131, 1 (1997), 233 - 240.
[7]	Skylaris, C.-K., Haynes, P. D., Mostofi, A. A., and Payne, M. C. Introducing onetep: Linear-scaling density functional simulations on parallel computers. 084119.