논문 정보

  • Title: scDiffusion: conditional generation of high-quality single-cell data using diffusion model
  • Authors: Erpai Luo, Minsheng Hao, Lei Wei, Xuegong Zhang
  • Affiliation: MOE Key Lab of Bioinformatics, Tsinghua University, Beijing
  • Journal: Bioinformatics, Volume 40, Issue 9, btae518 (2024)
  • DOI: 10.1093/bioinformatics/btae518

연구 배경 및 동기

단일세포 RNA 시퀀싱(scRNA-seq)은 세포 수준에서 생명 현상을 연구하는 데 핵심적인 기술이다. 그러나 충분한 양의 고품질 scRNA-seq 데이터를 확보하는 것은 여전히 어려운 과제다. 실험 비용이 높고, 특히 희귀 세포 유형(rare cell type)의 경우 충분한 수의 세포를 포착하기 어렵다.

이 문제를 해결하기 위해 생성 모델(generative model)을 활용하여 합성 scRNA-seq 데이터를 생성하려는 시도가 있어왔다. 기존 방법으로는 scGAN(GAN 기반), scDesign3(통계 학습 기반), SPARSim, SCRIP 등이 있으나, 생성된 데이터의 현실성이 부족하고, 특히 조건부 생성(conditional generation) — 특정 세포 유형이나 조직에 맞춰 데이터를 생성하는 것 — 에서 한계를 보였다.

한편, 이미지 생성 분야에서는 확산 모델(diffusion model)이 높은 충실도(fidelity)의 데이터 생성에서 압도적인 성능을 보여주고 있었다. 저자들은 이 확산 모델의 장점을 scRNA-seq 데이터 생성에 접목하여, scDiffusion이라는 새로운 생성 모델을 제안한다.

모델 아키텍처

scDiffusion 모델 아키텍처 — 오토인코더(SCimilarity), 디노이징 네트워크, 조건부 분류기로 구성된 잠재 확산 모델 Figure 1: scDiffusion의 전체 아키텍처 개요. 학습 단계에서 SCimilarity 인코더가 유전자 발현을 잠재 공간으로 변환하고, 확산 과정을 통해 노이즈를 추가한 뒤 디노이징 네트워크가 역방향 과정을 학습한다. 추론 단계에서는 가우시안 노이즈로부터 새로운 세포 임베딩을 생성한다. (출처: Luo et al., Bioinformatics 2024)

scDiffusion은 잠재 확산 모델(Latent Diffusion Model, LDM)과 파운데이션 모델(foundation model)을 결합한 구조로, 세 가지 핵심 모듈로 구성된다:

1. 오토인코더 (Autoencoder) — SCimilarity

  • 사전 학습된 파운데이션 모델 SCimilarity를 오토인코더로 활용
  • SCimilarity는 399개의 scRNA-seq 연구에서 수집한 2,270만 개 세포로 사전 학습된 인코더-디코더 네트워크
  • 유전자 발현 프로파일을 128차원 잠재 공간(latent space)으로 압축
  • 원시 분포를 가우시안 유사 분포로 변환하여 확산 과정에 적합한 형태로 정규화
  • 사전 학습 가중치를 기반으로 미세 조정(fine-tuning)하여, 처음부터 학습하는 것 대비 더 빠르고 우수한 성능 달성

2. 디노이징 네트워크 (Denoising Network)

  • 유전자 발현 데이터의 특성(긴 길이, 희소성, 비정렬)을 고려한 새로운 아키텍처 설계
  • CNN이나 Transformer 대신 스킵 연결 다층 퍼셉트론(skip-connected MLP) 기반 구조 채택
  • 스킵 연결 구조는 다양한 레벨의 특징을 보존하고 정보 손실을 줄임
  • 확산 과정의 역방향 과정(reverse process)을 학습하여 노이즈를 제거하고 의미있는 임베딩을 복원

3. 조건부 제어기 (Condition Controller) — 분류기 가이던스

  • 세포 유형 분류기(classifier)를 통한 분류기 가이던스(classifier guidance) 방식 적용
  • 분류기는 4층 MLP 구조로, 세포 유형이나 조직 유형 등의 조건 레이블을 예측
  • 디노이징 네트워크의 학습에 간섭하지 않고 별도로 학습되어 그래디언트를 제공
  • 다중 분류기를 동시에 사용하여 여러 조건 조합으로 데이터 생성 가능

학습 및 추론 파이프라인

학습 단계 (Training)

단계 내용 세부 사항
Step 1 오토인코더 미세 조정 SCimilarity 사전 학습 가중치 기반, 첫/마지막 레이어만 재학습
Step 2 디노이징 네트워크 학습 600,000~800,000 스텝, 학습률 어닐링 적용
Step 3 분류기 학습 ~200,000 이터레이션, 데이터셋 카테고리 수에 맞춰 설정

