[논문 리뷰] ViTally Consistent: Scaling Biological Representation Learning for Cell Microscopy

Kian Kenyon-Dean, Zitong Jerry Wang외 11명|arXiv (Cornell University)|2024. 11. 04.

Cell Image Analysis TechniquesBiochemistry, Genetics and Molecular Biology인용 수 3

한 줄 요약

이 논문은 Phenoprints-16M 데이터셋에서 수확한 80억 개의 정제된 현미경 영상 조각을 기반으로 사전 훈련한 19억 파rameter의 비전 트랜스포머(ViT-G/8) 마스크드 오토인코더를 소개한다. 이 모델은 생물학적 표현 학습에서 최신 기술 수준의 성능을 달성하며, 유전자 변형의 선형 가분성에서 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)로 사전 훈련함.
• 모든 트랜스포머 블록에서 표현을 평가하기 위해 생물학적으로 동기화된 선형 탐색 전략을 사용하여 하류 작업에 가장 우수한 성능을 보이는 중간 레이어를 선별함.
• 모델 일반화 및 생물학적 관련성을 향상시키기 위해 통계적으로 유의미한 양성 샘플의 정제된 데이터셋을 적용함.
• 유전자 그룹 분류 작업에 다중 과제 선형 탐색 프레임워크를 도입하여 전체 게놈 벤치마크와 강한 상관관계를 보이는 표현을 식별함.
• 100 GPU 시간 48,000시간 동안 훈련하여 10억 파rameter 규모를 달성함으로써 대규모에서의 강력한 표현 학습을 가능하게 함.
• 대규모 HCS 데이터에서 복제 일치성, 전체 게놈 관계 복구, 선형 가분성 지표를 사용하여 성능을 검증함.

실험 결과

연구 질문

• RQ1 정제된 현미경 데이터로 사전 훈련된 대규모 자기지도 학습 비전 트랜스포머 모델이 표현 학습에서 뛰어난 일치성과 생물학적 관련성을 달성할 수 있는가?
• RQ2 비전 트랜스포머의 중간 트랜스포머 블록이 세포 현미경 데이터에서 최종 레이어보다 더 생물학적으로 의미 있는 표현을 제공하는가?
• RQ3 모델 규모와 데이터 정제가 유전자 변형의 선형 가분성과 알려진 생물학적 관계 복구에 어떤 영향을 미치는가?
• RQ4 생물학적으로 동기화된 선형 탐색 작업이 대규모 전체 게놈 벤치마크 성능을 효과적으로 예측할 수 있는가?
• RQ5 양성 표현형 샘플의 정제된 데이터셋에 기반한 훈련이 하류 표현 품질에 어떤 영향을 미치는가?

주요 결과

• 제안된 MAE-G/8 모델은 이전의 ViT-L/8 MAE 대비 유전자 변형의 선형 가분성에서 60% 향상된 성능을 기록함.
• 전체 게놈 생물학적 관계 복구 및 복제 일치성 벤치마크에서 모두 최신 기술 수준의 성능을 달성함.
• 하류 생물학적 평가에서 중간 트랜스포머 블록이 최종 블록보다 일관되게 뛰어난 성능을 보이며, 더 관련성이 높은 생물학적 개념을 포착함을 시사함.
• 생물학적으로 동기화된 선형 탐색 작업 성능과 대규모 전체 게놈 벤치마크 성능 간에 강한 상관관계가 있음. 이는 효율적인 모델 평가를 가능하게 함.
• 훈련 FLOPs를 확장함에 따라 모든 벤치마크에서 뚜렷하고 일관된 선형 추세를 보이며 성능 향상이 확인됨. 이는 접근 방식의 확장 가능성을 확인함.
• 정제된 Phenoprints-16M 데이터셋은 알려진 유전자-유전자 관계의 복구 능력과 유전자 노크아웃 변형에 대한 임bedding 일치성 향상에 크게 기여함.