PACE: Geometry-Aware Bridge Transport for Single-Cell Trajectory Inference

Motivation

• Known: 기존 optimal transport 및 flow matching 기반 궤적 추론 방법들은 관측된 clock time에서 Euclidean 근접성 또는 OT 거리로 세포 간 coupling을 추정하고, Euclidean kinetic energy나 smoothness로 경로를 정규화한다. Metric Flow Matching과 같은 geometry-aware 방법도 존재하나 support-aware metric을 endpoint pairing 고정 이후에 적용한다.
• Gap: 발달 과정이 비동기적이어서 같은 실험적 시점에 샘플링된 세포들이 서로 다른 latent pseudotime 단계에 있을 경우, Euclidean 근접성 기반 coupling은 궤적을 잘못 정렬시킬 수 있으며, snapshot 분포만으로는 cross-time correspondence와 intermediate trajectory가 근본적으로 non-unique(ill-posed)하다는 문제가 존재한다.
• Why: RNA velocity, lineage tracing, metabolic labeling 같은 보조 측정 없이도 비동기적 발달 과정에서 생물학적으로 타당한 세포 상태 전이를 복원할 수 있다면, 발달·분화·면역활성화·종양진화·리프로그래밍 등 다양한 생물학적 과정의 기전 이해와 실험 설계에 폭넓게 활용될 수 있다.
• Approach: PACE는 국소 spatiotemporal 이웃에서 추정한 tangent 방향을 기반으로 state- 및 time-dependent anisotropic Riemannian metric을 구성하고, 이 metric이 유도하는 path-action cost 하에서 OT coupling 정제와 endpoint-preserving neural bridge fitting을 번갈아 수행한 뒤, 학습된 bridge dynamics를 global continuous-time velocity field로 distill한다.

Achievement

Figure 3. Velocity-alignment diagnostics on Ocean (Petrovi´c et al.,

1. Riemannian metric 설계: local PCA 기반 spatiotemporal tangent-space projection을 이용해 시간·상태에 따라 변하는 anisotropic metric을 구성하여, 국소적으로 지지되는 tangent 방향의 이동 비용은 낮게, normal 방향 이동 비용은 높게 부여했다.
2. 반복적 bridge-transport 절차: Riemannian path-action cost 하에서 cross-time coupling 정제와 endpoint-preserving neural bridge fitting을 alternating 방식으로 결합하여, correspondence와 interpolant path를 동일한 geometry-aware criterion 하에 공동으로 선택했다.
3. 정량적 성능 향상: 7개의 controlled/biological 데이터셋, 9개의 held-out reconstruction 실험에서 가장 강력한 baseline 대비 MMD, W1, W2를 평균 23.7% 감소시켰으며, embryoid body 분화 벤치마크에서 RNA-velocity alignment를 15.4% 개선했다. 이는 explicit cell pairing, lineage tracing, RNA velocity supervision 없이 달성되었다.
4. Ablation을 통한 구성요소 검증: 각 component(metric, coupling refinement, bridge distillation)에 대한 ablation을 통해 각 요소의 기여를 확인했다.

How

Figure 1. Overview of PACE. PACE uses local PCA to construct an anisotropic metric Gk(x, t) = I + αC(k)

• Proposition 3.1을 통해 snapshot만으로는 cross-time coupling과 intermediate trajectory가 non-unique함을 증명하여 문제를 variational 형태로 재정식화함
• 관측된 세포들의 국소 spatiotemporal neighborhood에서 local PCA로 tangent 방향을 추정하여 anisotropic Riemannian metric을 구성
• 이 metric으로 유도된 minimum Riemannian action을 endpoint 간 coupling cost로 정의하여 anisotropic bridge transport 문제를 설정
• OT coupling 정제 단계와 endpoint-preserving neural bridge fitting 단계를 alternating 방식으로 반복 수행
• 최종적으로 학습된 bridge dynamics를 distill하여 continuous-time population velocity field를 얻고, 이를 적분해 xt0에서 시계열 전체 궤적을 생성

Originality

• 관측 불가능한 실제 cell identity와 intermediate path 대신, 국소 spatiotemporal neighborhood가 제공하는 약한 prior(국소 지지 tangent 방향)를 이용해 admissible한 이동 방향을 정의하는 새로운 관점 제시
• Euclidean endpoint proximity 대신 state- 및 time-dependent anisotropic Riemannian metric으로 path-action cost를 정의하여 coupling과 path를 동시에 규정
• Metric Flow Matching과 달리 endpoint pairing을 고정하지 않고, geometry-aware action 하에서 coupling refinement와 bridge fitting을 alternating하여 correspondence 자체도 geometry에 의해 재조정되도록 설계

Limitation & Further Study

• Local PCA 기반 tangent 추정은 국소 이웃의 밀도와 샘플 수에 민감할 수 있어, sparse하거나 noise가 큰 고차원 데이터에서는 metric 추정이 불안정해질 위험이 있음
• Alternating 최적화 절차의 수렴성 및 계산 비용에 대한 분석이 제한적으로 보이며, 대규모 세포 수나 많은 시점(anchor time)으로의 확장성 검증이 추가로 필요함
• RNA velocity supervision 없이 학습하지만, 검증 자체는 RNA velocity와의 alignment로 이루어져 있어 다른 독립적 ground truth(lineage tracing 등)와의 비교가 추가되면 결과의 신뢰성이 높아질 것으로 보임
• Workshop paper 형태로 제출되어 있어 이론적 증명(Appendix D)과 실험 세부사항에 대한 더 깊은 검증이 후속 연구로 필요함

Evaluation

총평: 비동기적 발달 과정에서 Euclidean 근접성의 한계를 명확히 짚어내고, geometry-aware한 metric으로 coupling과 path를 통합적으로 재정의한 독창적이고 실용적인 접근으로, 다수의 데이터셋에서 일관된 성능 향상을 보여준 완성도 높은 연구이다.