Toward Better Geometric Representations for Molecule Generative Models

技術概要：LENSEs – 非滑らかな構造符号化のための潜在空間強化

問題定義
幾何学的表現条件付き生成（GeoRCG）は、分子表現のモデリングと 3 次元構造の生成を分離することで、3 次元分子生成における効率的なパラダイムとして登場しました。このフレームワークでは、事前学習された分子エンコーダ（例：UniMol、Frad）が潜在空間を定義し、表現拡散モデル（RDM）がこの空間から分子レベルのベクトルをサンプリングし、生成器がこのベクトルを条件として 3 次元構造を再構築します。

しかし、GeoRCG の有効性は、事前学習されたエンコーダが提供する潜在空間の品質によって根本的に制限されています。著者らは 2 つの重要なボトルネックを特定しました：

1. 非滑らかな表現多様体：主に座標ノイズ除去による物性予測のために最適化された事前学習エンコーダは、高いリプシッツ定数を持つ表現を生成します。分子配座のわずかな摂動が、潜在空間内で不釣り合いに大きなシフトを引き起こします。この非滑らかさは、RDM による分布のモデリングを困難にし、手動調整されたノイズ注入などの補償メカニズムを頻繁に必要とします。
2. 未活用の多レベル意味情報：事前学習エンコーダは、浅い層が局所的な構造モチーフ（例：官能基）を捉え、深い層がグローバルな意味を符号化する階層的な特徴を示します。既存の GeoRCG パイプラインは通常、最終層の表現のみを利用し、初期層に保持される豊かな意味情報を破棄しています。その結果、生成器はこれらの意味情報をゼロから再学習する必要があり、最適化の難易度が高まります。

手法：LENSEs
これらの課題に対処するため、著者らは LENSEs（非滑らかな構造符号化のための潜在空間強化）を提案しました。これは、生成訓練中に条件付け潜在空間を能動的に洗練させるフレームワークです。凍結されたエンコーダに受動的に依存するのではなく、LENSEs は 3 つの相補的なメカニズムを導入します：

1. 学習可能な層プーリングを備えた表現ヘッド：
   エンコーダからの単一の決定論的出力を使用する代わりに、LENSEs は多層特徴を統合します。学習可能な層プーリング機構を用いて、異なるエンコーダ深度（$H^{(l)}$）からの表現に重みを付けます。これらのプーリングされた特徴は、投影ヘッドを介して対角ガウス事後分布（$\mu, \sigma$）のパラメータにマッピングされます。潜在コード $z$ は再パラメータ化トリックを用いてサンプリングされます。このヘッドは、意味的に豊かでかつ生成に適した表現を抽出するために生成器と共同で訓練され、コンパクトな潜在空間を確保する KL 発散損失によって正則化されます。

2. 分子知覚損失（$L_{PERC}$）：
   計算機ビジョンにおける知覚損失に触発され、この損失は、ノイズ除去された（生成された）分子の高次意味表現を、凍結されたエンコーダの特徴空間内において、クリーンな真の分子（グラウンドトラウト）のそれと整合させます。生成された分子とクリーンな分子の両方をエンコーダに通し、プーリングされたグラフレベル表現間の距離を最小化することで、生成器は低レベルの座標誤差の最小化だけでなく、意味のある構造パターン（例：官能基）を保持するように促されます。真の表現にはストップグラディエント操作が適用されます。

3. 表現整合（REPA）損失（$L_{REPA}$）：
   生成器の隠れ状態がエンコーダの特徴と自然に整合するという観察に動機付けられ、この損失は生成器のノードレベル特徴を事前学習エンコーダのそれと明示的に整合させます。学習可能なプロジェクタが生成器の特徴をエンコーダの特徴空間にマッピングし、対応するノード間でコサイン類似度を最大化します。これにより、事前学習と生成との間の意味的ギャップが縮小され、訓練が加速され、化学的に意味のある知識の獲得が改善されます。

総合的な訓練目的関数は、ベース生成損失（$L_{gen}$）に KL 損失、知覚損失、および REPA 損失を組み合わせたものです。

主要な結果
著者らは、LENSEs を GEOM-DRUG および QM9 データセットで評価し、GeoRCG、CanonFlow、SemlaFlow などの最先端のベースラインと比較しました。

• 生成品質：困難な GEOM-DRUG データセットにおいて、LENSEs は最先端のパフォーマンスを達成しました。有効性 97.28%、分子安定性 98.51% を達成し、以前の最高記録である GeoRCG（有効性 95.3%、安定性 97.20%）および CanonFlow（有効性 95.9%）を上回りました。また、最も低いエネルギーおよび歪み指標を達成し、優れた物理的妥当性を示しました。
• 表現の滑らかさ：リプシッツ定数の分析により、洗練された潜在空間が生のエンコーダ空間よりも著しく滑らかであることが明らかになりました。GEOM-DRUG において、実証的なリプシッツ定数は 4.6 倍 低減され、表現多様体が拡散モデリングにより適していることが確認されました。
• 表現の豊かさ：表現の有効ランクは、生エンコーダの 65.32 から LENSEs の 88.42 に増加し、次元の崩壊の減少と、より豊かで構造化された潜在空間を示しています。
• 下流タスクでの有用性：凍結されたエンコーダと表現ヘッドを用いて QM9 物性に対する下流回帰タスクに使用した場合、LENSEs はすべての物性において生エンコーダを上回り、生成訓練プロセス自体が分子表現を洗練するための効果的な事前学習パラダイムとして機能することを示しました。
• サンプリング効率：LENSEs は、ベースラインと比較して少ない生成ステップで高い有効性と安定性を達成し、サンプリング効率の向上を示しました。

意義と主張
本論文は、高品質な 3 次元分子生成は、生成器のアーキテクチャだけでなく、条件付け表現の品質に決定的に依存すると主張しています。LENSEs は、事前学習エンコーダの潜在空間を、より滑らかで、意味的に情報豊富で、生成目的と整合性の取れたものへと体系的に改善できることを実証しました。

著者らは、生成品質と表現品質の間に双方向的な関係が確立されると提唱しています。具体的には、整合性目的を用いた生成訓練は、分子エンコーダのための潜在的な事前学習パラダイムとなり得ると示唆しています。これにより、表現はより予測的かつ下流タスクに有用なものへと洗練されます。この研究は、有用な分子表現は本質的に生成に適しているわけではなく、生成モデリングにおけるその潜在能力を最大限に引き出すためには明示的な洗練が必要であることを浮き彫りにしています。