FragmentNet: Adaptive Graph Fragmentation for Graph-to-Sequence Molecular Representation Learning

本論文は、分子を調整可能な粒度の化学的に妥当なフラグメントに分解する新規な適応型トークナイザーを採用するグラフからシーケンスへのモデルである FragmentNet を紹介し、このフラグメントレベルでの事前学習が、従来の原子レベルや硬直なルールベースのアプローチと比較して、下流の物性予測性能を大幅に向上させることを実証している。

要約

コンピュータに化学を理解させることを想像してみてください。従来、科学者たちはコンピュータに分子を見る方法を主に 2 つ教えてきましたが、どちらも欠点があります。

1. 「原子ごとの」アプローチ: これは、1 文字ずつ読んで小説を理解しようとするようなものです。「t」を見て、「h」を見て、「e」を見ることはできますが、「the」という単語全体を見逃してしまいます。化学において、これはコンピュータが個々の原子を見ることはできても、それらがどのように集まって機能的な部分（車のエンジンやドアノブのようなもの）を形成するかを理解することに苦労することを意味します。
2. 「硬直した規則」アプローチ: これは、定義済みで変更不可能な単語しか載っていない辞書を使うようなものです。新しい種類の単語が登場すれば、その辞書は対処できません。化学において、これは分子を断片に切り分けるために固定された規則を使用することを意味します。そこそこ機能しますが、柔軟性に欠け、自然界に見られる多様な化学構造に適応することができません。

FragmentNet の登場：「賢いレゴ」アプローチ

この論文は、分子についてコンピュータに教える新しい方法としてFragmentNetを紹介しています。単一の原子を見ることや硬直した規則を使用する代わりに、FragmentNet は学習型で適応的なトークナイザーを使用します。

分子を、レゴブロックで構築された巨大で複雑な構造だと考えてみてください。

• 従来の方法は、ブロックにあるすべての小さなプラスチックの突起（原子）を一つずつ見たり、構造をいくつかの事前に定義されたカテゴリーに無理やり当てはめたりしていました。
• FragmentNetは、その構造を見て、ブロックを意味のある塊にグループ化する方法を自ら学習します。特定のブロックの集まりが「車輪」を形成し、別の集まりが「座席」を形成し、さらに別の集まりが「エンジン」を形成すると判断するかもしれません。これらの塊が「フラグメント」です。

仕組み（3 つの魔法）

1. グループ化の学習（適応的トークナイザー）:
   このモデルは、ブロックをどのようにグループ化するかを単に推測するわけではありません。数百万の分子を研究し、どの原子のグループが化学的に通常くっつくかを学習します。ここで作成されるカスタム辞書では、「トークン」は単なる文字や原子ではなく、分子の化学的に有効な部分（機能性基全体など）となります。これは、コンピュータに「c-a-r」という文字を見るのではなく、「ing」が接尾辞であることを認識させたり、「car」が語根であることを理解させたりすることに似ています。

2. 地図の保持（空間的位置符号化）:
   3 次元のレゴ城を 1 次元の単語リスト（シーケンス）に変換すると、通常、部品が互いに対してどこにあるかという情報が失われます。FragmentNet は、すべてのフラグメントに特別な「GPS タグ」を追加することでこれを解決します。これらのタグは、コンピュータに「このエンジン部品はこの車輪部品に接続されており、座席から 3 ステップ離れている」と伝えます。これにより、分子がリストに扁平化されても、その形状を記憶し続けることができます。

3. 「穴埋め」ゲーム（マスクドフラグメントモデリング）:
   より賢くなるために、このモデルは「マッドリブス」やクロスワードパズルに似たゲームをプレイします。

   • コンピュータはフラグメントで構成された分子を見ます。
   • その中の 1 つのフラグメントを隠します（マスクします）。
   • 周囲の文脈に基づいて、その欠落した部分が何かを推測する必要があります。
   • 単一の原子ではなく、全体の塊（フラグメント）を推測するため、化学の「文法」を非常に速く学習します。「車輪」と「座席」が見えれば、欠落した部分は単なるランダムなプラスチックのブロックではなく、おそらく「エンジン」であると学習します。

論文が明らかにした結果

著者たちは、この新しい方法を、いくつかの標準的な化学テスト（薬の水中での溶解度や血液脳関門を通過できるかなどを予測するもの）において、従来の「原子ごとの」方法と比較してテストしました。

