NVIDIA PET：再学習なしのテンソルコア高速化、分子動力学ステップの短縮

NVIDIAが原子のエネルギーと力を予測するニューラルネットワークモデルPETを、再学習なしで高速化した結果を公開した。NVIDIAの評価では行列演算だけでなく、原子の位置と速度を更新する分子動力学の1ステップ全体も高速化し、既存モデルでシミュレーション時間を短縮できる可能性を示した。

この計算では小さな力の誤差がステップごとに蓄積し、原子の動きを乱すことがある。行列演算専用のGPU演算部であるテンソルコアを使う場合も、速度と数値的安定性を両立する必要がある理由だ（原文）。

NVIDIAの測定では、B300で大きな行列のBF16x9高速化が顕著だった（Figure 3）。A100でもTF32x3は提示されたすべての形状で同GPUのFP32より速かった（Figure 5）。デバイスが異なるため、両手法の優劣を意味するものではない。

BF16x9のperformant戦略は、cuBLASがより速いと予測した場合のみ使う。NVIDIAの測定ではステップ全体がコンパイル済みFP32比で192原子では1.2倍、17,280原子では1.7倍高速だった。NVIDIAが測定したTF32x3の高速化率は、大規模システムで同GPUのFP32を基準に1.5倍だった（Figure 4・6と説明）。

小規模システムではGPUの処理開始時間や近傍原子処理など、行列演算以外の処理の割合が大きく、利点が小さくなる。TF32x3には行列形状とGPUに合わせたカーネル調整と、モデルの精度および物理的安定性の要件に応じた検証が必要だ（制約と選択基準）。

NVIDIAの精度評価では、BF16x9が生む誤差は同じモデルを実行しても生じる固有の変動より小さかった。NVIDIAの測定では、TF32x3の力の誤差はFP32を再コンパイルして実行した際の差の約3倍だった（精度分析）。

539原子の合金表面の位置を安定した配置に調整する構造緩和には、PET・FIRE2・NVIDIA ALCHEMI Toolkitを使った。NVIDIAの3回の反復評価で、FP32は218ステップ、TF32x3は218~222ステップで力の許容誤差1.0×10⁻³ eV/Åに到達した。NVIDIAの測定でTF32変換を1回だけ行う経路はステップ当たりでは速かったが、1,000ステップでも収束しなかった（Figure 8と説明）。

NVIDIAが白金表面から一酸化炭素分子を引き離して評価した試験でも、TF32x3はFP32の滑らかなエネルギー曲線に追従した。NVIDIAが別途、粒子数・体積・エネルギーを固定した192原子の水の30 psシミュレーションを評価した結果では、FP32と同程度のエネルギー変動を示した。NVIDIAの評価で単一変換のTF32は、それぞれ曲線の揺れとエネルギーの一方向への偏りを示した（Figure 7と説明）。