MapPFN: 文脈における因果的擾乱マップの学習に関する技術的サマリー

本論文「MapPFN: Learning Causal Perturbation Maps in Context」は、生物学的システム（特に単細胞データ）における介入効果（遺伝子ノックアウトなど）を予測するための新しい深層学習フレームワーク「MapPFN」を提案するものです。既存の手法がトレーニングデータで見た生物学的文脈に限定されるという課題に対し、MapPFN は**イン・コンテキスト・ラーニング（文脈学習）と事前データ適合ネットワーク（PFN: Prior-Data Fitted Network）**の概念を組み合わせることで、未見の生物学的文脈や遺伝子セットへの適応を可能にします。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細にまとめます。

---

1. 問題定義と背景

背景

単細胞擾乱データ（Perturb-Seq など）は、遺伝子制御ネットワーク（GRN）の因果依存関係を解明する黄金標準ですが、実験的に全細胞状態と擾乱の空間を網羅することは不可能です。そのため、仮想細胞モデル（Virtual Cell Models）を用いて、実験前に介入効果を予測する研究が進んでいます。

既存手法の課題

• 文脈適応性の欠如: 既存の手法（CPA, CellOT, Meta Flow Matching など）は、トレーニング中に観測された生物学的文脈（細胞株や条件）に強く依存しています。トレーニングデータに含まれない新しい文脈や遺伝子セットに対しては、推論時に適応できず、性能が低下します。
• 非対分布のマップ: 単細胞シーケンシングは細胞を破壊するため、介入前後のデータはペアになっていません（unpaired）。これは、事前分布 $p(y_{obs})$ から事後分布 $p(y_{int})$ への分布間マッピング問題となります。
• 推論時の適応: 既存のモデルは、新しいデータセットに対してゼロショット（事前学習なし）で適応できず、通常は新しいデータでゼロからトレーニングする必要があります。

目標

トレーニングデータに含まれない生物学的文脈や任意の遺伝子セットに対して、推論時（テスト時）に観測された少数の介入実験データ（文脈）に基づいて適応し、未見の擾乱効果を予測するモデルの構築。

---

2. 手法：MapPFN

MapPFN は、**事前データ適合ネットワーク（PFN）**の枠組みを単細胞擾乱予測に初めて適用したモデルです。

2.1 基本的なアプローチ

• メタ学習とイン・コンテキスト・ラーニング:
  合成データ（シミュレーション）上で多数の異なる因果構造（SCM: Structural Causal Models）に対して事前学習を行います。推論時には、ターゲットとなる生物学的文脈における「観測データ」と「一部の既知の介入実験データ（コンテキストセット）」を入力として与え、Transformer のイン・コンテキスト・ラーニング能力を用いて、未見の擾乱に対する分布を予測します。
• 事前学習の戦略:
  実データではなく、合成データで事前学習を行います。これにより、実験コストやデータ量の制約から解放され、多様な因果メカニズムを網羅的に学習できます。

2.2 モデルアーキテクチャ

• Multi-modal Diffusion Transformer (MMDiT):
  Stable Diffusion 3 のアーキテクチャを基盤とし、細胞をトークンとして扱います。
  • 3 つのモダリティストリーム: ノイズ、細胞状態（観測/介入）、治療（One-hot エンコーディング）を別々のパラメータで処理し、ジョイントアテンションで相互作用させます。
  • 位置符号の削除: 細胞の順序は意味を持たないため、正弦波位置符号を削除し、アテンションの置換不変性（permutation invariance）に依存します。代わりに、モダリティやクエリ/コンテキストを区別する学習可能な埋め込みを導入します。
• 学習目標:
  条件付きフローマッチング（Conditional Flow Matching）を用います。ノイズから予測される事後分布への速度場（velocity field）を学習し、条件 $v_\theta^\tau(Y_\tau | t_q, Y_{obs}, C)$ を最適化します。

2.3 合成事前分布（Synthetic Biological Prior）

実データの不足を補うため、以下の手順で合成データを生成して事前学習を行います。

1. SCM のサンプリング: 遺伝子制御ネットワーク（GRN）を、スケーリングフリーなグラフ生成アルゴリズムや線形加性ノイズモデル（ANM）を用いてサンプリングします。
2. シミュレーション: SERGIO シミュレータを用いて、Hill 関数による非線形相互作用や技術的ノイズ（10x Chromium 配列に相当）を考慮した単細胞発現データを生成します。
3. 介入シミュレーション: 遺伝子ノックアウトを「do-操作」としてシミュレートし、介入後の分布を生成します。
4. 対数ペア（Counterfactual Paired）の活用: 推論時の非対データの問題を軽減するため、事前学習時には同じランダムシードを用いて対数ペア（介入前後の細胞が対応している状態）を生成し、因果効果のみを学習させるアプローチも採用しています。

---

3. 主要な貢献

