[論文レビュー] Enforcing Predictive Invariance across Structured Biomedical Domains

Wengong Jin, Regina Barzilay+1 more|arXiv (Cornell University)|Jun 6, 2020

Machine Learning in Materials Science参考文献 33被引用数 8

ひとこと要約

本稿では、分子スケルトンやタンパク質スーパーファミリーなどの複雑で構造的な環境において、予測の不変性を強制することで、バイオメディカルAIにおけるモデルの汎化性能を向上させるため、レグレット最小化（RGM）およびその構造的拡張（SRGM）を提案する。一般化を、未見の環境に後知恵でアクセスできるオラクルに対するレグレットを最小化する問題として定式化し、未学習ドメインの変化をシミュレートする勾配ベースの摂動を用いることで、QM9における分子性質予測（41.7 vs. 52.3 MAE）およびタンパク質安定性予測（0.793 vs. 0.73 Spearman’s ρ）で最先端の性能を達成する。

ABSTRACT

Many biochemical applications such as molecular property prediction require models to generalize beyond their training domains (environments). Moreover, natural environments in these tasks are structured, defined by complex descriptors such as molecular scaffolds or protein families. Therefore, most environments are either never seen during training, or contain only a single training example. To address these challenges, we propose a new regret minimization (RGM) algorithm and its extension for structured environments. RGM builds from invariant risk minimization (IRM) by recasting simultaneous optimality condition in terms of predictive regret, finding a representation that enables the predictor to compete against an oracle with hindsight access to held-out environments. The structured extension adaptively highlights variation due to complex environments via specialized domain perturbations. We evaluate our method on multiple applications: molecular property prediction, protein homology and stability prediction and show that RGM significantly outperforms previous state-of-the-art baselines.

研究の動機と目的

• 訓練データの分布を超えたバイオメディカルモデルの汎化に挑戦する。特に、分子スケルトンやタンパク質スーパーファミリーなどの構造的かつ疎な表現を持つ環境において有効であることを目的とする。
• 代表的表現学習に対する制約が弱いため、単一例や極めて構造的な環境で困難をきたす既存の不変性ベース手法（例：IRM）の限界を克服すること。
• 未学習の環境に後知恵でアクセスできるオラクルに対する一般化性能を数量化するレグレット最小化フレームワークを構築し、より強い不変性誘導を可能にする。
• ドメイン変化を明示的に定義せずに、表現の勾配ベースの摂動を用いることで、構造的ドメインにレグレットフレームワークを拡張する。
• 実世界のバイオメディカルベンチマーク、特に分子性質予測、タンパク質安定性および相同性分類タスクにおいて、提案手法の優位性を実証する。

提案手法

• 不変リスク最小化（IRM）を、ホールドアウトされた環境にアクセスできるか否かの予測器の差（レグレット）を最小化する問題に再定式化する。
• 訓練データのサブセットをサンプリングしてホールドアウト環境 $E_e$ を定義し、$E_e$ を用いて訓練した予測器とそうでない予測器の損失差としてレグレットを計算する。
• 構造的環境（例：分子スケルトン）では、分類器 $g$ がドメイン（例：スケルトンまたはトポロジー）を予測するようにし、表現 $\phi$ に対して勾配ベースの摂動を適用することで、未観測のドメイン変化をシミュレートする。
• 摂動 $\nabla_\phi \mathcal{L}_{\text{cls}}(g, \phi)$ はスケルトン固有の変化を強調し、摂動された例におけるレグレットは未学習のスケルトンタイプへの一般化を測定する。
• 複数のホールドアウト環境サンプルおよび摂動スキームにわたるレグレットを最小化するように、主モデルをエンドツーエンドで最適化する。
• 希少または疎な環境（例：10例未満のタンパク質スーパーファミリー）をクラスタリングして2つの粗い環境 $E_0, E_1$ にグループ化し、標準ベースラインの学習を安定化させる一方、SRGMは勾配摂動を用いて構造的ドメインを直接処理する。

実験結果

リサーチクエスチョン

• RQ1 疎なまたは構造的な環境において、IRMに比べてレグレット最小化がより強い不変性誘導バイアスを提供できるか？
• RQ2 すべての環境を明示的にモデル化せずに、表現の勾配ベースの摂動が、構造的バイオメディカルドメインにおける一般化をどのように向上させるか？
• RQ3 分子スケルトンやタンパク質トポロジーなどの構造的ドメイン情報を活用することで、分子およびタンパク質モデリングタスクにおける一般化性能がどの程度向上するか？
• RQ4 本手法は、外挿およびスケルトン分割評価プロトコルにおいて、IRM、MLDG、CrossGradなどの既存のドメイン一般化ベースラインを上回るか？
• RQ5 特に訓練分布とテスト分布が著しく乖離する場合、本手法はどの程度のドメインシフトのレベルにスケーリング可能か？

主な発見

• スケルトン分割評価下のQM9データセットにおいて、RGMは12種類の分子性質の平均MAEを、最良のベースライン（52.3）から41.7に低下させ、顕著な一般化性能の向上を示した。
• SRGMはタンパク質安定性予測タスクでSpearman相関係数0.793を達成し、前回SOTAの0.73およびERMベースラインの0.736を上回った。
• 相同性予測タスクでは、SRGMが23.8%の精度を達成し、最良のベースライン（ERM）の22.3%を上回り、大規模な進化的ギャップに対しても頑健であった。
• アブレーションスタディでは、RGMおよびCrossGradを上回る一貫性ある性能を、ドメインシフトの増加に伴って示した（例：8、7、6原子以下の分子で学習）。これは、極端な分布シフト下でも有効であることを確認した。
• HIVデータセットでは、SRGMが0.735の精度を達成し、次に良いベースライン（CrossGrad：0.708）を著しく上回った。これは、構造的ドメインモデリングの利点を示している。
• 本手法は、分子性質予測、タンパク質安定性、相同性検出といった多様なバイオメディカルタスクで一貫した向上を示し、低データ環境における構造的で一般化可能なモデルの有効性を検証した。