[논문 리뷰] Masked Autoencoders are Scalable Learners of Cellular Morphology

Oren Kraus, Kian Kenyon-Dean외 14명|arXiv (Cornell University)|2023. 09. 27.

Cell Image Analysis Techniques인용 수 6

한 줄 요약

이 논문은 특히 비전 트랜스포머 기반 모델인 마스킹된 오토인코더(MAE)가 거대한 고내용 스크리닝(HCS) 데이터셋으로부터 세포 형태를 효과적으로 학습하는 데에 스케일링될 수 있음을 보여준다. 9,300만 장이 넘는 현미경 영상로 훈련된 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)가 대규모 고내용 스크리닝(HCS) 데이터셋에 대해 감독 없이 세포 형태 표현을 효과적으로 학습할 수 있는지 조사하기.
• 전체 게놈 HCS 스크리너에서 알려진 생물학적 관계를 추론하는 데 있어 MAE 기반 모델과 약한 지도 학습 기반 모델 간의 성능을 비교하기.
• 모델 크기와 데이터셋 규모가 세포 영상 분석에서 표현 품질에 미치는 영향을 평가하기.
• 대규모 HCS 데이터를 사용하여 자기지도 세포 표현 학습 분야에서 새로운 최고 성능 기록을 수립하기.

제안 방법

• 인간 레이블이 없는 상태에서 표현을 학습하기 위해, RPI-93M(9,300만 개의 전용 현미경 영상 포함)를 포함한 네 개의 HCS 데이터셋에서 마스킹된 오토인코더(MAE)를 훈련함.
• 마스킹된 이미지 패치를 비마스킹된 패치로부터 복원하기 위해 비전 트랜스포머(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 마스킹된 오토인코더(MAE)는 감독 없이 대규모 HCS 데이터셋에서 의미 있는 세포 형태 표현을 효과적으로 학습할 수 있는가?
• RQ2 전체 게놈 HCS 스크리너에서 MAE 기반 모델은 약한 지도 학습 기반 모델 대비 알려진 생물학적 관계를 어떻게 추론하는가?
• RQ3 모델 크기와 데이터셋 규모를 증가시키면 세포 표현 학습에서 일관된 성능 향상이 이루어지는가?
• RQ4 HCS 데이터에서 MAE의 최적 성능를 얻기 위한 훈련 구성(최적화기, 배치 크기, 학습률)은 무엇인가?

주요 결과

• 9,300만 장의 현미경 영상에서 유일한 이미지 컷 35억 개 이상를 사용해 훈련된 ViT-L/8 MAE는 최고의 약한 지도 학습 기반 모델 대비 알려진 생물학적 관계 재확인에서 최대 28%의 상대적 향상을 달성함.
• MAE 기반 모델, 특히 ViT 변종의 성능은 모델 크기와 데이터셋 규모가 증가함에 따라 일관되게 향상되어 강력한 확장성 잠재력을 보임.
• ViT-L/8 모델은 CORUM, hu.MAP, Reactome, StringDB를 포함한 모든 벤치마크 데이터베이스에서 CNN 기반 MAE 및 약한 지도 학습 모델을 모두 압도함.
• Lion 최적화기와 대용량 배치를 사용한 훈련은 AdamW와 작은 배치를 사용한 경우에 비해 훈련 동역학과 후행 성능을 크게 향상시킴.
• 무작위 및 얕은 특징 기반 모델은 알려진 관계의 약 10%만 재확인함을 확인하여, MAE에서 학습된 표현의 우수성을 입증함.
• ImageNet 가중치에서의 사전 훈련은 도움이 되지 않았으며, 무작위 초기화에서 훈련된 모델이 더 뛰어난 성능을 보였으며, 자연 이미지와 HCS 데이터 간의 도메인 차이가 있음을 시사함.