Fine-tuning MLIP foundation models: strategies for accuracy and transferability

本論文は、多様な化学ベンチマークにおける機械学習型原子間ポテンシャル（MLIP）基盤モデルに対する7つのファインチューニング戦略を評価し、基盤モデルの品質や正しいエネルギー初期化といった前提条件が極めて重要である一方で、単一系への精度においてはナイーブなファインチューニングが最適であるのに対し、マルチヘッド・リプレイは広範な展開に向けた分布外（OOD）への堅牢性を特異に維持することを明らかにしている。

要約

想像してみてください。あなたは、岩石や金属、塩といった無機質な材料のみを使って完璧な料理を作る技術を、長年磨き続けてきたマスターシェフを持っています。このシェフは「基盤モデル（Foundation Model）」です。そして今、あなたはこのシェフに、非常に少ない新しいレシピを用いて、繊細な有機物のスープや生物学的なシチューのような、全く新しい料理を作らせようとしています。

ここで大きな疑問が生じます。既存のスキルを損なったり、これまでの料理の仕方を忘れさせたりすることなく、どのようにしてこの新しい料理をシェフに教えればよいのでしょうか？

この論文は、このマスターシェフに新しいレシピを教える（微調整する／ファインチューニングする）ための7つの異なる方法をテストした、大規模なキッチン実験です。研究者たちは、教え方の「手法」そのものよりも、3つの決定的な「調理前」のステップ――適切なシェフ選び、適切な基準値の設定、そして火加減の調整――が重要であることを突き止めました。

以下に、彼らの発見を分かりやすく解説します。

1. 3つの「事前チェック」（最も重要な部分）

新しいレシピを教え始める前に、3つのことを正しく行わなければなりません。もしここで失敗すれば、どんな教え方を用いても救いようはありません。

• 適切なシェフを選ぶ（基盤モデルの品質）:

  • 比喩: スフレの作り方を教えるために、お湯を沸かすことしか知らないシェフを雇うことはできないでしょう。
  • 発見: 元となるモデルの品質は、微調整の戦略よりも重要です。多様で膨大な無機材料のデータセット（「OMat24」のようなモデル）で訓練されたモデルは、新しく奇妙な化学反応を学ぶ能力が、古い小さなモデルよりもはるかに高いのです。たとえ同じ教え方を用いたとしても、「より優れた」基盤モデルの方が、常に優れた最終的な料理を生み出します。

• ゼロ地点を設定する（原子参照エネルギー / $E_0$）:

  • 比喩: 建物の高さを測る場面を想像してください。もし地上ではなく地下室から測り始めてしまったら、数値は狂い、建物が浮いているように見えたり、逆に埋まっているように見えたりします。化学においても、個々の原子の「重み」を差し引いて、モデルが「相互作用」についてのみを学習できるようにする必要があります。
  • 発見: 研究者たちは、このゼロ地点を設定するために、スマートで「モデルを意識した」方法を用いることが極めて重要であることを見出しました。もし適当な平均値で済ませてしまうと、モデルは不安定になります。見た目の数値（エラースコア）は良くても、現実世界の物理現象（例えば、風洞実験での建物の崩落など）をシミュレーションしようとした瞬間に、破綻してしまうのです。

• 火加減を調節する（ハイパーパラメータ）:

  • 比喩: 新しいスキルを習得するとき、転んでしまうほど速すぎてはいけませんし、かといって、いつまでも終わらないほど遅すぎてもいけません。
  • 発見: 教える手法によって、必要な「学習率」は異なります。例えば、モデルのごく一部のみを変更するLoRAという手法は、非常に速い学習率を扱うことができますが、2つのことを同時に教える手法には、非常にゆっくりとした穏やかなペースが必要です。

2. 7つの教育戦略

3つのチェックをパスした後、研究者たちは新しいレシピを教える7つの方法をテストしました。

1. ナイーブ・ファインチューニング（素朴な微調整）: 「ただ料理を続ける」。シェフの全知識を使い、新しいデータに対して訓練を続けます。
   • 結果: 特定の料理一つを完璧に学ぶには最適です。しかし、後で別の種類の料理を作ろうとしたとき、シェフが以前のスキルを忘れてしまっている可能性があります（これは「破滅的忘却」と呼ばれる問題です）。
2. レイヤー・フリージング（層の凍結）: 「基本は触らない」。シェフの基本的な包丁技術の知識はロックしておき、新しいソースの作り方だけを学ばせます。
   • 結果: 有用ですが、時に硬直的すぎます。新しい食材に適応する能力を制限してしまいます。
3. LoRA (Low-Rank Adaptation): 「カンニングペーパーを加える」。レシピ本全体を書き換える代わりに、新しいルールだけを記した小さくて効率的なメモ帳を、シェフのエプロンに付け加えます。
   • 結果: 特定のタスクに対して非常に効率的かつ正確であり、ナイーブな微調整と同様の結果をもたらします。
