[논문 리뷰] Masked Autoencoders for Microscopy are Scalable Learners of Cellular Biology

Oren Kraus, Kian Kenyon-Dean외 15명|arXiv (Cornell University)|2024. 04. 16.

Cell Image Analysis Techniques인용 수 10

한 줄 요약

논문은 Vision Transformer 백본을 갖춘 self-supervised masked autoencoders (MAEs)가 대규모 현미경 데이터세트로 확장되며 알려진 생물학적 관계를 회상하는 데 weakly supervised 방법을 능가함을 보여주고; 또한 서로 다른 채널 구성에 대해 일반화하는 채널-비특이적 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) 데이터에서 자기지도 학습을 확장하여 전이 가능한 세포 임베딩을 학습한다.
• MAEs와 weakly supervised learning(WSL) 및 ImageNet- pretrained baselines를 생물학적 관계 회상에서 비교한다.
• 추론 시 서로 다른 입력 채널 구성을 처리하는 채널-비특이적 MAE (CA-MAE)를 개발하고 평가한다.
• MAE 표현의 외부 데이터셋(다른 채널 및 실험 조건)에 대한 이전 가능성을 평가한다.
• 대규모 백본에 대한 MAE의 안정성을 높이기 위한 학습 역학 및 재구성 로스 분석을 수행한다.

제안 방법

• 점진적으로 더 큰 HCS 데이터세트(RxRx, RPI-52M, RPI-93M)에서 MU-Net 기반 및 ViT 기반 MAEs를 사전 학습한다.
• 채널별 자기 표준화를 입력으로 사용하고 256x256 크롭을 활용하며 8x8 및 16x16 패치와 75%/25% 마스크를 실험한다.
• Training의 안정성을 높이고 질감 재구성을 개선하기 위해 Fourier 도메인 재구성 로스를 도입한다 (alpha = 0.01 in L_MAE+ = (1−alpha)L_MAE + alpha L_F).
• MAE 학습을 위해 Large-scale ViT 인코더(ViT-S/B/L)를 채택하고 최종 레이어 패치 임베딩을 이미지 수준 표현으로 평균화한다.
• CA-MAE를 개발: 채널을 별도의 모달리티로 간주하고 공유 토크나이저와 채널별 디코더를 사용하여 서로 다른 채널 수와 순서로 추론이 가능하게 한다.
• Typical Variation Normalization(TVN) 및 염색체-팔 바이어스 보정을 적용하여 배치 보정 및 생물학적 관계 회상을 견고하게 한다.

실험 결과

연구 질문

• RQ1 더 큰 모델 백본과 더 큰 사전 학습 데이터가 HCS 데이터에서 알려진 생물학적 관계의 회상을 향상시킬까?
• RQ2 Fourier 도메인 재구성 로스가 학습을 안정화하고 대형 ViT 기반 MAE의 다운스트림 생물학적 관계 회상을 개선할까?
• RQ3 CA-MAE가 추론 시 서로 다른 채널 구성에 일반화될 수 있을까?
• RQ4 MAEs와 weakly supervised 방법 및 ImageNet- pretrained 모델이 교란(perturbation) 기반 생물학적 관계를 회상하는 데 얼마나 차이가 있을까?
• RQ5 MAE 유래 임베딩의 외부 데이터셋(JUMP-CP)으로의 전이 가능성과 Perturbation/Sibling 검색 성능은 어떤가?

주요 결과

• MAEs는 ImageNet 사전 학습 및 대부분의 WSL 모델보다 우수하며, 생물학적 관계 회상에서 최고의 WSL 모델 대비 최대 11.5% 상대 향상을 달성한다(ViT-L/8+가 RPI-93M에서 학습된 반면 ViT-L/16은 RxRx1-2M에서 학습).
• 생물학적 관계의 회상은 학습 FLOPS(모델 크기 × 데이터세트 크스)와 함께 확대되며, 더 큰 MAEs와 더 큰 데이터세트가 데이터베이스(CORUM, hu.MAP, Reactome, StringDB) 전반에서 더 높은 회상을 보인다.
• Fourier 도메인 재구성 로스를 도입하면 큰 MAE ViT의 학습이 안정화되고 로스 지형에서 샤들을 넘어선 crossing이 가능하며 강건한 이중 하강(double-descent) 동작을 보인다.
• 채널-비특이적 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 모델에서 경쟁력 있는 perturbation 검색을 보여주며; CA-MAE는 Perturbation 검색에서 CPJUMP1 부분집합에 대해 평균 정밀도 0.95와 같이 높은 성능을 달성한다.
• MAE 임베딩은 강력한 WSL 모델보다 CellProfiler 특징의 더 넓은 범위를 더 잘 예측하여 형태학적 표현 학습이 더 풍부함을 시사한다.