Intermediate Layers Encode Optimal Biological Representations in Single-Cell Foundation Models

Motivation

• Known: 현재 단일세포 파운데이션 모델 벤치마크는 보편적으로 최종층 임베딩을 추출하며, NLP와 vision 분야에서 중간층이 task-relevant information bottleneck을 통해 최종층을 능가할 수 있다는 연구가 존재한다.
• Gap: 단일세포 파운데이션 모델에서 층별 표현 품질의 체계적 분석이 부재했으며, 궤적 추론과 섭동 응답 예측 같은 생물학적 과제에서 최적 층이 과제와 맥락에 따라 변하는지 미검증 상태였다.
• Why: 층별 표현 분석을 통해 모델에서 feature를 추출하는 '위치'가 추출하는 '내용'만큼 중요함을 보여주며, 이는 벤치마크 및 다운스트림 응용에서 모델 성능을 향상시키는 실무적 지침을 제공한다.
• Approach: scFoundation(100M, 12층)과 Tahoe-X1(1.3B, 24층)에서 동결된 모델로 forward pass를 수행하여 각 transformer block 출력의 hidden state를 추출하고, 궤적 추론(diffusion pseudotime 상관성)과 섭동 응답 예측(representational similarity analysis)으로 층별 성능을 평가했다.

Achievement

• 궤적 추론의 과제-의존적 최적층: Tahoe-X1은 60% 깊이에서 최적(ρ=0.76)이며 최종층 대비 31% 우수, scFoundation은 단조증가 패턴을 보여 모델 크기와 아키텍처에 따른 층별 특성 차이 발견
• 섭동 응답의 맥락-의존적 변동성: 동일 모델에서 최적층이 T cell 활성화 상태에 따라 0%~96% 범위에서 변동하며, 정지기 세포에서는 첫 층이 더 우수한 성능 달성
• 최종층 우월성 가정의 도전: 중간층이 전이학습 생물학적 표현을 더 잘 보존하며 최종층은 사전학습 목표에 과도하게 특화되어 있음을 입증

How

• frozen 모델에서 각 transformer block 출력 후 hidden state 추출 (residual addition과 layer normalization 후)
• 교차-모델 비교를 위해 층 인덱스를 총 깊이로 정규화
• 궤적 추론: k-nearest neighbor graph (k=15, cosine similarity) 구성 후 layer-specific diffusion pseudotime(DPT) 계산 및 Spearman 상관으로 참조 pseudotime과 비교
• 섭동 응답: 각 섭동에 대해 DESeq2 기반 differential expression profile 및 centroid embedding 계산, biological reference(DE profile의 Spearman 상관)와 embedding-space 유사도 행렬(cosine 유사도)의 정렬도를 Spearman 상관으로 측정
• 평가 데이터: 조혈 분화의 LARRY 기반 cord blood LT-scSeq 데이터셋(궤적 추론) 및 CRISPRi perturb-seq CD4+ T cell 데이터셋(22M cells, 섭동 응답)

Originality

• 단일세포 파운데이션 모델에서 최초의 체계적 층별 표현 분석으로 NLP/vision의 중간층 우월성 발견을 생물학적 맥락에 적용
• 궤적 추론과 섭동 응답 예측이라는 상이한 생물학적 과제에서 층별 최적성의 일관성과 변동성을 비교 분석하는 이중 과제 설계
• T cell 활성화 상태에 따른 96% 포인트의 극단적 층 변동을 실증하여 '맥락-의존적 표현'의 개념 도입", '정지기 세포에서 첫 층의 우월성 발견으로 세포 상태와 층별 특성화의 관계성 규명

Limitation & Further Study

• 평가 과제가 궤적 추론과 섭동 응답 예측 2개에 제한되어 다른 다운스트림 과제(세포 타입 주석 등)에서의 일반화 가능성 미확인
• 데이터셋 규모 제약: CRISPRi 섭동 분석에서 D1 donor 단일 샘플과 조건당 상위 100개 섭동만 사용하여 생물학적 대표성 제한
• 기저 메커니즘 미분석: 중간층에서 최적 성능이 나타나는 이유(정보 병목, 기울기 흐름, 특성 진화 등)에 대한 해석 부족
• 후속 연구: (1) 다양한 세포 타입과 조직에서의 층별 특성 분석, (2) 층 선택의 자동화 방법 개발, (3) 최종층 특화의 기저 메커니즘 탐구, (4) 층별 임베딩을 통합하는 앙상블 전략 검토

Evaluation

총평: 단일세포 파운데이션 모델 평가에서 최종층 임베딩의 보편적 가정을 체계적으로 도전하며, 과제와 맥락별 층별 최적성의 변동을 실증적으로 입증한 중요한 연구이다. 벤치마크 구축과 모델 응용 시 층 선택의 중요성을 부각하는 실무적 기여가 있으나, 메커니즘 분석과 평가 과제의 확대가 필요하다.