학습 과정에서 SCimilarity 모델이 유전자 발현 프로파일을 잠재 공간으로 임베딩한 후, 확산 과정을 통해 각 임베딩에 노이즈를 추가하여 노이즈 임베딩 시퀀스를 생성한다. 이 노이즈 임베딩이 디노이징 네트워크의 학습 데이터가 되며, 동시에 분류기는 임베딩으로부터 조건 레이블을 예측하도록 학습된다.

추론 단계 (Inference)

  1. 가우시안 노이즈를 입력으로 받음
  2. 디노이징 네트워크가 T=1,000 스텝에 걸쳐 반복적으로 노이즈를 제거
  3. 분류기 가이던스 또는 Gradient Interpolation 전략으로 생성 방향 제어
  4. 생성된 잠재 임베딩을 디코더에 입력하여 최종 유전자 발현 프로파일 복원

Gradient Interpolation 전략

scDiffusion의 가장 독창적인 기여 중 하나는 Gradient Interpolation 전략이다. 이는 이산적인 세포 상태(discrete cell states)로부터 연속적인 세포 발달 궤적(continuous developmental trajectory)을 생성할 수 있는 새로운 제어 전략이다.

기존의 직접 보간(linear interpolation)과 달리, Gradient Interpolation은 확산 모델이 학습한 세포 분포를 활용하여 생물학적으로 더 타당한 중간 상태를 생성한다. 구체적으로:

  • 두 개의 조건(예: 발달 시점 day 3과 day 4.5)을 분류기에 입력
  • 두 조건 사이의 그래디언트를 보간하여 중간 상태를 연속적으로 생성
  • 시퀀싱 간격 사이의 빈 구간을 채워 더 포괄적인 발달 타임라인 제공

실험 데이터셋

데이터셋 세포 수 유전자 수 세포 유형 용도
Tabula Muris 54,865 23,433 55 비조건부/단일조건부 생성
PBMC68k 68,579 ~2,000+ 11 비조건부/단일조건부 생성
Human Lung PF - - - 비조건부 생성
Tabula Sapiens - - - 다중조건부 생성
Mouse embryo - - - Gradient Interpolation

주요 결과

1. 비조건부 생성 (Unconditional Generation)

scDiffusion이 조건 없이 생성한 세포 데이터의 품질을 scGAN, scDesign3, SPARSim, SCRIP과 비교했다.

메트릭 scDiffusion scGAN scDesign3 SPARSim SCRIP
SCC (↑) 0.984 유사 유사 낮음 낮음
MMD (↓) 0.018 유사 유사 높음 높음
LISI (↑) 0.887 유사 유사 낮음 낮음
RF AUC (→0.5) 0.697 유사 유사 낮음 낮음
  • SCC (Spearman Correlation Coefficient): 실제와 생성 데이터 간 유전자 발현 상관관계. 0.984로 매우 높은 유사도
  • MMD (Maximum Mean Discrepancy): 분포 간 거리. 0.018로 실제 데이터와 매우 근접
  • LISI (Local Inverse Simpson’s Index): 실제와 생성 세포의 혼합 정도. 0.887로 높은 혼합도
  • RF AUC: 랜덤 포레스트가 실제/생성 세포를 구분하는 정확도. 0.5에 가까울수록 구분 불가능

UMAP 시각화에서도 scDiffusion이 생성한 세포들이 실제 세포들과 잘 겹치는 것을 확인할 수 있었다.

2. 단일 조건부 생성 (Single-Conditional Generation)

세포 유형 분류기를 가이드로 사용하여 특정 세포 유형의 데이터를 조건부로 생성했다.

  • Tabula Muris: 세포 유형별로 동일 수의 세포를 조건부 생성. UMAP에서 실제 세포와 시각적으로 중첩
  • 희귀 세포 유형도 성공적으로 생성:
    • Thymus 세포 (전체의 2.5%): 정확하게 생성
    • CD34+ 세포 (전체의 0.4%): 정확하게 생성
  • KNN 평가: 모든 세포 유형에서 AUC가 0.5 근처로, KNN 모델이 실제와 생성 세포를 구분하지 못함
  • CellTypist 분류: 생성 세포의 평균 분류 정확도 0.93 (실제 세포: 0.98)
    • 비교: scDesign3 생성 세포는 평균 0.99, scGAN 생성 세포는 평균 0.04 (구분 불가)

3. 다중 조건부 생성 (Multi-Conditional Generation)

세포 유형 + 조직 유형 두 가지 분류기를 동시에 사용하여, 학습 데이터에 없었던 조합의 세포를 생성하는 실험을 수행했다.

