Discrete Diffusion for Single-Cell Gene Expression Modeling

이 논문은 이산적 유전자 발현 데이터를 직접 모델링하는 '이산 세포 모델 (DCM)'을 제안하여, 기존 연속적 잠재 공간 기반 방법론보다 단일 세포 전사체 모델링 성능을 획기적으로 개선하고 유전적 교란에 대한 조건부 생성 능력을 입증했습니다.

핵심

🧬 핵심 비유: "레고 블록 vs 점토"

이 논문의 핵심은 **"세포 데이터를 어떻게 표현할 것인가?"**에 대한 질문에서 시작합니다.

1. 기존 방법 (점토 모델):

   • 지금까지의 AI 모델들은 세포 속 유전자의 수 (데이터) 를 연속적인 점토로 취급했습니다.
   • 유전자 A 의 발현량이 '5'개라고 해도, AI 는 이를 '4.9'나 '5.1'처럼 소수점 단위로 표현하고 계산했습니다.
   • 문제점: 실제로는 유전자가 '5 개'일지 '6 개'일지 정수 (Integer) 로만 존재합니다. 점토처럼 흐르는 값으로 만들면, 실제로는 존재할 수 없는 '5.3 개의 유전자' 같은 엉뚱한 상태를 만들어내게 됩니다. 마치 레고 블록을 쌓는 데 점토를 쓰면서 "이제 0.5 개의 블록을 쌓자"라고 하는 것과 비슷합니다.

2. 새로운 방법 (DCM, 레고 모델):

   • 이 논문이 제안한 **DCM(Discrete Cell Models)**은 데이터를 레고 블록처럼 정수 (0, 1, 2, 3...) 그대로 다룹니다.
   • 유전자가 '없음 (0)'인지 '하나 (1)'인지, '두 개 (2)'인지를 명확하게 구분합니다.
   • 장점: 실제 생물학적 현상 (유전자가 켜지거나 꺼지는 것) 을 훨씬 더 정확하게, 자연스럽게 모방할 수 있습니다.

---

🚀 이 모델이 하는 일: "세포의 미래를 예측하는 시간 여행"

이 모델은 확산 (Diffusion) 기술을 사용합니다. 이를 쉽게 비유하자면 다음과 같습니다.

• 전진 과정 (소음 추가): 깨끗한 세포 데이터 (유전자 발현량) 에 점점 더 많은 '소음'을 섞어서, 결국 아무것도 없는 상태 (MASK) 로 만들어버립니다. (예: 레고 성을 부수고 조각난 블록을 섞는 과정)
• 역행 과정 (복원): AI 는 이 섞여버린 조각난 블록들을 보고, "아, 원래는 이렇게 조립되어 있었구나!"라고 추론하며 원래의 세포 상태로 되돌리는 과정을 학습합니다.

기존 모델은 이 복원 과정을 '점토'로 했다면, DCM 은 '레고 블록'을 하나씩 정확히 맞춰서 원래 모양으로 되돌립니다.

---

🏆 왜 이것이 중요한가요? (성과)

이 모델은 두 가지 주요 테스트에서 기존 최고의 모델들 (scVI, scLDM 등) 을 압도했습니다.

1. 세포의 전체적인 모양 (분포) 을 더 잘 잡습니다:

   • 비유: 새로운 세포를 만들 때, 기존 모델은 "대충 비슷하게 생겼네"라고 했지만, DCM 은 "정확히 같은 얼굴, 같은 옷차림"을 만들어냅니다.
   • 결과: Dentate Gyrus(치상회) 데이터셋에서 기존 모델보다 약 5 배 더 정확한 결과를 보여주었습니다.

2. 유전자 조작 (변형) 을 예측하는 데 탁월합니다:

   • 비유: "만약 이 유전자를 끄면 (Knockout), 세포는 어떻게 변할까?"라는 질문에 답할 때, DCM 은 실제 실험 결과와 가장 유사한 변화를 예측했습니다.
   • 결과: 유전자를 조작했을 때 세포가 어떻게 반응하는지 예측하는 데서도 가장 낮은 오차율을 기록했습니다.

---

💡 왜 '정수'로 다루는 게 더 나을까요?

논문의 저자들은 다음과 같은 이유를 듭니다.