• 結果: 「賢いレゴ」アプローチ（FragmentNet）がほとんどの場合で勝利しました。
• 理由: 文脈を学習したからです。全体のフラグメントで訓練されたことで、コンピュータは特定の原子のグループが一緒に機能することを理解し、より良い予測につながりました。
• ボーナス機能: この論文はまた、モデルがこれらの塊を理解しているため、新しい有効な分子を作成するために、ある「レゴの塊」を別のものに簡単に交換できることも示しています。これは、車を分解せずに、エンジンを外して別のエンジンを取り付けるようなものです。

注意点（限界）

この論文は、その限界について率直に述べています。予算の制約のため、この実験は 1 台のラップトップ（MacBook Pro）で実行されました。大規模な AI モデルが使用する数十億に及ぶデータと比較して、使用されたデータセットは比較的小さく（200 万分子）、また「塊の大きさ」のレベルも 2 つ（非常に小さな部分対中程度の大きさの部分）のみでテストされました。

要約

FragmentNet は、個々の原子をじっと見つめるのではなく、意味のある「単語」（フラグメント）を認識し、それらの単語がどのように組み合わさって文を形成するかを理解することで、コンピュータに化学を読ませる新しいツールです。これにより、コンピュータははるかに優れた化学の学生となり、分子の挙動に関するより正確な予測が可能になります。

技術概要

技術的サマリー：FragmentNet

問題提起

分子表現学習は、従来、分子を個々の原子としてトークン化するか、BRICS などの硬直したルールベースのフラグメント分解を利用する手法に依存してきました。これらのアプローチには重大な限界があります：

1. 原子レベルのトークン化は、広範な化学的コンテキストを捉えられず、「ネガティブ転移」を引き起こすことが多く、事前学習済みモデルが単純なベースラインよりも劣る結果をもたらします。個々の原子をマスクすると、化学的に一貫性のない環境が生まれ、結合規則や官能基間の相互作用の学習を阻害します。
2. ルールベースのフラグメント化は柔軟性に欠け、多様な化学空間への汎化が困難です。
3. シーケンスベースの手法（SMILES トークン化など）は、分子グラフに固有の重要なトポロジー情報を失う傾向があります。

既存のグラフへのマスク言語モデル（MLM）戦略は、しばしば原子をマスクしますが、これにより化学的整合性が崩れます。一方、サブグラフをマスクする手法（SimSGT など）は、それらの間の相互作用を明示的にモデル化しないため、長距離依存性の捕捉が制限されます。

手法

著者らは、グラフトポロジーとシーケンスモデル間のギャップを、適応的で学習されたトークナイザを通じて埋めるために設計されたグラフからシーケンスへのモデル、FragmentNet を導入します。

1. 適応的、学習型トークナイザ

ルールベースの手法とは異なり、FragmentNet はデータ駆動型のトークナイザを採用し、分子グラフを調整可能な粒度の化学的に有効なフラグメントに分解します。

• 反復的なペアワイズマージ：トークナイザは個々の原子から開始し、学習コーパスから導出された学習済みのマージ履歴に基づいて、接続されたペアを反復的にマージします。
• 粒度制御：マージ反復回数（$T$）がトークンサイズを制御します。分子は再学習なしに最初の $t$ 回のマージ（$t \le T$）を使用してトークナイズでき、タスク固有の粒度最適化を可能にします。
• ** dangling 結合の処理**：切断された結合は「ダミー原子」（原子番号 0）で表現されます。フラグメントは、切断された結合の数と種類（例：1 つの切断された単結合を持つ炭素対 2 つ）によって区別されます。
• 一意性：立体異性体と互変異性体を区別するために、著者らはWeisfeiler-Lehman（WL）グラフハッシュアルゴリズムを使用し、非同型グラフに固有のハッシュを割り当てます。

2. 階層的エンコーダ（VQVAE + GCN）

このモデルは、ハイブリッドエンコーダを使用して原子レベルとフラグメントレベルの機能を統合します：

• VQ-VAE：離散的な原子レベルの機能を量子化された潜在空間にエンコードします。
• GCN：離散的なフラグメント内の隣接ノードからの機能を集約し、構造的な関係を捕捉します。
• 統合：原子埋め込みを平均化してフラグメント表現を形成し、それを GCN の出力と結合して圧縮されたフラグメントレベルの機能埋め込みを生成します。

3. 化学的意識を持つ空間的位置符号化（SPEs）

グラフをシーケンスに直列化する際に分子トポロジーを保持するために、FragmentNet は 3 種類の位置符号化を採用します：

