[论文解读] 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），特别是基于视觉Transformer的模型，能够有效扩展至大规模高内涵筛选（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.

研究动机与目标

• 探究通过掩码自编码器进行自监督学习是否能够有效扩展至大规模高内涵筛选（HCS）数据集，以实现细胞形态表征学习。
• 比较基于MAE的模型与弱监督学习基线在全基因组HCS筛选中推断已知生物学关系的性能表现。
• 评估模型规模与数据集规模对细胞图像分析中表征质量的影响。
• 利用大规模HCS数据建立自监督细胞表征学习的新SOTA（最先进）水平。

提出的方法

• 在四个HCS数据集上训练掩码自编码器（MAE），包括包含9300万张专有显微镜图像的RPI-93M数据集，以在无需人工标注标签的情况下学习表征。
• 采用视觉Transformer（ViT）和U-Net架构进行掩码自编码，通过未掩码的图像块重建被掩码的图像块。
• 采用大批次训练，使用Lion优化器，并对学习率、权重衰减和随机深度等超参数进行优化，以实现最佳收敛与性能。
• 采用确定性推理流程：对于每张2048×2048×6的孔板图像，生成64个非重叠的256×256×6图像块，通过MAE编码器分别编码，最终对得到的嵌入向量取平均，生成最终表征。
• RPI-52M和RPI-93M使用全局批量大小16,384，RxRx3使用4,096，采用循环余弦学习率调度，且不使用梯度裁剪。
• 在下游任务中，使用整理的数据库（CORUM、hu.MAP、Reactome、StringDB）评估模型对已知生物学关系的召回性能，并与随机、浅层特征及弱监督基线进行比较。

实验结果

研究问题

• RQ1 掩码自编码器是否能够有效扩展至大规模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和弱监督模型。
• 与使用小批量的AdamW相比，采用Lion优化器和大批次训练显著改善了训练动态与下游性能。
• 随机和浅层特征基线仅能召回约10%的已知关系，证实了MAE学习表征的优越性。
• 使用ImageNet预训练权重并未带来收益；从随机初始化开始训练的模型表现更优，表明自然图像与HCS数据之间存在显著的领域差异。