[论文解读] 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

一句话总结

论文表明：自监督掩码自编码器（MAEs）以 Vision Transformer 骨干网络扩展到大规模显微数据集，并在回忆已知生物关系方面优于弱监督方法；此外还提出了通道无关 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 中，alpha = 0.01）。
• 采用大规模 ViT 编码器（ViT-S/B/L）进行 MAE 训练，平均最终层补丁嵌入以获得图像级表示。
• 开发 CA-MAE：将通道视为独立模态，使用共享的分词器和逐通道解码器，使推断时能够处理不同通道数量和顺序。
• 应用典型变异归一化（TVN）和染色体臂偏置修正，以实现批量校正和稳健的生物关系召回。

实验结果

研究问题

• RQ1 当使用更大模型骨干和更大规模的预训练数据集来扩展 MAEs 时，是否会提升在 HCS 数据中对已知生物关系的回忆能力？
• RQ2 傅里叶域重构损失是否能稳定大规模 ViT 基 MAE 的训练并提升下游生物关系回忆？
• RQ3 通道无关的 MAE（CA-MAE）能否在推断时泛化到具有不同通道配置的显微数据集？
• RQ4 MAE 与弱监督方法和 ImageNet 预训练模型在回忆基于扰动的生物关系方面的比较？
• RQ5 MAE 派生嵌入在外部数据集如 JUMP-CP 上的可迁移性及其 Perturbation/Sibling 检索性能如何？

主要发现

• MAEs 超越预训练的 ImageNet 和大多数 WSL 模型，在回忆生物关系方面相对最强 WSL 模型最高提升达 11.5%（以 ViT-L/8+ 在 RPI-93M 上训练 vs ViT-L/16 在 RxRx1-2M 上训练）。
• 已知生物关系的回忆随训练 FLOPS（模型大小×数据集大小）而扩展；更大规模的 MAE 和数据集在各数据库（CORUM、hu.MAP、Reactome、StringDB）上获得更高回忆。
• 引入傅里叶域重构损失稳定大规模 MAE ViT 的训练，能够跨越损失面的鞍点，产生稳健的双下降行为。
• 通道无关的 MAEs 能使模型能够泛化到具有不同通道配置的数据集；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 子集上的 Perturbation 检索平均精度 0.95）。
• MAE 嵌入比强烈的 WSL 模型更能预测广泛的 CellProfiler 特征，表明更丰富的形态学表示学习。