はじめに

こんにちは、Preferred Networksの森です。今回は、先日Journal of Chemcal Information and Modelling誌で出版された原子の種類と座標から結合や原子の状態を推定する手法「CoTAR」を紹介します。論文のタイトルは "CoTAR: Topology and Atomic State Reconstruction in Condensed Phases" です。

この研究の出発点は、2025年夏のインターンシップで橘川さんが取り組んだ、GNN（グラフニューラルネットワーク）に基づく凝集系の結合推定です。そこでは、周囲の原子環境から結合次数を推定するモデルを、水素原子や周期境界条件を含むMDの構造に適用しました。一方、ラジカルや電荷を持つ分子の扱い、トラジェクトリの時間方向の情報の利用が課題として残っていました。

CoTARは、この取り組みを発展させたものです。冒頭の画像のように、原子の種類と座標から、どの原子がどのように結合しているかを復元します。結合次数に加えて形式電荷と不対電子を推定し、前後のフレームの情報も利用することで、復元した分子構造を古典力場のMDにつなげられるかを検証しました。

uMLIPで得た構造を、次の計算に使うために

汎用機械学習ポテンシャル（uMLIP）は、分子の結合をあらかじめ定義せずにMDを実行できることが特徴です。ただ、得られた原子座標から分子を同定したり、古典力場を設定したりするには、どの原子が結合しているかという情報を別途用意する必要があります。この結合のつながりを分子トポロジーと呼びます。

たとえば、uMLIPで得た構造から、より長い時間の物性評価を古典MDで行いたい場合を考えてみます。古典力場では結合や角度などに応じて相互作用を設定するため、座標を渡すだけでは計算を始められません。結合次数や形式電荷を含む分子情報を復元し、力場を割り当てられる状態にする必要があります。

分子情報は解析にも役立ちます。同じ分子に属する原子がわかれば、分子ごとの移動や、特定の官能基の周囲に集まる分子を調べやすくなります。構造復元は、計算手法の切り替えと、その後の解析の両方に関わっています。

原子間距離を使えば、結合の有無はある程度推定できます。しかし、液体や溶液では水素結合や溶媒和によって、共有結合していない原子同士も近づきます。また、価数のルールから結合次数を決めようとしても、余分な結合候補があると整合する割り当てが見つからない場合があります。先行研究でGNNを導入したのは、このような場面で周囲の環境も含めて判断するためです。

結合・原子状態・時間方向の整合性をまとめて扱う

CoTARでは、まず近くにある原子同士を結んだグラフを作ります。この段階では、グラフの線は単に近接関係を表しており、共有結合と決めているわけではありません。入力は元素と座標で、必要に応じて系全体の総電荷も指定します。

このグラフに対し、GNNの一種で、周囲の情報をやり取りするMessage Passing Neural Network（MPNN）を用いて、各原子対を「結合なし・単結合・二重結合・三重結合・芳香族結合」の5状態に分類します。同時に、各原子の形式電荷と不対電子の有無も予測します。先行研究で難しかったイオンやラジカルを扱うため、原子の状態を明示的な予測対象に加えた点が工夫です。

同じC–Oの原子対でも、周りの元素や配置によって判断が変わります。MPNNは情報のやり取りを繰り返し、直接隣り合う原子のさらに周囲まで、結合や原子状態の推定に利用します。

ただし、これらを別々に当てるだけでは、組み合わせたときに価数が合わない、といった問題が起こります。そこで、結合と原子状態の整合性や、総電荷からのずれを損失関数に組み込みました。van der Waals半径に基づく距離の情報も予測に加え、化学的に整合した構造を学習しやすくしています。なお、ここで推定する形式電荷は構造式上の整数の電荷であり、力場に用いる部分電荷は後段で別途割り当てます。

さらに、トラジェクトリでは隠れマルコフモデル（HMM）を使います。非反応性のMDでは原子座標が揺らいでも結合は基本的に保たれるため、1フレームだけ結合状態が変わる予測はノイズの可能性があります。HMMは各フレームの予測確率から、時間方向にもっともらしい状態の並びを求め、一時的な誤判定を減らします。単一フレームからも推論できますが、その場合はこの補正を行いません。

古典MDのトラジェクトリで基本性能を確認

学習と評価には、次の4カテゴリ、計128系を使いました。グラフの凡例もこの分類に対応しています。

混合した分子同士の近接に加え、イオン系では電荷、ラジカル系では不対電子も正しく扱えるかを調べます。おおむね300 K、1 atmの条件で、各系は数千原子程度です。今回は基礎的な検証として、結合の形成・切断が起きない系を対象にしています。

まず、古典MDのトラジェクトリでモデルを学習します。正解ラベルは、系の構築に用いたSMILESからRDKitで取得した結合種、形式電荷、不対電子の情報です。原子の対応関係を保つことで、同じ分子情報を各フレームの教師データに利用できます。学習と評価は各系のフレームを分けて行い、学習した系での熱揺らぎに対して、安定に分子構造を復元できるかを確認しました。

