[論文レビュー] Masked Autoencoders for Microscopy are Scalable Learners of Cellular Biology

Oren Kraus, Kian Kenyon-Dean+15 more|arXiv (Cornell University)|Apr 16, 2024

Cell Image Analysis Techniques被引用数 10

ひとこと要約

本論文は、Vision Transformer（ViT）バックボーンを用いた自己教師ありマスク付きオートエンコーダー（MAE）が大規模な顕微鏡データセットへ拡張可能で、既知の生物学的関係の想起において弱教師付き法を上回ることを示す；さらに、異なるチャネル構成に対して一般化するチャネル非依存MAE（CA-MAE）を導入する。

ABSTRACT

Featurizing microscopy images for use in biological research remains a significant challenge, especially for large-scale experiments spanning millions of images. This work explores the scaling properties of weakly supervised classifiers and self-supervised masked autoencoders (MAEs) when training with increasingly larger model backbones and microscopy datasets. Our results show that ViT-based MAEs outperform weakly supervised classifiers on a variety of tasks, achieving as much as a 11.5% relative improvement when recalling known biological relationships curated from public databases. Additionally, we develop a new channel-agnostic MAE architecture (CA-MAE) that allows for inputting images of different numbers and orders of channels at inference time. We demonstrate that CA-MAEs effectively generalize by inferring and evaluating on a microscopy image dataset (JUMP-CP) generated under different experimental conditions with a different channel structure than our pretraining data (RPI-93M). Our findings motivate continued research into scaling self-supervised learning on microscopy data in order to create powerful foundation models of cellular biology that have the potential to catalyze advancements in drug discovery and beyond.

研究の動機と目的

• 高含有スクリーニング顕微鏡（HCS）データで自己教師あり学習をスケールさせ、転用可能な細胞埋め込みを学習する。
• MAEを弱教師付き学習（WSL）およびImageNetプリトレーニングベースラインと生物学的関係の想起で比較する。
• 推論時に異なる入力チャネル構成を扱えるチャネル非依存MAE（CA-MAE）を開発・評価する。
• 異なるチャネルおよび実験条件を持つデータセットへのMAE表現の伝達性を評価する。
• 大規模バックボーンのためのMAEの安定性を改善するため、トレーニングダイナミクスと再構成損失を分析する。

提案手法

• 段階的に大きくなるHCSデータセット（RxRx、RPI-52M、RPI-93M）でMU-NetベースおよびViTベースのMAEを事前学習する。
• チャネルごとに自己標準化した入力として256x256のクロップを使用；8x8および16x16パッチと75%/25%のマスクを実験。
• 訓練を安定化させ、テクスチャ再構成を改善するためにフーリエ領域の再構成損失を組み込む（L_MAE+ = (1−alpha)L_MAE + alpha L_F、α = 0.01）。
• MAEトレーニングを伴うLarge-scale ViTエンコーダ（ViT-S/B/L）を採用し、最終層パッチ埋め込みを平均化して画像レベル表現を得る。
• CA-MAEを開発：チャネルを共有トークナイザと各チャネルデコーダを用いた別個のモダリティとして扱い、異なるチャネル数と順序での推論を可能にする。
• バッチ補正と堅牢な生物学的関係想起のためにTypical Variation Normalization（TVN）と染色体アームバイアス補正を適用する。

実験結果

リサーチクエスチョン

• RQ1 MAEをより大きなモデルバックボーンとより大規模な事前学習データセットでスケールさせると、HCSデータの既知の生物学的関係の想起は改善されるか？
• RQ2 フーリエ領域の再構成損失は大規模ViTベースMAEのトレーニングを安定化させ、下流の生物学的関係の想起を改善するか？
• RQ3 CA-MAEは推論時に異なるチャネル構成を持つ顕微鏡データセットへ一般化できるか？
• RQ4 MAEは弱教師付き法およびImageNetプリトレーニングモデルと比較して、摂動ベースの生物学的関係の想起においてどうか？
• RQ5 MAE由来の埋め込みはJUMP-CPなどの外部データセットへどの程度転移可能で、Perturbation/Sibling検索性能に影響するか？

主な発見

• MAEsは事前学習済みImageNetおよびほとんどのWSLモデルより優れており、最良のWSLモデルに対して相対的に最大11.5%の改善を実現（ViT-L/8+がRxRx1-2Mで訓練、ViT-L/16がRxRx1-2Mに訓練された場合の比較）。
• 生物学的関係の想起はトレーニングFLOPS（モデルサイズ×データセットサイズ）と共にスケールする；より大きなMAEとデータセットは、CORUM、hu.MAP、Reactome、StringDBなどのデータベースでより高い想起をもたらす。
• フーリエ領域の再構成損失を組み込むことで大規模MAE ViTのトレーニングが安定化し、損失景観のサドル点を越えることができる；頑健なダブルデセント挙動を示す。
• チャネル非依存MAEは異なるチャネル構成のデータセットへ一般化を可能にする；CA-MAE ViT-L/16+がRPI-93Mで強力なクロスデータセット想起を達成（例：CORUM 0.614、hu.MAP 0.424、Reactome 0.264、StringDB 0.478）。
• JUMP-CPへの転移はCA-MAEとMAEモデルで競争力のある摂動検索を示す；CA-MAEはCPJUMP1サブセットで平均適合度0.95の摂動検索を達成。
• MAE埋め込みは強力なWSLモデルよりもCellProfiler特徴の予測に優れており、形態的表現学習がより豊かであることを示す。