[論文レビュー] Masked Autoencoders are Scalable Learners of Cellular Morphology

Oren Kraus, Kian Kenyon-Dean+14 more|arXiv (Cornell University)|Sep 27, 2023

Cell Image Analysis Techniques被引用数 6

ひとこと要約

この論文は、特にビジョントランスフォーマーを用いたマスクド自己符号化器（MAE）が、大規模なハイコントtentスクリーニング（HCS）データセットから細胞の形態を学習する際に効果的にスケーリングできることを示している。9300万枚以上の顕微鏡画像で訓練されたViT-L/8 MAEは、弱い教師ありベースラインと比較して、既知の生物学的関係を推論する際、最大28%の相対的改善を達成した。これは、MAEが細胞フェノミクスのための強力でスケーラブルな自己教師付き学習者として確立されたことを示している。

ABSTRACT

Inferring biological relationships from cellular phenotypes in high-content microscopy screens provides significant opportunity and challenge in biological research. Prior results have shown that deep vision models can capture biological signal better than hand-crafted features. This work explores how self-supervised deep learning approaches scale when training larger models on larger microscopy datasets. Our results show that both CNN- and ViT-based masked autoencoders significantly outperform weakly supervised baselines. At the high-end of our scale, a ViT-L/8 trained on over 3.5-billion unique crops sampled from 93-million microscopy images achieves relative improvements as high as 28% over our best weakly supervised baseline at inferring known biological relationships curated from public databases. Relevant code and select models released with this work can be found at: https://github.com/recursionpharma/maes_microscopy.

研究の動機と目的

• 自己教師付き学習としてのマスクド自己符号化器（MAE）が、細胞形態の表現学習のための大規模ハイコントtentスクリーニング（HCS）データセットに、効果的にスケーリング可能かどうかを調査すること。
• 全ゲノムHCSスクリーンからの既知の生物学的関係を推論する際、MAEベースのモデルと弱い教師あり学習ベースラインの性能を比較すること。
• モデルサイズとデータセット規模が、細胞画像解析における表現品質に与える影響を評価すること。
• 大規模HCSデータを用いて、自己教師付き細胞表現学習の新しい最良状態を確立すること。

提案手法

• 人為的ラベルなしで表現を学ぶために、RPI-93M（9300万枚の独自の顕微鏡画像を含む）を含む4つのHCSデータセットでマスクド自己符号化器（MAE）を訓練した。
• 画像パッチのマスク化と、未マスク化パッチからの再構成を実行するため、ビジョントランスフォーマー（ViT）およびU-Netアーキテクチャを用いたマスクド自己符号化を採用した。
• 収束と性能最適化のため、学習率、正則化係数、確率的深さのハイパーパramータを最適化した大バッチ学習を実施した。
• 決定論的な推論パイプラインを採用：各2048×2048×6のウェル画像に対して、64個の重複しない256×256×6のクロップを生成し、それぞれをMAEエンコーダーで符号化し、得られた埋め込みを平均化して最終表現を算出した。
• RPI-52MおよびRPI-93Mではグローバルバッチサイズ16,384、RxRx3では4,096を採用し、勾配クリッピングなしでサイクルコサイン学習率を適用した。
• キュレートされたデータベース（CORUM、hu.MAP、Reactome、StringDB）を用いて、既知の生物学的関係の再検索性能を評価し、ランダム、浅い特徴量、弱い教師ありベースラインと比較した。

実験結果

リサーチクエスチョン

• RQ1 マスクド自己符号化器（MAE）は、教師なし条件下で、大規模なHCSデータセットに効果的にスケーリングされ、意味のある細胞形態の表現を学習可能か？
• RQ2 全ゲノムHCSスクリーンにおいて、MAEベースのモデルは弱い教師あり学習ベースラインと比較して、既知の生物学的関係をどれほど効果的に推論できるか？
• RQ3 モデルサイズとデータセット規模の増加は、細胞表現学習において一貫した性能向上をもたらすか？
• RQ4 HCSデータに対するMAEの最適なパフォーマンスを達成するためのトレーニング設定（最適化手法、バッチサイズ、学習率）は何か？

主な発見

• 9300万枚の顕微鏡画像から得られる35億個以上のユニークな画像クロップを用いて、ViT-L/8 MAEを学習したところ、最良の弱い教師ありベースラインと比較して、既知の生物学的関係の再検索で最大28%の相対的改善が達成された。
• MAEベースのモデル、特にViTバージョンは、モデルサイズとデータセット規模の増加に伴い、一貫して性能向上を示し、強力なスケーラビリティを示した。
• ViT-L/8モデルは、CORUM、hu.MAP、Reactome、StringDBのすべてのベンチマークデータベースで、CNNベースのMAEや弱い教師ありモデルをすべて上回った。
• Lion最適化手法と大バッチサイズを用いた学習は、AdamWと小バッチサイズを用いた場合と比較して、トレーニングのダイナミクスと下流性能を顕著に改善した。
• ランダムおよび浅い特徴量ベースラインは、既知の関係の約10%しか再検索できず、MAEから得られる学習済み表現の優位性を確認した。
• ImageNetの重みからの事前学習は効果がなく、ランダム初期化からの学習が優れたパフォーマンスを達成した。これは、自然画像とHCSデータの間でドメインシフトが生じていることを示唆している。