• 0 과 1 의 차이는 큽니다: 유전자가 '0 개'인 것 (아예 꺼짐) 과 '1 개'인 것 (약간 켜짐) 의 차이는 생물학적으로 매우 큽니다. 하지만 '100 개'와 '101 개'의 차이는 그냥 잡음일 뿐입니다.
• 기존 모델의 한계: 점토 (연속) 모델은 이 '0 과 1 의 큰 차이'와 '100 과 101 의 작은 차이'를 구분하기 위해 AI 가 힘들게 학습해야 합니다.
• DCM 의 장점: 레고 (이산) 모델은 처음부터 이 차이를 구조적으로 이해하고 있습니다. 0 은 0, 1 은 1 로 명확하니까요.

---

📝 한 줄 요약

"기존 AI 는 세포 데이터를 흐르는 물 (점토) 로 다뤄서 부정확한 결과를 냈다면, 이 새로운 모델 (DCM) 은 세포 데이터를 정확한 레고 블록으로 다뤄서, 실제 세포와 거의 구별할 수 없을 정도로 정교한 가상의 세포를 만들어냅니다."

이 연구는 생물학 데이터를 모델링할 때, 데이터의 본질 (정수/이산) 에 맞는 방식을 선택하는 것이 얼마나 중요한지 보여주며, 향후 '가상의 세포 (Virtual Cell)'를 만드는 기초 기술로 큰 기대를 받고 있습니다.

기술 요약

1. 문제 정의 (Problem)

단일 세포 전사체학 (Single-cell transcriptomics) 데이터는 본질적으로 **이산적 (discrete)**이고 **희소 (sparse)**한 mRNA 카운트 (정수 값) 로 구성됩니다. 그러나 기존의 생성 모델 (scVI, scGPT, scLDM 등) 은 이러한 정수 카운트 데이터를 **연속적인 잠재 공간 (continuous latent space)**으로 매핑하여 모델링하는 방식을 주로 사용했습니다.

이러한 연속적 완화 (continuous relaxation) 접근법의 한계는 다음과 같습니다:

• 불가능한 상태 할당: 연속 공간은 정수가 아닌 값에 확률 질량을 할당하여, 실제 측정 불가능한 상태에 모델 용량을 낭비합니다.
• 비대칭적 거리 측정: 카운트 공간에서 0 과 1 의 차이 (발현 유무) 는 100 과 101 의 차이 (샘플링 노이즈) 와 생물학적으로 완전히 다릅니다. 연속적인 유클리드 거리는 이러한 비대칭성을 데이터로부터 학습해야 하므로 본질적인 표현의 한계가 있습니다.
• 정보 이론적 간극: 데이터 생성 과정이 이산적일 때, 연속적 완화는 이산 공간 내의 구조를 학습하는 대신 이산화 경계를 학습하도록 강요합니다.

따라서, 원시 카운트 데이터를 이산 도메인에서 직접 모델링하여 이러한 표현적 한계를 극복할 필요가 있습니다.

2. 방법론 (Methodology)

저자들은 **Discrete Cell Models (DCM)**을 제안하며, 이는 Score Entropy Discrete Diffusion (SEDD) 을 단일 세포 전사체 데이터에 직접 적용한 프레임워크입니다.

• 데이터 표현: 각 세포의 발현 프로필을 $M$개의 유전자에 대한 이산 시퀀스로 표현합니다. 각 유전자의 발현 수준은 $X = \{0, 1, \dots, K\}$ 범위의 이산 토큰으로 처리됩니다.
• 전향 확산 과정 (Forward Diffusion):
  • 연속 시간 마르코프 과정을 사용하여 깨끗한 발현 프로필을 점진적으로 오염시킵니다.
  • 모든 토큰이 특수한 'MASK' 상태로 전이되는 흡수 (absorbing) 확산 구조를 사용합니다.
• 역방향 확산 과정 (Reverse Process):
  • Concrete Score를 학습합니다. 이는 연속 확산의 $\nabla_x \log p_t(x)$에 해당하는 이산적 비율로, 인접한 이산 상태 간의 확률 분포 비율을 추정합니다.
  • Score 네트워크 $s_\theta$는 해밍 거리 1 인 시퀀스 간의 비율을 예측하도록 설계됩니다.
• 학습 목표 (Training Objective):
  • 흡수 확산의 경우, DWDSE (Denoising Weighted Score Entropy) 목표 함수가 가중 크로스 엔트로피 (Weighted Cross-Entropy) 손실로 단순화됩니다.
  • 이는 $L = \mathbb{E}[\text{CrossEntropy}(p_\theta(x_0 | x_t, t, c), x_0)]$ 형태로, tractable 한 가능도 기반 학습을 가능하게 합니다.
