[論文レビュー] ViTally Consistent: Scaling Biological Representation Learning for Cell Microscopy

Kian Kenyon-Dean, Zitong Jerry Wang+11 more|arXiv (Cornell University)|Nov 4, 2024

Cell Image Analysis TechniquesBiochemistry, Genetics and Molecular Biology被引用数 3

ひとこと要約

本論文では、Phenoprints-16Mデータセットから収集された80億の精錬済み顕微鏡画像クロップを用いて事前学習された19億パラメータのビジョントランスフォーマー（ViT-G/8）マスクド自己符号化器（MAE）であるViTally Consistentを紹介する。生物学的表現学習において最先端の性能を達成し、遺伝子ペルトゥルベーションの線形分離性が60%向上し、優れたレプリケートの一貫性と全ゲノム関係の再検出性能を示しており、中間トランスフォーマーブロックが最終層よりも生物学的に意味のある特徴を抽出していることが判明した。

ABSTRACT

Large-scale cell microscopy screens are used in drug discovery and molecular biology research to study the effects of millions of chemical and genetic perturbations on cells. To use these images in downstream analysis, we need models that can map each image into a feature space that represents diverse biological phenotypes consistently, in the sense that perturbations with similar biological effects have similar representations. In this work, we present the largest foundation model for cell microscopy data to date, a new 1.9 billion-parameter ViT-G/8 MAE trained on over 8 billion microscopy image crops. Compared to a previous published ViT-L/8 MAE, our new model achieves a 60% improvement in linear separability of genetic perturbations and obtains the best overall performance on whole-genome biological relationship recall and replicate consistency benchmarks. Beyond scaling, we developed two key methods that improve performance: (1) training on a curated and diverse dataset; and, (2) using biologically motivated linear probing tasks to search across each transformer block for the best candidate representation of whole-genome screens. We find that many self-supervised vision transformers, pretrained on either natural or microscopy images, yield significantly more biologically meaningful representations of microscopy images in their intermediate blocks than in their typically used final blocks. More broadly, our approach and results provide insights toward a general strategy for successfully building foundation models for large-scale biological data.

研究の動機と目的

• 細胞顕微鏡画像に対して、実験レプリケート間で生物学的に意味のある一貫性のある表現を学習する大規模な基盤モデルを開発すること。
• 高スループットスクリーニングデータにおける遺伝子ペルトゥルベーションの線形分離性と全ゲノム生物学的関係の再検出性能を向上させること。
• 最終層に依存するのではなく、下流の生物学的評価に最適な中間層を特定すること。
• 自己教師付き学習における生物学的画像処理のスケーラブルで生物学的に根拠のあるトレーニングおよびプロービング戦略を確立すること。
• モデル規模と精錬済みデータが、表現学習における一貫性と生物学的関連性を顕著に向上させることを示すこと。

提案手法

• 精錬済みのPhenoprints-16Mデータセットから得た80億以上の顕微鏡画像クロップを用いて、ビジョントランスフォーマー（ViT-G/8）をマスクド自己符号化器（MAE）として事前学習する。
• 全トランスフォーマーブロックにおける表現の評価に生物学的に根拠のある線形プロービング戦略を用い、下流タスクに最適な中間層を選択する。
• 統計的に有意な陽性サンプルの精錬済みデータセットを用いて、モデルの一般化性能と生物学的関連性を向上させる。
• 遺伝子群分類タスクにおけるマルチタスク線形プロービングフレームワークを採用し、全ゲノムベンチマークと強い相関を示す表現を同定する。
• 48,000 H100 GPU時間のトレーニングにより、数十億パラメータ規模を達成し、大規模スケールでの強固な表現学習を可能にする。
• 大規模HCSデータにおいて、レプリケートの一貫性、全ゲノム関係再検出性能、線形分離性のメトリクスを用いて性能を検証する。

実験結果

リサーチクエスチョン

• RQ1 精錬済み顕微鏡画像データで事前学習された大規模自己教師付きビジョントランスフォーマー・モデルは、表現学習における優れた一貫性と生物学的関連性を達成できるか？
• RQ2 ビジョントランスフォーマーにおける中間トランスフォーマーブロックは、細胞顕微鏡画像データにおいて最終層よりも生物学的に意味のある表現を生成するか？
• RQ3 モデル規模とデータ精錬が、遺伝子ペルトゥルベーションの線形分離性と既知の生物学的関係の再検出に与える影響は何か？
• RQ4 生物学的に根拠のある線形プロービングタスクは、大規模全ゲノムベンチマークにおける性能を効果的に予測できるか？
• RQ5 陽性表現型サンプルの精錬済みデータセットでのトレーニングが、下流の表現品質に与える影響は何か？

主な発見

• 提案されたMAE-G/8モデルは、以前のViT-L/8 MAEと比較して、遺伝子ペルトゥルベーションの線形分離性で60%の向上を達成した。
• モデルは、全ゲノム生物学的関係再検出とレプリケート一貫性ベンチマークの両方で最先端の性能を達成した。
• 下流の生物学的評価において、中間トランスフォーマーブロックが常に最終ブロックを上回り、より関連性の高い生物学的概念を捉えていることが示された。
• 生物学的に根拠のある線形プロービングタスクの性能は、大規模全ゲノムベンチマークの性能と強く相関しており、効率的なモデル評価が可能である。
• トレーニングFLOPsのスケーリングは、すべてのベンチマークで改善された性能と強く一貫した線形的トレンドを示し、本手法のスケーラビリティを裏付けた。
• 精錬済みのPhenoprints-16Mデータセットは、既知の遺伝子-遺伝子関係の再検出性能と遺伝子ノックアウトペルトゥルベーションの埋め込み一貫性の両方を顕著に向上させた。