3D Masked Autoencoders are Robust Learners of Volumetric and Multimodal Cellular Representations for Microscopy

기술 요약: 세포 표현 학습을 위한 3D 마스크드 오토인코더 (3D Masked Autoencoders)

문제 정의
형광 현미경은 세포의 아세포 구조(subcellular organization)를 포착하는 풍부하고 고차원적인 데이터인 볼륨형 z-스택(예: OpenCell, WTC-11) 형태의 데이터를 제공한다. 그러나 Subcell 및 DINO4Cell과 같은 기존의 단일 세포 분석용 표현 학습 방식은 주로 2D 최대 강도 투영(maximum-intensity projection) 또는 개별 슬라이스 상에서 작동한다. 이러한 관행은 내재된 깊이 해상도 구조 정보를 유실하며, 이는 모델이 차별적인 세포 표현을 학습하는 능력을 제한할 수 있다. 또한, 순수 시각 모델은 단백질 서열에 인코딩된 생물학적 사전 지식(biological priors)과의 통합이 결여되어 있는 경우가 많다. 본 논문은 네이티브 3D 모델링과 단백질 언어 모델과의 멀티모달 정렬(multimodal alignment)이 이러한 한계를 극복하여 단백질 국소화(localization) 및 단백질-단백질 상호작용(PPI) 예측과 같은 다운스트림 태스크의 성능을 개선할 수 있는지 조사한다.

방법론
저자들은 볼륨형 현미경 데이터를 위한 마스크드 오토인코더(MAE) 프레임워크의 체계적인 비교 및 향상을 제안한다.

• 아키텍처 변형:

  • MAE-2D: 3D 볼륨을 2D 최대 강도 투영($I_{2D}$)으로 압축한다. 모델은 사인 함수 기반 위치 임베딩(sinusoidal positional embeddings)을 사용하는 2D ViT 백본을 통해 중첩되지 않는 2D 패치를 처리한다.
  • MAE-3D: 전체 3D 볼륨($I_{3D}$)을 유지하고 이를 3D 패치($p_z \times p_x \times p_y$)로 나눈다. 모델은 인코더-디코더 파이프라인 전반에 걸쳐 깊이 차원을 보존하는 3D 사인 함수 기반 위치 임베딩과 3D ViT 백본을 활용한다.
  • 채널 교차 주의 집중(Channel Cross-Attention, CCA): 채널 간 상호작용을 명시적으로 모델링하기 위해, 각 채널이 별도의 토큰 스트림이 되는 이중 스트림 인코더-디코더를 도입한다. 채널 전반에 걸쳐 공유된 무작위 마스킹 패턴이 적용된다. 채널 $c_i$의 공간 인덱스 $i$에 있는 토큰이 동일한 인덱스의 채널 $c_j$에 있는 토큰을 쿼리하도록 교차 주의 집중(cross-attention)이 구현된다. 단일 키(single-key) 제약을 처리하기 위해 소프트맥스(softmax) 연산 대신 시그모이드 게이팅(sigmoid gating) 메커니즘이 사용된다.
  • 주파수 영역 정규화(FFT Loss): 미세한 아세포 구조를 보존하기 위해, 원본 이미지와 재구성된 이미지의 크기 스펙트럼(magnitude spectra) 사이의 주파수 영역에서 L1 손실을 적용한다. 이 손실은 MAE-2D의 경우 2D 축에 대해, MAE-3D의 경우 세 개의 볼륨 축 모두에 대해 계산된다. 학습 초기 단계의 안정성을 위해 선형적으로 램프업(ramp up)된다.

• 멀티모달 정렬:
  본 프레임워크는 사전 학습된 단백질 언어 모델(ESM2)을 통합하여 시각적 특징을 단백질 서열 정보와 정렬한다.

  • 디코더 컨디셔닝(Decoder Conditioning): 동결된(frozen) ESM2 임베딩은 디코더 차원으로 투영된 후, 각 채널별 디코더 시퀀스에 단일 토큰으로 삽입된다(위치 임베딩은 0으로 설정). 이를 통해 마스킹된 패치들이 셀프 어텐션(self-attention)을 통해 단백질 문맥(context)을 참조할 수 있게 한다.
  • 대조 학습 목적 함수(Contrastive Objective): 대칭적 InfoNCE 손실($L_{CLIP}$)은 정규화된 이미지 임베딩을 단백질 임베딩과 정렬한다. 이 정렬은 가시적인 토큰(visible tokens)에 대해서만 계산되어 강건함을 촉진한다.
  • 전체 손실 함수: 최종 목적 함수는 MSE 재구성 손실, FFT 손실, 그리고 CLIP 손실의 조합이다: $L = L_{MSE} + w_{FFT}L_{FFT} + w_{CLIP}L_{CLIP}$.