예를 들어, 학습 데이터에 유방(mammary gland)의 B세포는 있었지만, 흉선(thymus)의 기억 B세포(memory B cell)나 비장(spleen)의 대식세포(macrophage)는 없었다. scDiffusion은 이러한 분포 외(out-of-distribution) 조합도 성공적으로 생성했다:

생성 조건 CellTypist 분류 정확도 (생성) CellTypist 분류 정확도 (실제)
유방 B세포 98% 92%
흉선 기억 B세포 96.75% 96.91%
비장 대식세포 96.63% 99.53%

마커 유전자(marker gene) 발현 수준에서도 생성된 세포가 실제 세포와 유사한 패턴을 보였다.

4. Gradient Interpolation을 통한 발달 궤적 재구성

마우스 배아 발달 데이터에서 day 0~8 시점의 세포(day 3.5와 day 4 제외)로 scDiffusion을 학습한 뒤, day 3과 day 4.5 사이의 중간 발달 상태를 Gradient Interpolation으로 생성했다.

메트릭 scDiffusion (Gradient Interpolation) 직접 보간 (Linear Interpolation)
Mean MMD (↓) 0.3217 0.5206
Mean LISI (↑) 0.4488 0.3217
State 6 MMD 0.047 0.1289
State 8 MMD 0.0545 0.1167

scDiffusion은 치료(treatment) 정보 없이 학습되었음에도, 직접 보간보다 유의미하게 우수한 성능을 보였다. 이는 확산 모델이 세포 분포의 다양한 양상을 잘 포착하고, 중간 상태를 적절히 보간할 수 있음을 시사한다.

의의 및 한계

의의

  1. 파운데이션 모델 활용: SCimilarity라는 대규모 사전 학습 모델을 오토인코더로 활용하여, 유전자 발현의 보편적 표현(universal representation)을 확보
  2. 다중 조건부 생성: 여러 분류기를 동시에 활용하여 학습 데이터에 없던 조건 조합의 세포도 생성 가능
  3. Gradient Interpolation: 이산적 발달 시점 사이의 연속적 궤적을 생성하는 독창적 전략
  4. 희귀 세포 유형 생성: 전체의 0.4%에 불과한 세포 유형도 정확하게 생성 가능
  5. 확장성: 이론적으로 멀티오믹스(multi-omics) 데이터 등 다른 유형의 단일세포 데이터에도 적용 가능 (후속 연구 scDiffusion-X로 이어짐)

한계 및 후속 연구 방향

  • 학습 비용: 오토인코더, 디노이징 네트워크, 분류기를 순차적으로 학습해야 하며, 디노이징 네트워크만 60~80만 스텝이 필요
  • 분류기 의존성: 조건부 생성을 위해 별도의 분류기를 학습해야 하는 점은 한계. 후속 연구인 cfDiffusion은 분류기 없는(classifier-free) 가이던스로 이 문제를 해결
  • 조건 레이블 필요: 조건부 생성에는 세포 유형 등의 레이블 정보가 필수적이므로, 레이블이 없는 데이터에 대한 적용이 제한적
  • 평가 지표의 한계: 생성된 단일세포 데이터의 품질을 종합적으로 평가하기 위한 표준화된 벤치마크가 아직 부족

개인적 소감

scDiffusion은 이미지 생성에서 혁명적인 성과를 보인 확산 모델을 단일세포 유전체학에 접목한 흥미로운 시도다. 특히 인상적인 부분은:

  1. 파운데이션 모델과 확산 모델의 결합: 2,270만 개 세포로 학습된 SCimilarity의 표현력과 확산 모델의 생성력을 결합한 전략이 효과적이다. 유전자 발현 데이터의 고차원성과 희소성이라는 고유한 문제를 잠재 공간으로의 매핑을 통해 우아하게 해결했다.

  2. OOD 생성 능력: 학습 데이터에 없는 조건 조합(예: 특정 조직의 특정 세포 유형)을 생성할 수 있다는 점은 실용적 가치가 매우 크다. 실험적으로 얻기 어려운 조건의 세포 데이터를 in silico로 생성할 수 있기 때문이다.

  3. Gradient Interpolation: 발달 생물학 관점에서, 시퀀싱 시점 사이의 공백을 채워 연속적인 발달 궤적을 재구성할 수 있다는 것은 강력한 도구다.

노화 연구(aging research) 관점에서, scDiffusion과 같은 모델은 노화 과정에서의 세포 유형 변화를 모델링하는 데 활용될 수 있을 것이다. 특히 노화에 따라 감소하는 특정 면역세포나 줄기세포 데이터를 증강(augmentation)하거나, Gradient Interpolation을 통해 노화의 연속적인 전사체 변화 궤적을 재구성하는 데 적용할 수 있는 가능성이 있다.

References