[論文レビュー] In-Context Learning for Few-Shot Molecular Property Prediction

Christopher Fifty, Jure Leskovec+1 more|arXiv (Cornell University)|Oct 13, 2023

Machine Learning in Materials Science被引用数 5

ひとこと要約

この論文では、微調整を伴わずに、(分子, 特性)のデモンストレーションの文脈から直接分子特性を推論する、少サンプル分子特性予測のための新しいイン・コンテキスト・ラーニングフレームワーク、Context Aware Molecule Prediction (CAMP) を紹介する。CAMP は、分子とラベルの埋め込みを順序で処理するトランスフォーマー・エンコーダーを用い、小規模なサポートサイズにおいて最先端の性能を達成するとともに、プロトタイプ・ネットワークよりも4.5倍高速な推論を実現する。

ABSTRACT

In-context learning has become an important approach for few-shot learning in Large Language Models because of its ability to rapidly adapt to new tasks without fine-tuning model parameters. However, it is restricted to applications in natural language and inapplicable to other domains. In this paper, we adapt the concepts underpinning in-context learning to develop a new algorithm for few-shot molecular property prediction. Our approach learns to predict molecular properties from a context of (molecule, property measurement) pairs and rapidly adapts to new properties without fine-tuning. On the FS-Mol and BACE molecular property prediction benchmarks, we find this method surpasses the performance of recent meta-learning algorithms at small support sizes and is competitive with the best methods at large support sizes.

研究の動機と目的

• 自然言語に限定されてきたイン・コンテキスト・ラーニングを、創薬の分野で極めて重要な少サンプル学習が求められる分子特性予測へと拡張すること。
• 微調整や勾配更新なしに、新しい分子特性に迅速に適応できる手法を開発すること。
• 既存の少サンプル学習手法の分子科学分野における限界を克服し、大規模スクリーニングタスクにおいて低遅延かつ高精度な推論を可能にすること。
• シーケンスモデリングを用いることで、非言語分野（例：分子科学）に対してもイン・コンテキスト・ラーニングの原則が成功裏に適用可能かどうかを検証すること。
• 位置埋め込みとデモンストレーションの順序が、分子の少サンプル学習におけるモデルの一般化性能に与える影響を評価すること。

提案手法

• CAMP は、分子エンコーダー（例：MPNN）を用いてクエリおよびデモンストレーション分子をエンコードし、ラベルエンコーダーで特性測定値をエンコードする。
• 分子とラベルの埋め込みを連結してシーケンスとし、それをトランスフォーマー・エンコーダーに供給して、クエリ分子の特性を予測する。
• モデルは1回の順伝播処理で推論を実行し、微調整や重心計算を経由せずに、デモンストレーションの文脈から直接ラベルを予測する。
• アーキテクチャはデモンストレーションの順序に不変であり、アブレーションスタディでは位置埋め込みが性能を低下させ、過学習を促進することが示された。
• LLMベースのイン・コンテキスト・ラーニングとは異なり、CAMP は大規模事前学習を必要とせず、小規模な少サンプルデータセット（約5 GB）からランダム初期化で学習する。
• アブレーションスタディでは、埋め込みを交互に配置するナイーブ ICL と CAMP の連結シーケンスアプローチを比較し、分離された埋め込みの連結が学習の安定性と性能を向上させることを示した。

実験結果

リサーチクエスチョン

• RQ1 イン・コンテキスト・ラーニングの原則は、言語以外の分野（例：分子科学）の分子特性予測に成功裏に適応可能か？
• RQ2 分子設定において、分子とラベルの埋め込みを用いたシーケンスモデリングアプローチは、勾配ベースおよびメトリックベースの少サンプル学習ベースラインを上回るか？
• RQ3 文脈内のデモンストレーションの順序がモデル性能に与える影響は何か？位置埋め込みは有益か、逆効果か？
• RQ4 小規模データセットからランダム初期化で学習したモデルが、微調整なしに競争力のある性能を達成できるか？
• RQ5 埋め込み連結や位置埋め込みといったアーキテクチャ的選択が、学習の安定性と一般化性能に与える影響は何か？

主な発見

• CAMP は、FS-Mol および BACE ベンチマークにおいて、小規模なサポートサイズで最近のメタラーニングアルゴリズムを上回り、優れた少サンプル一般化性能を示した。
• 大規模なサポートサイズでは、CAMP はプロトタイプ・ネットワークを含む最良のベースラインと同等の性能を達成したが、推論速度は約4.5倍速かった。
• 位置埋め込みを含まないモデルは、位置埋め込みを含むバージョン（0.5817）よりも低い検証損失（0.5746）を達成しており、位置埋め込みが性能を損なうことが示された。
• 位置埋め込みを用いた学習では、最終的な訓練損失が 0.3181 に達した一方、位置埋め込みなしのバージョンは 0.3374 に留まり、過学習が顕著に進行した。
• 分子とラベルの埋め込みを交互に配置するナイーブ ICL は、収束が不安定で、CAMP の連結シーケンスアプローチに比べて性能が劣った。
• CAMP の性能はデモンストレーションの順序に対して頑健であり、分子予測におけるイン・コンテキスト・ラーニングが順序に依存しない（置換不変である）という仮説を支持した。