• ホップベース符号化：最短経路距離を通じて相対的な接続性を捕捉します。
• WL 絶対位置符号化：異性体を区別するために、グラフ構造に基づいて一意の役割 ID を割り当てます。
• クーロン行列符号化：逆二乗則の距離と原子電荷に基づく相互作用をモデル化します。
  これらは集約され、Transformer に包括的な空間的コンテキストを提供します。

4. マスクされたフラグメントモデリング（MFM）

事前学習の目的は、個々の原子ではなく、化学的に有効なフラグメント全体をマスクすることです。

• プロセス：フラグメントが [MASK] トークンに置き換えられ、モデルはマスクされていないフラグメントのコンテキストを使用して元のフラグメントを予測します。
• 利点：これは、NLP における多単語句の再構築に類似して、化学的に意味のあるコンテキストを保持し、結合規則や機能的な関係の学習を促進します。
• 構成：著者らはコンテキストを保持するためにシーケンスあたり 1 つのトークンのみマスクすることを制限し、200 万の分子で学習しました。

5. アーキテクチャ

SPE とMolecular Descriptor CLS トークン（RDKit 記述子から派生）で強化された直列化されたフラグメント埋め込みは、Transformer エンコーダによって処理されます。プロパティ予測ヘッドは、下流タスクのためにシーケンス全体に対して最大プーリングを使用します。

主要な貢献

1. 新規学習型適応トークナイザ：構造的な接続性を保持しながら分子グラフを化学的に有効なフラグメントに分解する手法であり、調整可能な粒度を可能にします。
2. 空間的位置符号化：分子グラフトポロジーをシーケンス互換形式で捕捉する符号化セット（Hop、WL、Coulomb）であり、効果的なグラフからシーケンスへのモデリングを可能にします。
3. 粒度に関する実証的研究：トークナイズ粒度が重要な設計選択であることを示す実証です。論文は、フラグメントレベルのトークナイズが MFM 事前学習と組み合わされた場合、原子レベルのトークナイズよりも大多数のプロパティ予測タスクで優れていることを示しています。

結果

モデルは、MoleculeNet とマラリアベンチマークで、スケフォールド分割（80-10-10）を使用して評価されました。

• 事前学習の影響：MFM で事前学習された FragmentNet は、一貫して未事前学習モデルを上回りました。
• フラグメント対原子：MFM 事前学習により、フラグメントレベルのバリアント（100 回のマージ反復）は、7 つのデータセットの 5 つ（BBBP、Tox21、ToxCast、BACE、ESOL、Lipo、Malaria）で原子レベルのバリアント（0 回のマージ反復）を上回りました。事前学習なしでは、原子レベルのトークナイズがしばしば優れており、粗いトークナイズの恩恵は事前学習を通じて特に解き放たれることを示唆しています。
• 解釈可能性：アテンションマップは、溶解度（ESOL）に対するヒドロキシ基へのアテンションヘッドや、抗マラリア活性に対するキナゾリンコアへのアテンションなど、化学的に直感的なパターンを明らかにし、既知のファーマコフォアと整合していました。
• フラグメント交換：学習されたトークナイザは、部分構造一致なしに化学的に有効なアナログ（イブプロフェンの変更など）を生成するフラグメント交換モジュールを可能にし、分子編集における有用性を示しました。

意義と主張

本論文は、トークナイズ粒度が分子表現を改善するための主要なレバーであると主張しています。原子レベルからフラグメントレベルのモデリングへ移行することで、FragmentNet は原子レベルのマスクに共通するネガティブ転移の問題に対処し、より高次の構造的モチーフを捕捉します。

著者らは、このアプローチが「化学的に情報に基づいている」ことを強調し、標準的な Transformer モデルと比較してシーケンス長を短縮し、計算コストを低下させます。200 万の分子と小規模な語彙を用いた単一のラップトップという控えめな設定で学習されたにもかかわらず、事前学習されたフラグメントモデルは、未事前学習のバリアントに対して大幅な改善を示しました。

この研究は、適応的学習型トークナイズとマスクされたフラグメントモデリングの組み合わせが、分子表現学習のための実行可能かつ効果的な戦略であることを確立し、改善された下流性能と強化された化学的解釈可能性を提供します。著者らは、実験の規模（単一のラップトップ、小規模データセット）に関する限界を認め、将来の研究では特定タスクに対する最適な粒度の探索と、より大規模なモデルおよびデータセットへのスケーリングを提案しています。