概要
GROVER は、論文「Self-Supervised Graph Transformer on Large-Scale Molecular Data」に付属する PyTorch 実装です。ダウンロード可能な事前学習済みモデルと、分子表現を学習し、ラベル付きの物性予測タスクに適応させるコードを組み合わせています。README には GROVER base と large のチェックポイントに加え、11 種類の分類・回帰データセットでファインチューニングされたモデルが掲載されています。ワークフローは、表現学習、下流タスクの学習、予測、フィンガープリント抽出、評価にわたります。
事前学習はラベルなしの分子データから開始し、意味モチーフラベル、原子および結合のコンテキスト語彙、分割済みのデータセット構造を必要とします。文書に記載された準備手順では、特徴量配列、語彙辞書、個別に分割されたグラフデータが生成されます。ファインチューニングでは smiles 列を含む CSV を使用でき、.npz ファイルに保存された追加の分子記述子も利用できます。ファインチューニング済みチェックポイントは、出力ファイルに書き込む予測に利用できます。予測値は、設定されたフォールド数とアンサンブルメンバーにわたって平均されます。
表現ベースのワークフローでは、プーリングされた原子埋め込み、プーリングされた結合埋め込み、またはその連結から得られるフィンガープリントを GROVER から出力できます。追加の分子特徴量を指定した場合は、それらを付加することもできます。そのため、このプロジェクトは特定タスクの予測だけでなく、独立した下流解析で学習済み分子表現を評価する用途にも適しています。
README は、再現に関係する制限事項を示しています。index_select_nd に関連する非決定的な挙動により、ファインチューニング結果の厳密な再現は困難です。また、この実装では元の実装と異なり、メッセージパッシング層に密結合が追加されています。プロジェクトは評価手順とファインチューニング済みチェックポイントを提供しますが、新しいユーザーのシステム上で再現性が確立されているわけではありません。また、このプロジェクトは Tencent の公式サポート対象製品ではないことも明記されています。
主な機能
- 意味モチーフラベルと原子・結合のコンテキスト語彙を用いた自己教師あり `GTransformer` 事前学習。文書には、単一 GPU および Horovod を用いたマルチ GPU のワークフローが記載されています。
- ダウンロード可能な GROVER base および large の事前学習済みチェックポイントと、11 種類の分子ベンチマークデータセットでファインチューニングされたチェックポイント。
- `smiles` 列を含むラベル付き CSV データセットでのファインチューニング。`.npz` 配列として保存された `rdkit_2d_normalized` 分子特徴量を任意で利用できます。
- ファインチューニング済みモデルによる予測。設定されたフォールド数とアンサンブルメンバーの出力を平均します。
- 原子埋め込み、結合埋め込み、またはその連結を用いた分子フィンガープリントの出力。追加の分子特徴量も任意で利用できます。
- 文書に記載された分類用 AUC の設定、回帰のデフォルトである RMSE、QM7 および QM8 の再現に用いる MAE によるチェックポイント評価。
用途
- 評価案：事前学習済み GROVER チェックポイントをラベル付きの分子分類または回帰データセットに適応させ、ホールドアウト分割で性能を評価する。
- 評価案：原子、結合、およびそれらを組み合わせた GROVER フィンガープリントを、独立した下流の物性予測解析の入力として比較する。
- 評価案：提供されたファインチューニング済みベンチマークチェックポイントと文書記載の指標を用いて、README に報告された実験との一致度を調べる。
- 評価案：モチーフラベル、コンテキスト語彙、データ分割を準備したうえで、ラベルなしの分子集合から分子表現を事前学習する。
使い方
- 公式 READMEを読み、事前学習、ファインチューニング、予測、フィンガープリント生成、評価のいずれかを選びます。依存関係と環境に関する説明を確認してください。Python 3.6.8 が指定され、
requirements.txtとDockerfileが参照されています。 - 事前学習済みモデルを出発点にする場合は、GROVER base チェックポイントを入手するか、README 内の large モデルへのリンクを選びます。事前学習済み表現チェックポイントと、タスク固有のファインチューニング済みモデルを区別してください。
- 該当する README の例に従って入力を準備します。ファインチューニングには
smilesを含む CSV が必要です。事前学習にはさらに、モチーフラベル、原子・結合の語彙、データセットの分割が必要であり、単一 GPU で実行する場合も同様です。 - 文書記載のファインチューニングまたはフィンガープリントのワークフローに従います。学習に追加の分子特徴量を使用した場合は、予測入力用に対応する特徴量を生成してください。解析の目的に応じて、原子、結合、またはそれらを組み合わせたフィンガープリントを選びます。
- 保存済みチェックポイントを、文書記載のタスク指標と分割設定で評価します。これは提案されるローカル評価として扱い、記載された非決定的な挙動と密結合の違いを考慮してください。コードのライセンスは、チェックポイントやデータセットの条件とは別に確認してください。