핵심 기여

1. 체계적인 3D vs. 2D 비교: 저자들은 네이티브 3D MAE 모델이 2D 최대 투영 및 슬라이스 기반 변형 모델보다 모든 다운스트림 태스크에서 일관되게 우수한 성능을 보임을 입증하였다. 이는 2D 모델이 동일한 수의 슬라이스로 제한되었을 때도 마찬가지이다.
2. CCA 및 FFT를 통한 향상: 본 논문은 채널 교차 주의 집중과 주파수 영역 정규화가 3D 공간적 문맥을 활용하고 구조적 충실도를 보존하는 데 핵심적인 구성 요소임을 확인하였다.
3. 멀티모달 통합: 이미지 특징을 단백질 서열 임베딩(ESM2 활용)과 정렬하는 것이 표현 품질을 크게 향상시킨다는 점을 입증하였으며, 이러한 이득은 2D 모델보다 3D 모델에서 훨씬 더 크게 나타났다.

결과
실험은 OpenCell 데이터셋(1,310개 단백질, 약 6,300개 볼륨)을 사용하여 5-겹 교차 검증(five-fold cross-validation)을 통해 수행되었다.

• 단백질 국소화(Protein Localization):

  • MAE-3D vs. MAE-2D: MAE-3D는 모든 마스크 비율에서 MAE-2D보다 일관되게 우수한 성능을 보였다. 중앙 슬라이스 범위(45–55)로 제한하더라도 MAE-3D는 우월한 성능을 유지하였다 (AUCmicro: 0.91 vs. 0.90).
  • 멀티모달 이득: 가장 우수한 3D 모델(MAE-3D⋆)은 AUCmicro 0.952, F1micro 0.742를 달성하여, 이전 방식들보다 각각 +0.003 및 +0.10 증가한 SOTA(state-of-the-art) 결과를 기록하였다. ESM2의 통합은 2D 변형 모델에 비해 3D 모델에서 더 큰 성능 향상을 제공하였다.
  • 어텐션 분석: 시각화 결과, ESM2를 통합하면 관련 z-슬라이스 전반에 걸쳐 단백질 관련 구조(예: 액틴 세포 골격)에 주의 집중이 집중되는 반면, ESM2가 없거나 2D 모델인 경우 주의 집중이 분산되거나 깊이 특이성이 부족함을 보여주었다.

• 단백질-단백질 상호작용 (PPI):

  • MAE-3D⋆는 ROC-AUC 0.865를 달성하여, 이전 방법들(Subcell: 0.85, DINO4Cell: 0.84)보다 최대 +0.025 높은 성능을 보였다.
  • 주목할 점은, 3D 모델이 훨씬 더 큰 규모의 HPA 데이터셋에서 사전 학습된 모델들과 비교하여, 상대적으로 작은 규모의 OpenCell 데이터셋(~6k 볼륨)만으로도 이러한 결과를 달성했다는 점이며, 이는 볼륨형 모델링의 효율성을 강조한다.

의의 및 주장
본 논문은 전체 3D 볼륨 문맥을 보존하는 것이 2D 투영보다 단일 세포 현미경에 대해 더 정보가 풍부하고 차별적인 표현을 생성한다는 점을 주장한다. 저자들은 볼륨 형태(volumetric morphology)와 단백질 수준의 의미론(semantics) 사이의 시너지는 네이티브 3D 아키텍처를 사용할 때 가장 강력하다고 상정한다. 저자들은 네이티브 3D 모델링과 멀티모달 정렬의 결합이 세포 이미징을 위한 강력한 파운데이션 모델을 개발하는 데 필수적이라고 결론짓는다. 또한, 이러한 이점이 대규모 사전 학습 없이도 유효하다는 점을 강조하며, 아키텍처 선택(3D vs. 2D)과 멀티모달 통합이 성능의 핵심 요소임을 시사한다.