グラフの横軸には、次の4手法を並べています。RDKitは化学構造を扱うソフトウェアライブラリですが、ここでは、それを用いて構成したルールベースの復元手法を「RDKit」と表記しています。

vdWとvdW＋HMMは結合の有無のみを推定するため、評価時にはすべて単結合、原子は中性・不対電子なしとしています。また、RDKitの比較手法には1フレーム5秒の制限を設け、未完了による未確定の割り当ても問題数に含めます。

結果の読み方を図1で説明します。左の「Valid snapshots」は、化学的な問題や未確定の割り当てが残っていないスナップショットの割合で、高いほどよい指標です。正解の分子構造との完全一致率を表すものではありません。右の「Problems per snapshot」は、1フレームあたりの問題数で、低いほどよい指標です。灰色の「Macro avg.」は、4カテゴリを等しい重みで平均した値を示します。純液体が多いため、単純に全データをまとめず、少数のイオン系やラジカル系の性能も反映させています。

論文の古典MD評価では、CoTARのValidは4カテゴリ平均で94.0%でした。距離に基づくvdW法や、価数のルールを用いるRDKitベースの比較手法に対して、より多くのフレームで化学的に整合した構造が得られています。問題のないフレームを増やすとともに、フレーム内に残る問題数も減らすことが、後段の処理を自動化するうえで重要になります。

少量のデータでuMLIPのトラジェクトリに適応

実際に使いたいのは、結合情報を持たないuMLIPのトラジェクトリです。ただし、古典MDとuMLIPでは、同じ分子でも結合長や局所構造の分布が完全に同じとは限りません。そこで、古典MDで学習したモデルを、少量のuMLIPデータで調整しました。

各系20フレーム、計2,560フレームを集めて1つのモデルをまとめて追加学習し、各系の別の180フレームで評価しています。これは、対象系の少量のデータを使って適応させたときの性能を調べる設計です。

図2の左を見ると、uMLIPの評価では、CoTARのValidは4カテゴリ平均で82.1%でした。距離に基づくvdW法の36.8%、vdW＋HMMの37.6%、RDKitベースの比較手法の58.2%と比べて、多くの構造で化学的に整合した分子情報が得られました。右に示す平均問題数も大きく減少し、CoTARでは4カテゴリ平均で1フレームあたり1件未満になっています。

一方、uMLIPの評価は古典MDの評価より難しく、特にイオン系のValidは61.0%にとどまっています。結合のつながりに加えて、電荷を含む原子状態を正しく扱うことが課題として残りました。

構成要素を外して比較した別の実験では、化学的な整合性を促す損失項が、問題の少ない構造を得るうえで特に効いていました。HMMは、そこから一時的な誤判定を減らす補正として働きます。先行研究で課題となっていた原子状態と時間方向の情報を組み合わせることに、実用上の意味があると考えています。

復元した分子構造で古典MDを実行する

最後に、uMLIPから復元した構造を起点に、次の古典MDを実行できるかを確認しました。図1の構造復元性能とは別に、後段のシミュレーションへの接続を調べる検証です。

CoTARの予測からOpenFF 2.0.0の力場を設定し、OpenMMでMDを行ったところ、128系中110系、85.9%で後段の検証を完了しました。HMMによる補正前は83.6%だったため、一時的な誤判定を減らすことが、古典MDまでつながる系を増やすことにも寄与しています。

計算が完了した系では、正解のトポロジーを使った古典MDと、密度、拡散係数、動径分布関数（RDF）がよく対応しました。密度のR²は0.982、安定に評価できた拡散係数の対数スケールでのR²は0.908でした。これは同じ古典力場の設定での比較であり、実験やuMLIPの物性を再現したという評価ではありませんが、構造復元を挟んでも、系の詰まり方や分子の動き、局所構造を保てるかを確認できます。

今後の課題と展望

現時点では、イオン系の形式電荷の推定に課題が残っています。学習に使わなかった6系への追加学習なしの評価では、非イオン系で良好な結果が得られた一方、2つのイオン混合系では形式電荷の誤りが残りました。学習データに含まれるイオンの種類は限られており、有機カチオンなど異なる電荷環境への対応が必要です。HMMは一時的な誤判定を減らせますが、継続して同じ電荷を間違えるような誤りは直せません。

また、今回の検証は結合が変わらない系に限っています。先行研究でも取り上げた反応中間体や高分子の分解過程に適用するには、反応を含むデータを用意し、実際の結合変化を正しく追えるかを確かめる必要があります。

CoTARでは、先行研究を基盤に原子状態と時間方向の情報を取り込み、多くの系で古典MDへの接続を確認しました。こうした構造復元を整えることが、uMLIPを解析や物性評価まで含めて使いやすくするうえで大切だと考えています。

評価用の実装はGitHubで公開しています。公開版は動作を試すための縮小版で、同梱モデルは論文の評価に使ったモデルとは異なります。詳しくは論文と合わせてご覧ください。