1. 文脈学習による分布マッピングの定式化:
   擾乱予測を、テスト時の介入文脈（観測データ＋既知の介入データ）に基づいた分布マッピング問題として定式化し、単一の事前学習済みモデルで未見の生物学的文脈や任意の遺伝子セットに適応可能にしました。
2. 擾乱予測のための最初の PFN（MapPFN）の導入:
   合成生物学的事前分布（in silico 遺伝子ノックアウト）で事前学習された PFN を提案しました。これにより、実験的擾乱データの可用性に依存せず、多様な因果構造に対するメタ学習が可能になりました。
3. 合成ベンチマークでの性能向上:
   構造的因果モデル（SCM）の制御された合成ベンチマークにおいて、既存のベースラインをすべての指標で上回りました。特に、介入文脈と対数事前分布（counterfactual prior）が予測精度に独立して寄与することを示しました。
4. 実データでのゼロショットおよび微調整性能:
   異なる生物学的特性を持つ 2 つの単細胞擾乱データセット（黒色腫、白血病）において、実データでゼロから学習したベースラインと同等の「ゼロショット」で差異発現遺伝子（DEG）を回復しました。さらに、実データで微調整（Fine-tuning）を行うことで、すべてのデータセットでベースラインを一貫して上回りました。

---

4. 実験結果

4.1 合成データ（線形 SCM）での評価

• 指標: Wasserstein 距離 ($W_2$), MMD, RMSE, 擾乱識別スコア (PDS), 大きさ比率 (MR)。
• 結果: MapPFN はすべての指標で最良の性能を示しました。特に、ベースラインが観測分布に収束する（Identity Collapse）傾向があったのに対し、MapPFN は介入効果を正確に捉え、大きさ比率（MR）が 1.0 に近い値を示しました。

4.2 実単細胞データでの評価

• データセット:
  • 黒色腫（Frangieh et al., 2021）: 3 つの生物学的文脈、248 遺伝子。
  • 白血病（Papalexi et al., 2021）: 1 つの文脈、26 遺伝子。
• ゼロショット性能:
  合成データのみで事前学習した MapPFN は、実データでゼロから学習したベースライン（CPA, CondOT など）と同等の DEG 回復精度（AUPRC）を達成しました。
• 微調整後の性能:
  合成事前学習後に実データで微調整した MapPFN は、すべての指標（PDS, MR, AUPRC など）でベースラインを凌駕しました。これは、合成データでのメタ学習が実世界の多様な生物学的文脈に一般化することを示しています。

4.3 消融実験（Ablation Study）

• 介入文脈の重要性: 介入データ（コンテキスト）を条件付けしない場合（観測データのみ）、性能が全指標で低下しました。介入データは因果グラフの同定可能性（identifiability）を高め、未見の擾乱に対するマッピングを改善します。
• 対数ペア事前分布: 非対データ（実データに近い設定）ではなく、対数ペアデータで事前学習した場合、収束が速く、最終性能も向上しました。

4.4 スケーラビリティ

• 遺伝子セットの拡張: 学習時に 50 遺伝子しか見ていなくても、テスト時拡張（TTA）を用いることで 100 遺伝子以上のセットに適応し、予測精度を維持・向上させました。
• データ量の増加: 推論時にコンテキストに含まれる細胞数や実験数が増えると、性能が単調に向上しました（イン・コンテキスト・ラーニングの特性）。

---

5. 意義と将来展望

意義

• 仮想細胞モデルの新たなパラダイム: 従来の「特定のデータセットに特化したモデル」から、「合成データでメタ学習し、実データで即座に適応する汎用モデル」への転換を示しました。
• 実験コストの削減: 限られた実データ（少数の介入実験）だけで、未見の条件や遺伝子に対する効果を高精度に予測できるため、創薬ターゲットの探索や実験計画の効率化に寄与します。
• 因果推論と深層学習の融合: PFN と因果推論の概念を組み合わせ、分布マッピング問題に対して理論的根拠（同定可能性の向上）を持つアプローチを確立しました。

限界と将来の課題

• 事前分布への依存: 性能は合成データの事前分布の質に依存します。より多様な生物学的メカニズムをカバーする事前分布の構築が必要です。
• 計算コスト: 事前学習には GPU 時間がかかります（約 36 GPU 時間）が、一度学習すれば微調整は短時間で完了します。
• 拡張性: 現在のモデルは単一遺伝子のノックアウト（ハード介入）に特化しており、組み合わせ擾乱や薬剤（化学的介入）への対応は今後の課題です。

結論

MapPFN は、単細胞擾乱予測において、合成データによる事前学習と推論時の文脈学習を統合することで、未知の生物学的文脈への強力な一般化能力を実現しました。これは、文脈適応型の仮想細胞基盤モデル（Context-Adaptive Virtual Cell Foundation Models）への重要な一歩となります。