• 조건부 생성 (Conditional Generation):
  • 세포 유형 (Cell type) 과 유전적 교란 (Perturbation, 예: 유전자 녹아웃) 정보를 조건으로 포함합니다.
  • 교란 레이블은 단백질 언어 모델을 통해 임베딩되며, AdaLN (Adaptive LayerNorm) 메커니즘을 통해 Score 네트워크에 주입됩니다.
• 아키텍처:
  • DiT (Diffusion Transformer) 백본을 사용하며, Flash Attention 을 통해 약 17k 개의 유전자 (전체 발현 프로필) 를 효율적으로 처리합니다.
  • 발현되지 않거나 선택되지 않은 유전자는 'PAD' 토큰으로 처리하고 마스크 어텐션을 적용합니다.

3. 주요 기여 (Key Contributions)

1. 이산 확산 프레임워크의 도입: 단일 세포 데이터 생성을 위해 연속 잠재 공간 대신 원시 카운트 데이터를 직접 이산 도메인에서 모델링하는 최초의 프레임워크 (DCM) 를 제안했습니다.
2. 조건부 및 무조건부 생성 지원: 하나의 엔드 - 투 - 엔드 아키텍처로 세포 유형, 교란 조건 등을 결합하여 복잡한 생물학적 시나리오 (예: 특정 세포 유형의 유전적 교란에 대한 전사적 반응) 를 정밀하게 모델링합니다.
3. 효율성 및 확장성: 별도의 인코더 - 디코더 파이프라인 (VAE) 이 필요 없는 단일 구조로, 기존 SOTA 모델보다 파라미터 수가 적으면서도 우수한 성능을 달성합니다.

4. 실험 결과 (Results)

저자들은 Dentate Gyrus (무조건부 생성) 와 Replogle (조건부 교란 예측) 두 가지 벤치마크에서 DCM 을 평가했습니다.

• 무조건부 생성 (Dentate Gyrus 데이터셋):

  • W2 거리: 기존 최상위 연속 확산 모델인 scLDM 대비 약 2 배 개선 (scLDM: 10.615 $\to$ DCM: 5.913).
  • MMD2RBF: scLDM 대비 5 배 개선 (scLDM: 0.102 $\to$ DCM: 0.019).
  • 이는 DCM 이 전사체 분포의 전역적 기하학적 정렬과 미세한 통계적 유사성 모두에서 탁월함을 보여줍니다.

• 조건부 생성 (Replogle 교란 데이터셋):

  • W2 거리: 모든 베이스라인 (scVI, CPA, scGPT, STATE, scLDM) 을 능가하는 최고 성능을 기록했습니다 (scLDM 대비 13% 개선).
  • MMD2RBF: scLDM 에 비해 다소 낮았으나, Parse 1M 벤치마크에서는 경쟁력 있는 성능을 보였습니다.
  • 해석: W2 점수의 우수성은 평균 발현 프로필과 유전자 수준의 분산을 정확히 복원함을 의미하며, MMD2RBF 의 차이는 고차원 의존성 구조 (유전자 간 상관관계) 모델링의 복잡성에서 기인한 것으로 분석됩니다.

5. 의의 및 결론 (Significance & Conclusion)

• 생물학적 모델링의 패러다임 전환: 이 연구는 생성 모델의 상태 공간이 모델링하려는 생물학적 측정치 (이산적이고 희소한 카운트) 의 구조와 일치할 때 표현력이 향상됨을 입증했습니다.
• 가상 세포 (Virtual Cell) 의 발전: 연속적 근사에 의존하지 않고 직접 이산 데이터를 학습함으로써, 유전적 교란에 대한 더 정확한 전사체 예측이 가능해졌습니다.
• 미래 방향: 현재 조건부 생성 시 고차원 상관관계 (MMD2RBF) 모델링에 약간의 한계가 있으나, 이는 주의 기반 (attention-based) 조건부 메커니즘 개선 등을 통해 해결 가능한 과제로 남겼습니다.

결론적으로, DCM은 단일 세포 전사체 모델링 분야에서 이산 확산 (Discrete Diffusion) 이 강력한 기초 모델 (Foundational Model) 로서의 가능성을 제시하며, 생물학적 데이터의 본질적인 특성을 존중하는 새로운 방향성을 제시합니다.