기술 요약: 단일 세포 표현 학습을 위한 프로토타입 유도 사전-후학습 (Prototype Guided Post-pretraining)

문제 정의

단일 세포 표현 학습 (SCRL) 은 고차원 유전자 발현 데이터를 해독하여 세포 조절 논리를 규명하는 것을 목표로 합니다. 최근 대규모 언어 모델에서 영감을 받아 유전자를 토큰으로, 세포를 문장으로 간주하는 기반 모델 (foundation models) 이 가능성을 보였으나, 두 가지 근본적인 한계에 직면해 있습니다:

1. 긴 꼬리 (Long-Tailed) 세포 유형 분포: 단일 세포 데이터셋은 극심한 클래스 불균형을 보이며, 드물지만 생물학적으로 중요한 세포 집단 (예: 질병을 유발하는 줄기 세포) 은 샘플의 극히 작은 부분만을 차지합니다. 이는 모델이 풍부한 집단에 편향됨에 따라 드문 세포 유형에 대해 노이즈가 많거나 붕괴된 임베딩을 초래합니다.
2. 공변량 이동 (Covariate Shifts): 배치 효과 (batch effects) 와 다양한 시퀀싱 깊이와 같은 기술적 인공물은 훈련 데이터와 추론 데이터 간의 분포 격차를 만듭니다. 기존 모델들은 종종 보편적인 조절 논리를 포착하기보다는 배치별 패턴에 과적합되어, 서로 다른 플랫폼이나 실험 프로토콜 간 일반화 능력이 제한됩니다.

현재의 미세 조정 (fine-tuning) 전략들은 초기 사전 학습 단계에서 확립된 불안정한 잠재 임베딩 매니폴드에서 시작하기 때문에 이러한 문제를 완전히 해결하지 못하는 경우가 많습니다.

방법론: CellRefine

이러한 과제를 해결하기 위해 저자들은 초기 사전 학습 단계와 하류 미세 조정 사이에 작동하도록 설계된 CellRefine이라는 사후-사전 학습 (post-pretraining) 프레임워크를 소개합니다. CellRefine 은 생물학적으로 근거한 귀납적 편향을 주입하여 잠재 임베딩 공간을 더 안정적이고 차별화된 구조로 재구성합니다.

이 방법은 네 가지 구성 요소를 결합한 다면적 목적 함수 ($L_{Total}$) 를 사용합니다:

1. 프로토타입 유도 학습 (마커 유전자 프로그램):

   • 개념: 알려진 마커 유전자 집합을 구조적 사전 지식으로 활용합니다. 완전한 마커 집합은 종종 알려지지 않았으므로, 이 방법은 각 세포 유형에 대한 "프로토타입"을 생성하기 위해 선별된 마커 하위 집합을 사용합니다.
   • 구현: 세포 유형 $c_i$에 대한 마커 유전자는 세포 계통도 (ontology) 계층 구조 (구체적인 아형에서 더 넓은 계통으로) 를 기반으로 함수 $H(c_i)$를 사용하여 정렬됩니다. 이러한 정렬된 시퀀스는 토큰화되어 프로토타입 임베딩 ($z_m$) 으로 인코딩됩니다.
   • 손실 (Loss): 프로토타입 유도 정규화 손실 ($L_{prototype}$) 은 세포의 임베딩 ($z_c$) 과 해당 프로토타입 임베딩 간의 코사인 유사도를 최대화하면서 다른 프로토타입과의 유사도를 최소화합니다. 이를 통해 표현에 생물학적 구조를 주입합니다.

2. 계통 인식 정규화 (Lineage-Aware Regularization):

   • 개념: 계통도에서 공통 조상을 공유하는 밀접하게 관련된 세포 유형들이 임베딩 공간에서 충분히 분리되지 않는 문제를 해결합니다.
   • 구현: 계통 인식 분리 손실 ($L_{lineage}$) 은 동일한 부모 계통과 계통도 수준을 공유하는 세포 쌍 간의 과도한 유사도에 페널티를 부과하여, 모델이 세밀한 생물학적 정체성을 구분하도록 장려합니다.

3. 가우시안 혼합 변분 인코딩 (GMVE):

   • 개념: 표현 품질을 향상시키기 위해 잠재 공간을 클러스터링된 구조로 명시적으로 모델링합니다.
   • 구현: 가우시안 혼합 변분 인코더가 모델에 추가됩니다. KL 발산 손실 ($L_{GMVE}$) 은 학습된 가우시안 혼합 사전 분포에 대해 변분 사후 분포를 정규화하여 클러스터링된 잠재 구조를 강제합니다.

4. 마스크 언어 모델링 (MLM):

   • 모델이 일반적인 유전자 발현 패턴을 계속 학습하도록 보장하기 위해 표준 사전 학습 목적이 유지됩니다.

최종 목적 함수는 다음과 같습니다:
$$L_{Total} = L_{MLM} + \lambda_1 L_{prototype} + \lambda_2 L_{lineage} + \lambda_3 L_{GMVE}$$

