LoCoFlow: Local-Context Flow Matching for Single-Cell Dynamics and Fate Inference

기술 요약: LoCoFlow

문제 정의

단일 세포 시퀀싱은 이산적인 시점에서의 세포 상태에 대한 고해상도 스냅샷을 제공하지만, 측정의 파괴적 특성으로 인해 개별 세포 간의 진정한 종단적 대응 관계를 관찰하는 것은 불가능하다. 스냅샷으로부터 연속적인 발달 궤적을 재구성하는 것은 여러 가능한 해가 존재하는 부적정 문제(ill-posed problem)이다. 또한, 기존 방법들은 세포를 고립된 벡터로 취급하는 경우가 많아, 세포 상태 공간의 국소적 매니폴드 기하학(local manifold geometry)이나 계보 분화의 내재적 불확실성을 충분히 표현하지 못한다. 유사한 전사 상태를 가진 세포들이 서로 다른 운명 결정 경향을 보일 수 있음에도 불구하고, 기존의 최적 운송(optimal transport) 및 흐름 매칭(flow-matching) 접근 방식은 일반적으로 개별 세포 수준에서 시간적 결합을 추론하며, 응집된 단위로서 국소 세포 이웃의 진화를 모델링하는 견고한 프레임워크가 부족하다.

방법론: LoCoFlow

저자들은 로컬 포인트 클라우드 기하학(local point-cloud geometry)과 조건부 운명 불확실성(conditional fate uncertainty)을 통합하여 단일 세포 역학을 모델링하는 프레임워크인 LoCoFlow(Local-Context Flow Matching)를 소개한다. 이 방법은 두 가지 주요 단계로 작동한다.

1. 국소 패치 구축 및 계층적 대응 (Local Patch Construction and Hierarchical Correspondence)

LoCoFlow는 고립된 세포를 매칭하는 대신, 앵커 세포 주변에 중첩되는 "패치"(국소 이웃)를 구축한다.

• 패치 정의: 발현 데이터의 경우 패치는 최근접 이웃 그래프 상의 측지선 이웃(geodesic neighborhoods)이며, 공간 데이터의 경우 측정된 좌표에서의 고정 크기 이웃이다.
• 계층적 운송: 이 방법은 2단계 최적 운송(OT) 체계를 사용한다. 먼저, 패치 중심과 공분산(위치와 분산 모두를 포착)을 정렬하기 위해 Bures 거리를 사용하여 패치 수준의 운송 문제를 해결한다. 그다음, 유지된 패치 쌍 내에서 엔트로피 정규화된 세포 수준의 운송 계획을 계산한다.
• 소프트 할당 (Soft Assignment): 이러한 계획들의 집계는 하드 할당 대신 각 소스에 대한 타겟 세포의 조건부 분포를 생성하여, 여러 가능한 후속 대응 관계를 보존한다.

2. 2단계 훈련 파이프라인

• 단계 1: 기하학적 정규화 브릿지 운송 (Geometry-Regularized Bridge Transport):
  • 계층적 운송을 통해 도출된 쌍을 이룬 엔드포인트 세트를 사용하여, 모델은 소스 패치와 타겟 패치 사이의 "브릿지"(중간 경로)를 학습한다.
  • 신경망은 잔차 속도장(residual velocity field) $f_\theta$를 예측하며, 이는 선형 보간법과 결합되어 경로가 정확히 엔드포인트를 통과하도록 보장한다.
  • 기하학적 사전 지식 (Geometric Priors): 손실 함수에는 운동 에너지 항과 패치의 상대적 변형(팽창/수축)을 제한하되 공통 이동(common translation)은 제한하지 않는 방사형 페널티(radial penalty)가 포함된다. 또한, 추정된 국소 매니폴드 기하학에 직교하는 속도에 대해서도 페널티를 부여한다.
• 단계 2: 조건부 속도 증류 (Conditional Velocity Distillation):
  • 추론 과정 중에는 미래 상태를 알 수 없으므로, 브릿지(두 엔드포인트를 모두 필요로 함)를 직접 사용할 수 없다. LoCoFlow는 브릿지 속도를 조건부 속도장 $V_\omega$로 증류한다.
  • 계층적 잠재 변수 (Hierarchical Latents): 모델은 계층적 잠재 구조를 사용하여 소스 컨텍스트 $Z_0$로부터 타겟 코드 $c_i$를 생성하는 법을 학습한다. 여기에는 인구 집단 수준의 결과를 포착하는 패치 공유 잠재 변수 $\eta_P$와 개별 변이를 포착하는 세포별 잠재 변수 $\eta_i$가 포함된다.
  • 훈련: 타겟 인코더는 먼저 관측된 엔드포인트를 코드 $c^\star$로 매핑하는 법을 학습하고, 생성기는 소스 컨텍스트와 잠재 변수로부터 이러한 코드를 생성하는 법을 학습한다. 속도 네트워크는 생성된 코드에 조건화된 브릿지 속도와 일치하도록 훈련된다.
  • 추론: 추론 시에는 잠재 코드가 사전 분포(미래 관측값 없이)로부터 추출되어 고정된 상태로 적분된다. 이를 통해 운명 불확실성을 나타내는 다양한 궤적을 샘플링할 수 있다.