4. マルチヘッド・リプレイ（多頭再生）: 「二つの帽子を被ったシェフ」。シェフに2つの帽子を与えます。一つは新しい料理用、もう一つは馴染みのある古い料理用です。彼らは両方を同時に練習します。
   • 結果: これが「安全性」において勝者です。 以前のスキルを忘れさせないことを一貫して防げる唯一の方法です。新しい料理にも、古い料理にも高いレベルを維持させることができます。
5. 疑似ラベル・リプレイ（擬似ラベル再生）: 「合成シェフ」。実際の古いレシピを使う代わりに、シェフ自身が予測した「古いレシピの記憶」を使って練習します。
   • 結果: 本来の古いデータを用意する必要がないため、柔軟性が高く、うまく機能します。
6. リプレイ ＋ LoRA: カンニングペーパーと二つの帽子を組み合わせた方法。
   • 結果: 良好ですが、「二つの帽子（マルチヘッド）」単体でも十分な性能を発揮することが多々ありました。

3. 大きなまとめ

• 車輪の再発明はしない: もし、特定の狭いタスク（例：塩水のシミュレーションのみ）のためのモデルが必要なら、ナイーブ・ファインチューニングが最も速く、簡単な方法です。
• 過去を忘れない: もし、元の訓練内容を忘れずに、奇妙で新しい状況（例：新しいタイプのバッテリーや複雑な生物学的分子）にも対応できるモデルが必要なら、必ずマルチヘッド・リプレイを使用してください。これこそが、モデルを堅牢に保ち、「忘却」から守る唯一の戦略です。
• テクニックよりも品質: この論文は、完璧な微調整アルゴリズムを選ぶことよりも、高品質な基盤モデルを選び、エネルギーの参照値を正しく設定することに時間を割く方が重要であると強調しています。もし土台が弱い、あるいは数学的な設定が間違っていれば、世界最高の教え方を用いたとしても、結果は伴いません。

要約すると： 化学のための最高のAIを手に入れるには、スマートな基盤から始め、数学的なルールを正しく設定し、もしAIに汎用性を持たせ、物忘れを防ぎたいのであれば、「二つの帽子（マルチヘッド・リプレイ）」による教え方を選択してください。

技術概要

技術要約：MLIP基盤モデルのファインチューニング

問題提起
機械学習による原子間ポテンシャル（MLIP）基盤モデルは、多様な化学系にわたって転移する能力を示しており、タスク固有のポテンシャルをゼロから訓練するというリソース集約的なプロセスを回避するワークフローを提供する。しかし、コミュニティには、これらのモデルを「どのように」「いつ」ファインチューニングすべきかについての体系的な指針が欠けている。初期の報告では、単純な（naive）ファインチューニングはしばしば「破滅的忘却」を招くと示唆されており、その結果、大規模言語モデルから発展した制約付き手法（層の凍結や低ランク適応（LoRA）など）の採用が進んできた。本論文は、これらの制約が必要なのか、あるいは初期の失敗は他の要因（より弱い基盤モデル、不適切な原子基準エネルギー（$E_0$）の初期化、不安定な訓練手順など）によるものなのかを調査するものである。本研究の目的は、ファインチューニングの結果（特にターゲットタスクの精度と分布外（OOD）への堅牢性）を決定づける主要な要因を明らかにすることである。

手法
著者らは、7つの異なるファインチューニング戦略を、5つの化学的に多様なベンチマーク、3世代の基盤モデル、およびサイズが5桁に及ぶ訓練セットにわたって評価している。

• 評価されたファインチューニング戦略:

  1. Naive（単純）: 継続的な勾配降下法による全パラメータ更新。
  2. Layer Freezing（層の凍外 - バリアントあり）: 埋め込み層やメッセージパッシング層を凍結し、リードアウトのみを訓練する、あるいは埋め込み層と最初のメッセージパッシング層を凍結する。
  3. Low-Rank Adaptation (LoRA): 事前学習済みの重みを凍結したまま、スカラーおよび等変線形層の両方に訓練可能な低ランク分解を注入する。
  4. Multihead Replay（マルチヘッド・リプレイ）: ターゲットデータと、事前学習データまたは擬似ラベル付きデータからなるリプレイデータセットを同時に最適化し、個別のリードアウトヘッドを使用する。
  5. Pseudolabel Replay（擬似ラベル・リプレイ）: リプレイのソースを基盤モデル自身によって生成されたラベルにすることで、リプレイソースを元の事前学習コーパスから切り離した、マルチヘッド・リプレイの変種。
  6. Replay + LoRA: マルチヘッド・リプレイとLoRAの組み合わせ。

