셀제파: 단일 세포 전사체를 위한 잠재 표현 학습 모델

초록

Cell‑JEPA는 마스크된 유전자 발현을 직접 복원하는 대신, 마스크된 입력으로부터 교사 네트워크가 만든 셀 임베딩을 예측하도록 설계된 학생‑교사 구조를 도입한다. 이를 통해 90% 이상인 드롭아웃을 가진 희소 데이터에서도 잡음에 강한 세포 정체성 표현을 학습한다. zero‑shot 세포 유형 클러스터링에서 AvgBIO 0.72를 기록해 scGPT(0.53)보다 36% 상대적 향상을 보였으며, 단일 세포 라인 교란 예측에서는 절대 상태 재구성은 개선했지만 효과 크기 추정은 기존 모델과 비슷했다.

상세 분석

Cell‑JEPA는 기존 단일 세포 foundation model인 scGPT의 구조를 그대로 유지하면서, JEPA(Joint Embedding Predictive Architecture) 개념을 도입해 두 개의 트랜스포머 인코더(학생, 교사)를 병렬로 운영한다. 교사 인코더는 마스크되지 않은 전체 토큰을 입력받아 안정적인 셀‑레벨 임베딩( 토큰)을 생성하고, 학생 인코더는 마스크된 토큰만을 사용해 동일한 임베딩을 예측한다. 교사 파라미터는 학생 파라미터의 지수 이동 평균(EMA)으로 업데이트되므로, 학습 초기에 교사는 빠르게 변하고 후기에 점점 안정화된다. 이 설계는 “표현‑레벨” 예측 손실(L_JEPA)을 도입함으로써, 모델이 원시 카운트값이 아닌 고차원 잠재 공간에서 일관성을 유지하도록 강제한다.

기술적으로는 두 손실을 가중합한다. 첫 번째는 기존 scGPT와 동일한 유전자‑레벨 재구성 손실(L_rec)로, 마스크된 유전자의 이산화된 발현값을 MSE로 복원한다. 두 번째는 코사인 유사도 기반 L_JEPA로, 학생이 예측한 임베딩 ˜e와 교사의 고정된 임베딩 e 사이의 코사인 거리를 최소화한다. L_JEPA는 stop‑gradient를 적용해 교사 임베딩이 역전파에 영향을 받지 않게 함으로써, 학생이 잡음이 많은 입력에서도 의미 있는 셀 정체성을 학습하도록 만든다.

데이터 전처리 단계에서도 중요한 설계가 있다. 각 셀은 비제로 유전자만을 추출해 (gene_id, expression) 쌍의 희소 시퀀스로 표현하고, 표현값은 셀별 50개의 분위수 구간으로 이산화한다. 또한, 최대 600개의 유전자를 무작위로 서브샘플링해 시퀀스 길이를 일정하게 맞추면서, 학습 에폭마다 다른 부분 집합을 관찰하도록 하여 자연스러운 드롭아웃 효과를 만든다. 이러한 전처리는 모델이 다양한 부분 관측에 대해 강건하게 학습하도록 돕는다.

실험 결과는 세 가지 다운스트림 태스크에서 검증된다. (i) 지도 학습 파인튜닝을 통한 세포 유형 클러스터링에서는 Cell‑JEPA가 기존 scGPT 대비 평균 Adjusted Rand Index와 NMI 모두에서 유의미하게 상승했다. (ii) zero‑shot 전이 설정에서는 사전 학습된 임베딩만으로 K‑Nearest Neighbor 분류를 수행했을 때 AvgBIO 점수가 0.72로, scGPT의 0.53을 크게 앞섰다. 이는 사전 학습 단계에서 얻은 잠재 공간이 실제 생물학적 구조를 잘 반영한다는 증거다. (iii) 단일 세포 라인 교란 예측에서는 절대적인 발현 재구성(RMSE)에서는 개선이 있었지만, 교란에 따른 효과 크기(Δ expression) 추정 정확도는 scGPT와 비슷했다. 이는 표현 학습이 잡음 억제와 전반적 상태 복원에 강점을 보이지만, 미세한 교란 효과를 정량화하려면 별도의 교란‑특화 모델이 필요함을 시사한다.

전체적으로 Cell‑JEPA는 “표현‑레벨 예측”이라는 새로운 자기지도 학습 목표를 도입해, 고드롭아웃, 고희소성 데이터를 다루는 단일 세포 전사체 분석에 적합한 견고한 임베딩을 제공한다. 이는 향후 다양한 조직, 종, 그리고 교란 실험에 걸친 전이 학습에 유용한 기반이 될 것으로 기대된다.