3D Masked Autoencoders are Robust Learners of Volumetric and Multimodal Cellular Representations for Microscopy

技術要約：細胞表現学習のための3D Masked Autoencoders

問題提起
蛍光顕微鏡法は、細胞のサブセルラー（細胞内小器官レベル）の組織化を捉えた、高次元で豊かなデータ（例：OpenCell、WTC-11のようなボリューム形式のzスタック）を提供します。しかし、SubcellやDINO4Cellといった単一細胞解析のための既存の表現学習アプローチは、主に2Dの最大輝度投影（maximum-intensity projection）や個別のスライス操作に依存しています。この手法は、固有の深度分解構造情報を破棄してしまうため、モデルが識別性の高い細胞表現を学習する能力を制限する可能性があります。さらに、純粋な視覚モデルは、タンパク質配列にエンコードされた生物学的事前知識との統合を欠いていることがよくあります。本論文では、ネイティブな3Dモデリングとタンパク質言語モデルとのマルチモーダルな整列（alignment）が、これらの限界を克服し、タンパク質局在化やタンパク質間相互作用（PPI）予測といったダウンストリームタスクを改善できるかどうかを調査しています。

手法
著者らは、ボリューム形式の顕微鏡データに対するMasked Autoencoder (MAE) フレームワークの体系的な比較と強化を提案しています。

• アーキテクチャのバリエーション:

  • MAE-2D: 入力される3Dボリュームは2D最大輝度投影 ($I_{2D}$) に集約されます。モデルは、正弦波位置エンコーディングを用いた2D ViTバックボーンを使用して、重なりのない2Dパッチを処理します。
  • MAE-3D: 完全な3Dボリューム ($I_{3D}$) が保持され、$p_z \times p_x \times p_y$ の3Dパッチに分割されます。モデルは、エンコーダー・デコーダーのパイプライン全体を通じて深度次元を保持する、3D正弦波位置エンコーディングを備えた3D ViTバックボーンを利用します。
  • Channel Cross-Attention (CCA): チャネル間の相互作用を明示的にモデル化するために、各チャネルを独立したトークンストリームとするデュアルストリームのエンコーダー・デコーダーを導入しています。チャネル間で共有されたランダムなマスキングパターンが適用されます。クロスアテンションは、チャネル $c_i$ における空間インデックス $i$ のトークンが、同じインデックス $i$ におけるチャネル $c_j$ のトークンに対してクエリを行うように実装されています。シングルキー制約を処理するために、softmax操作の代わりにシグモイド・ゲーティング機構が採用されています。
  • 周波数領域の正則化 (FFT Loss): 微細なサブセルラー構造を保持するために、元の画像と再構成された画像の振幅スペクトルの間で、周波数領域におけるL1損失が適用されます。この損失は、MAE-2Dでは2D軸に対して、MAE-3Dでは全3つのボリューム軸に対して計算されます。学習を安定させるため、訓練中に線形に増大（ramp up）させます。

• マルチモーダルな整列:
  本フレームワークは、事前学習済みのタンパク質言語モデル（ESM2）を統合し、視覚的特徴をタンパク質配列情報と整列させます。

  • デコーダーの条件付け: 固定されたESM2埋め込みは、デコーダーの次元に投影され、各チャネル固有のデコーダーシーケンスに単一のトークンとして挿入されます（位置エンコーディングはゼロ）。これにより、マスクされたパッチが自己アテンションを介してタンパク質のコンテキストに注意を向けることが可能になります。
  • 対照学習目的関数: 対称的なInfoNCE損失 ($L_{CLIP}$) が、正規化された画像埋め込みをタンパク質埋め込みと整列させます。この整列は、堅牢性を高めるために可視トークンに対してのみ計算されます。
  • 総損失: 最終的な目的関数は、再構成のための平均二乗誤差（MSE）、FFT損失、およびCLIP損失を組み合わせたものです：$L = L_{MSE} + w_{FF_T}L_{FF_T} + w_{CLIP}L_{CLIP}$。

主な貢献

1. 体系的な3D vs 2Dの比較: 著者らは、ネイティブな3D MAEモデルが、2Dの最大投影やスライスベースのバリアントと比較して、ダウンストリームタスクにおいて一貫して優れた性能を示すことを実証しました。これは、2Dモデルが同じ数のスライスに制限されている場合でも同様です。
2. CCAとFFTによる強化: 著者らは、チャネル間クロスアテンションと周波数領域の正則化が、3D空間コンテキストを活用し、構造的忠実度を保持するための重要な要素であることを特定しました。
3. マルチモーダルな統合: 本研究は、画像特徴をタンパク質配列の埋め込み（ESM2経由）と整列させることが、表現の質を大幅に向上させることを確立しました。その恩恵は、2Dモデルよりも3Dモデルにおいて顕著でした。

結果
実験は、OpenCellデータセット（1,310種類のタンパク質、約6,300ボリューム）を用い、5分割交差検証で行われました。

• タンパク質局在化:

  • MAE-3D vs MAE-2D: MAE-3Dは、すべてのマスク比率において、MAE-2Dを一貫して上回りました。中央のスライス範囲（45–55）に制限した場合でも、MAE-3Dは優れた性能を維持しました（AUCmicro: 0.91 vs 0.90）。
  • マルチモーダルの恩恵: 最良の3Dモデル（MAE-3D⋆）は、AUCmicro 0.952、F1micro 0.742 という最先端の結果を達成し、従来の手法をそれぞれ+0.003および+0.010改善しました。ESM2の統合は、2Dバリアントと比較して、3Dモデルにおいてより大きな性能向上をもたらしました。
  • アテンション解析: 可視化により、ESM2を統合することで、アテンションが関連するzスライスにわたってタンパク質に関連する構造（例：アクチン細胞骨格）に集中することが示されました。一方、ESM2なしのモデルや2Dモデルは、アテンションが分散しているか、あるいは深度の特異性に欠けていました。

• タンパク質間相互作用 (PPI):

  • MAE-3D⋆はROC-AUC 0.865 を達成し、従来の手法（Subcell: 0.85、DINO4Cell: 0.84）を最大+0.025上回りました。
  • 特筆すべきは、3Dモデルが、はるかに大規模なHPAデータセットで事前学習されたモデルと比較して、より小さなOpenCellデータセット（約6kボリューム）でこれらの結果を達成したことです。これは、データが限られた設定におけるボリューム形式モデリングの効率性を強調しています。

意義と主張
本論文は、完全な3Dボリュームコンテキストを保持することが、2D投影よりも、単一細胞顕微鏡における、より情報豊富で識別性の高い表現を生み出すと主張しています。また、体積的な形態学とタンパク質レベルのセマンティクスとの相乗効果は、ネイティブな3Dアーキテクチャを使用する場合に最も強固になると述べています。著者らは、ネイティブな3Dモデリングとマルチモーダルな整列の組み合わせが、細胞イメージングにおける堅牢な基盤モデルの開発に不可欠であると結論付けています。彼らは、これらの利点が大規模な事前学習なしでも保持されることを強調しており、アーキテクチャの選択（3D vs 2D）とマルチモーダルな統合が、性能における決定的な要因であることを示唆しています。