• ベンチマーク: 本研究は、OMat24の事前学習ドメイン（周期的な無機バルク）からの乖離が増大するシステムを網羅している：

  • リチウムアルジロライト電解質（無機周期固体）。
  • 水溶液NaCl（イオン溶液）。
  • 氷の多形（分子性固体）。
  • $\text{SN}_2$ 反応（気相反応化学）。
  • SPICE 生体分子（有機/生体分子コンフォーマー）。

• 技術的実装: 著者らはMACEコードベースにおいて以下の3つの新機能を実装した：

  • 等変メッセージパッシングアーキテクチャに適応したLoRA（スカラーおよび等変線形層の両方をカバー）。
  • 事前学習されたベースラインをターゲットデータに整合させるための、モデル認識型の原子基準エネルギー（$E_0$）再推定。
  • リプレイソースを事前学習コーパスから分離するための、擬似ラベル化されたリプレイ。

• 評価指標: 標準的な点（pointwise）のエネルギーおよび力の誤差に加え、本研究では動的および外挿的な挙動を調査している。これには、分子動力学（MD）からの動径分布関数（RDF）、Nudged Elastic Band（NEB）反応プロファイル、MD安定性テスト、および短距離反発の失敗を検出するためのランダム構造探索（RSS）が含まれる。

主な結果

1. 前提条件が戦略の選択を支配する: 本研究では、基盤モデルの品質、正しい$E_0$の初期化、および適切に選択されたハイパーパラメータが、特定のファインチューニング戦略よりも日常的に大きな影響を与える前提条件であることを発見した。

   • 基盤モデルの品質: 新しい基盤モデル（例：OMat24ベース）は、固定されたファインチューニングレシピを用いても、古いモデル（MPTrajベース）よりも一貫して高いOOD転移性能を示す。
   • $E_0$の初期化: 「平均化された」$E_0$値を使用すると、誤差が大幅に増大し、MDの不安定性（例：氷モデルが50 ps以内に崩壊するなど）を招く。「再推定された」$E_0$（事前学習モデルのゼロ点をターゲットデータに整合させたもの）は、安定性と転移性に極めて重要であり、多くの場合、ファインチューニングアルゴリズムの選択よりも優れた結果をもたらす。
   • ハイパーパラメータ: Naiveなファインチューニングには、学習率の低下とEMA減衰の増加が必要である。LoRAはより高い学習率を許容する。マルチヘッド・リプレイは、競合する更新信号を避けるために、大幅に低い学習率を必要とする。重み減衰（weight decay）は、パラメータを事前学習済みの解から引き離さないよう、ゼロに設定すべきである。

2. 目的別のパフォーマンス:

   • 分布内特化（単一システム）: 狭いタスク（例：$\text{SN}_2$障壁、水溶液NaCl溶媒和）において、ほとんどの戦略（Naive、LoRA、Multihead）は高い精度を達成し、ゼロから訓練されたモデルを一貫して上回る。単一システムへの適用においては、Naiveなファインチューニングが最高の収束性を提供する。
   • 分布外（OOD）への堅牢性: 関連するが未知の組成（例：非アルジロライト電解質）や異なる化学系（例：生体分子）への転移を評価する場合、Multihead Replay（オリジナルまたは擬似ラベル付きデータを使用）のみが、一貫してOODへの堅牢性を維持できる唯一のアプローチである。これは、ターゲットタスクを学習しながら、事前学習分布における精度を維持し、効果的に破滅的忘却を防ぐ。
   • 凍結とLoRA: パラメータ効率の面では効果的であるが、テストされたシナリオにおける溶媒和の特徴への適応や、広範な化学的堅牢性の維持において、層の凍結やLoRAは限界を示した。

意義と主張
本論文は、MLIPにおけるNaiveなファインチューニングの脆弱性の認識は、手法の本質的な限界ではなく、主に不適切なセットアップに起因すると主張している。著者らは以下のように述べている：

• Naiveなファインチューニングは、基盤モデルが高品質であり、$E_0$が正しく再推定されている限り、単一システムへのアプリケーションにおいて実行可能であり、しばしば優れた出発点となる。
• Multihead Replayは、ファインチューニング分布の外側における基盤モデルの挙動を保持する必要がある、より広範な展開において不可欠な戦略である。
• Pseudolabelled Replayは、リプレイ用の構造的に多様なデータセットの使用を可能にし、リプレイソースへの依存関係を取り除くことで、実用的な利点を提供する。

本研究は、実践者にとって、特定の制約付きファインチューニングアルゴリズムを選択することよりも、最強の利用可能な基盤モデルに投資し、正しい原子基準エネルギーの整合性を確保することの方が、より重要な設計上の選択であることを確立している。本研究は、ファインチューニングをニッチなオプションから、システム固有の開発におけるデフォルトの出発点へと移行させる、MLIP基盤モデル展開のための体系的なフレームワークを提供している。