주요 기여

• 새로운 사후-사전 학습 단계: 사전 학습 후 직접적인 미세 조정에만 의존하는 대신, 작업별 미세 조정 전에 기반 모델의 잠재 매니폴드를 정제하는 중간 적응 단계 (CellRefine) 를 제안합니다.
• 생물학적 사전 지식 통합: 마커 유전자 프로그램과 세포 계통도와 같은 기계론적 생물학적 지식을 체계적으로 최적화 과정에 직접 통합하여 표현 학습을 유도합니다.
• 긴 꼬리 및 공변량 이동 처리: 생물학적 사전 지식에 기반을 둔 모델이 풍부한 세포 유형에 대한 편향을 완화하고 기술적 이동에 대한 강건성을 향상시킴을 입증합니다.

실험 결과

저자들은 혈액, 췌장, 간, 뇌, 폐, 심장 등 다양한 조직에 걸친 11 개의 데이터셋에서 3 가지 핵심 계산 생물학 과제에 대해 CellRefine 을 평가했습니다:

1. 세포 정체성 예측:

   • 동역역 (On-Domain): CellRefine 은 10 개 데이터셋에서 최신 기준 모델 (전체 미세 조정, LoRA, 표준 MLM 정렬 포함) 일관적으로 능가했습니다. 혈액 (Blood) 과 췌장 (Pancreas) 과 같은 불균형 데이터셋에서 가장 강력한 기준 모델 (MLM $\to$ FF) 대비 매크로-F1 점수가 최대 **15%**까지 향상되었습니다.
   • 이역역 (Out-of-Domain, Zero-Shot): 교차 모달 전이 (scRNA-seq 에서 훈련, 공간 전사체학에서 테스트) 에서 CellRefine 은 더 강건하고 일반화 가능한 표현을 나타내는 우수한 recall@k 성능을 보여주었습니다.

2. 공간 전사체학 보간 (Spatial Transcriptomic Imputation):

   • CellRefine 은 측정되지 않은 유전자 발현 예측을 개선하여 간 공간 데이터셋에서 기준 모델 대비 더 높은 피어슨 상관관계 및 코사인 유사도 점수를 달성했습니다.

3. 유전자 교란 반응 예측:

   • 교란 하의 유전자 발현 변화 예측에서 CellRefine 은 기준 모델 대비 일관되게 (비록 작지만) 개선을 보였으며, 테스트된 모든 방법 중 전체적으로 가장 좋은 성능을 달성했습니다.

Ablation 연구:

• 세 가지 정규화 항 ($L_{prototype}$, $L_{lineage}$, $L_{GMVE}$) 중 하나를 제거하면 성능이 저하되어 다목적 접근법이 시너지 효과를 낸다는 것이 확인되었습니다.
• 샘플 효율성: 퓨샷 (few-shot) 학습 시나리오 (5~50 샷) 에서 CellRefine 은 표준 미세 조정에 비해 낮은 샘플 수에서 훨씬 더 가파른 성능 향상을 보여주어, 제한된 레이블 데이터를 더 효과적으로 활용함을 시사합니다.

정성적 분석:

• 잠재 공간의 UMAP 시각화는 CellRefine 이 임베딩 매니폴드를 재구성하여, MLM 만으로 훈련된 모델에 비해 드문 "꼬리" 세포 집단 (예: 특정 T 세포 아형, 드문 간 세포 유형) 을 더 구별 가능하게 만들고 주요 세포 클러스터와의 얽힘을 줄였음을 보여줍니다.

중요성 및 주장

이 논문은 생물학적 정보에 기반한 귀납적 편향이 기반 모델의 하류 성능을 크게 향상시킬 수 있음을 입증함으로써, 단일 세포 표현 학습을 위한 유망한 방향을 제시한다고 주장합니다.

• 임상적 유용성: 드문 세포 유형을 더 잘 포착하고 공변량 이동에 대한 강건성을 향상시킴으로써, 표준 모델에서 종종 가려지는 드물거나 과도기적인 세포 상태의 질병 메커니즘 및 치료 표적을 식별할 가능성을 높입니다.
• 방법론적 전환: 데이터 기반 사전 학습 후 미세 조정에만 의존하는 것을 반대하며, 학습 과정을 안정화하기 위해 마커 유전자 및 계통도와 같은 도메인 지식을 명시적으로 통합하는 중간 단계를 옹호합니다.
• 한계: 저자들은 이 접근법이 선별된 마커 유전자 프로그램에 의존하며, 이는 새로운 세포 유형의 경우 불완전하거나 노이즈가 있을 수 있음을 겸손하게 지적합니다. 또한, 성능 향상이 반드시 파라미터 효율성 향상을 의미하는 것은 아닙니다.

요약하자면, CellRefine 은 대규모 사전 학습과 단일 세포 생물학 데이터의 구체적이고 종종 불균형한 현실 사이의 격차를 해소하는 프레임워크를 제공하여, 더 차별화되고 일반화 가능한 세포 표현을 산출합니다.