주요 기여

1. 국소 컨텍스트 표현 (Local-Context Representation): 역동적 모델링의 근본 단위를 고립된 발현 벡터에서 국소 인구 상태(포인트 클라우드)로 전환하여, 이웃 기하학의 진화를 명시적으로 모델링한다.
2. 계층적 운송 (Hierarchical Transport): 세포 집단의 전역적 움직임과 내부 구조적 변화(예: 이질성, 확장)를 모두 존중하는 결합된 패치 수준 및 세포 수준의 최적 운송 메커니즘을 도입한다.
3. 조건부 운명 불확실성 (Conditional Fate Uncertainty): 기하학적 경로 제약과 결과의 불확실성을 분리하는 2단계 증류 프로세스를 개발하여, 미래 관측값 없이도 여러 가능한 운명 궤적을 샘플링할 수 있게 한다.
4. 섭동 분석 프레임워크 (Perturbation Analysis Framework): 유전자 섭동 시뮬레이션에서 초점 세포 효과(focal-cell effects)와 이웃(배경) 효과를 구분하는 메커니즘을 제공하여, 국소적 맥락이 예측된 반응을 어떻게 형성하는지 밝혀낸다.

실험 결과

논문은 EB 분화, iPSC 리프로그래밍, C. elegans 발달, Axolotl 뇌 재생의 세 가지 데이터셋에 대해 LoCoFlow를 평가한다.

• 홀드아웃 재구성 (Held-Out Reconstruction): EB 및 iPSC 데이터셋에서 LoCoFlow는 대부분의 고차원 설정(예: PCA50)에서 OT-CFM, MFM, PACE와 같은 베이스라인보다 낮은 Wasserstein 거리($W_1, W_2$)를 달로한다. 그러나 논문은 LoCoFlow가 때때로 더 높은 Maximum Mean Discrepancy(MMD)를 보인다는 점을 언급하며, 이는 예측이 세포 상태 공간의 특정 영역에 더 집중될 수 있음을 시사한다.
• 운명 예측 (LARRY 데이터): 계보 추적된 LARRY 데이터를 사용하여, LoCoFlow는 평균 raw Brier 에러를 0.7990(scDiffEq)에서 0.6620으로 줄였으며, 이는 전체 운명 확률 벡터를 예측하는 정확도가 향상되었음을 나타낸다.
• 섭동 반응 (Perturbation Responses):
  • 기하학적 패치 (Geometric Patches): 기하학적으로 정의된 패치에서 유전자를 섭동할 때, 배경 전용(background-only) 및 모든 구성원 개입(all-member interventions)은 정렬된 반응(예: 단핵구에서 호중구로의 운명 변화)을 생성한다.
  • 무작위 미니배치 (Random Minibatches): 이와 대조적으로, 무작위로 조립된 패치를 사용하는 모델은 초점 세포 반응과 배경 반응이 상충하는 양상을 보여, 국소적 맥락이 확률적 운명 추론을 위한 유용한 표현임을 입증한다.
• 공간 보간 (Axolotl): LoCoFlow는 Axolotl 뇌 재생에서 OT-CFM에 비해 개선된 공간 예측($W_2$)을 보여주지만, 저자들은 예측된 공간 분포가 소스 시점과 시각적으로 매우 유사함을 언급하며 추가적인 공간 정렬의 필요성을 시사했다.

의의 및 주장

이 논문은 LoCoFlow가 국소 세포 집단 구조가 추론된 궤적을 어떻게 형성하는지를 연구하고 운명의 불확실성을 명시적으로 표현하는 견고한 프레임워크를 제공한다고 주장한다. 저자들은 다음을 강조한다:

• 국소 컨맥스트의 중요성: 기하학적 패치와 무작위 미니배치 사이에서 관찰된 뚜렷한 반응 패턴은 국소적 맥락이 확률적 운명 추론을 위한 유용한 표현이라는 가설을 뒷받침한다.
• 모델링 가정: 저자들은 스냅샷 데이터가 세포의 역사를 유일하게 식별할 수 없다는 점을 인정하며 신중한 입장을 유지한다. LoCoFlow의 기하학적 및 잠재적 구성 요소는 실제 계보 관계를 드러내기보다는 모델링 가정을 부과하는 것이다.
• 검증 가능한 가설: 이 프레임워크는 국소적 맥락이 섭동에 대한 예측된 반응을 어떻게 형성하는지에 대한 검증 가능한 가설을 생성하며, 세포 고유의 변화와 환경적 영향(background effects)을 구분한다.

이 연구는 궤적 추론의 근본적인 모호성을 해결한다고 주장하는 것이 아니라, 세포 상태의 매니폴드 기하학에 더 잘 부합하고 운명 불확실성을 정량화하고 샘플링할 수 있는 구조화된 방법을 제공하는 데 목적이 있다.