FreeBridge: Variational Schrödinger Bridges for Cellular Transition Dynamics

기술 요약: FreeBridge – 세포 전이 역학을 위한 변분 슈뢰딩거 브릿지(Variational Schrödinger Bridges)

문제 정의
고함량 이미징 어세이(High-content imaging assays)는 화학적 및 유전적 섭동에 대한 세포의 반응을 정량화하는 데 필수적입니다. 그러나 근본적인 한계는 이미지 획득 시점에 세포가 화학적으로 고정되어 있어, 개별적인 시간적 궤적을 관찰할 수 없다는 점입니다. 결과적으로 섭동 모델링은 제어(control) 집단과 처리(treated) 집단 사이의 확률적 수송(stochastic transport)을 추론하는 문제로 축소되며, 이때 알려진 것은 오직 말단 경계 분포(marginal endpoint distributions)뿐입니다.

최근의 생성 모델(예: 확산 모델, 플로우 매칭)은 제어 및 섭동 종단 간의 강력한 정렬을 달est하지만, 종단 일치성(endpoint consistency)만으로는 중간 진화 과정을 결정하기에 불충분합니다. 동일한 주변 분포를 연결하는 확률 과정은 무수히 많으며, 추가적인 제약 조건이 없다면 이러한 모델들은 관찰된 단일 세포 형태학(morphologies)에 의해 지지되지 않는 잠재 공간(latent space) 영역을 통과할 수 있습니다. 기존 방법들은 분포 수준의 목적 함수를 최적화하여 중간 세포 구조를 미결정 상태로 남겨두는 경우가 많으며, 이는 FID나 KID와 같은 지표를 만족할 수는 있으나 아티팩트(artifact)를 생성할 위험이 있습니다.

방법론: FreeBridge 프레임워크
저자들은 유효한 세포 상태의 정의와 전이 메커니즘을 명시적으로 분리하는 "셀 엔진(Cell Engine)"으로 정식화된 제약된 확률적 수송 프레임워크인 FreeBridge를 제안합니다. 이 방법은 종단 데이터만을 감독(endpoint-only supervision) 하에 작동하며, 세 가지 핵심 구성 요소로 이루어져 있습니다:

1. 상태 단위 (명시적 단일 세포 기하학):

   • FreeBridge는 이미지를 집단 수준의 혼합물로 취급하는 대신, 원자적 상태를 인스턴스 분할된 단일 세포 표현(Cellpose 사용)으로 정의합니다.
   • 고정된 사전 학습된 인코더 $\phi$는 이 크롭(crop)들을 잠재 공간 $z = \phi(x)$로 임베딩합니다. 이 인코더는 전이 학습 동안 고정되어, 전이 규칙을 학습하기 전에 정의된 고정된 잠재 기하학을 구축합니다.
   • 경험적 지원 뱅크(empirical support bank) $Z_{emp}$가 관찰된 세포의 잠재 임베딩으로부터 구축됩니다. 이 뱅크는 생물학적으로 유효한 형태학의 "매니폴드(manifold)"를 정의합니다.

2. 전이 규칙 (슈뢰딩거 브릿지 수송):

   • 제어 분포 $\mu_z$와 섭동 분포 $\nu_{c,z}$ 사이의 전이는 제어된 확률 미분 방정식(SDE)으로 모델링됩니다: $dZ_t = u_t(Z_t) dt + \sigma dW_t$.
   • 목표는 $t=1$에서 초기 분포를 대상 분포로 수송하는 드리프트 필드(drift field) $u_t$를 찾는 것입니다.
   • 표준 최적 수송(optimal transport)과 달리, 이 정식화는 경로를 규제하기 위해 **태스크 특화 상태 비용(task-specific state cost)**을 포함합니다.

3. 기하학적 규제 (경험적 지원 제약):

   • 확률 과정이 잠재 공간의 "공백(voids)"을 통과하는 것을 방지하기 위해, FreeBridge는 지원 비용 $V(z) = \lambda \cdot d_{emp}(z)^2$를 도입합니다. 여기서 $d_{emp}(z)$는 경험적 지원 뱅크 $Z_{emp}$에서의 최근접 이웃까지의 거리입니다.
   • 전체 목적 함수는 운동 에너지(드리프트 크기)와 지원 비용의 합을 최소화합니다:
     $$\min_u \mathbb{E} \left[ \int_0^1 \frac{1}{2\sigma^2}\|u_t(Z_t)\|^2 dt + \lambda \int_0^1 d_{emp}(Z_t)^2 dt \right]$$
   • 이 규제는 쌍을 이룬 세포 수준의 궤적이 필요하지 않음에도 불구하고, 중간 상태를 경험적으로 관찰된 형태학 매니폴드에 고정시킵니다.

주요 기여

• 정식화: 세포 상태의 명시적 지정과 확률적 수송을 명확히 분리하는 단일 세포 전이 모델링을 위한 슈뢰딩거 브릿지 정식화를 도입했습니다.
• 기하학적 근거 제시: 비모수 밀도 추정(경험적 지원 뱅크에 대한 최근접 이웃 거리)에 기반한 새로운 규제 항을 제안하여 중간 궤적을 생물학적으로 타당한 영역 내로 제한했습니다.
• 통합 프로토콜: 인스턴스 분할된 단일 세포 크롭을 사용하여 여러 데이터셋(BBBC021, RxRx1, JUMP)에 걸친 통합 평가 프로토콜을 구축함으로써 공정한 비교를 보장했습니다.
• 성능: 기하학적 제약이 종단 충실도(endpoint fidelity)뿐만 아니라 섭동 메커니즘의 의미론적 유지(Mechanism-of-Action)를 개선함을 입증했습니다.

결과
BBBC021, RxRx1, JUMP 데이터셋에 대해 평가한 결과, FreeBridge는 다음과 같은 성과를 보였습니다:

• 종단 충실도 (Endpoint Fidelity): 최첨단 베이스라인(PhenDiff, IMPA, CellFlux)과 비교하여 경쟁력 있거나 향상된 Fréchet Inception Distance (FID) 및 Kernel Inception Distance (KID) 점수를 달성했습니다. BBBC021에서 가장 낮은 FID(15.35)와 KID(1.10)를 기록했습니다.
• 중간 일관성 (Intermediate Consistency): 경험적 형태학 지원 범위를 벗어나는 중간 상태의 비율을 측정하는 지표인 **지원 위반율 ($R_{viol}$)**을 유의미하게 감소시켰습니다. BBBC021에서 FreeBridge는 $R_{viol}$을 0.11로 낮추었으며, 이는 PhenDiff(0.38) 및 CellFlux(0.31)보다 낮은 수치입니다.
• 의미론적 유지 (Semantic Retention): 절제 연구(Ablation studies) 결과, 기하학적 규제를 제거하면 작용 기전(MoA) 분류 정확도가 감소함(BBBC021에서 72.08%에서 65.08%로)을 보여주었습니다. 이는 중간 진화를 제약하는 것이 메커니즘 신호를 보존하는 데 도움이 됨을 나타냅니다.
• 정성적 분석: 시각화를 통해 FreeBridge가 베이스라인 모델들이 흔히 보이는 텍스처 평활화(texture smoothing)나 불완전한 패턴을 보이는 것과 달리, 화합물 특유의 형태학적 특징(예: 핵 압축, 세포질 수축)을 더 잘 보존함을 확인했습니다.

의의 및 주장
본 논문은 생물학적 전이를 모델링하는 데 있어 경계 정렬 그 이상의 것이 필요하며, **중간 진화의 명시적인 기하학적 근거(geometric grounding)**가 필수적이라고 주장합니다. 경험적으로 지지되는 세포 매니폴드에 확률적 수송을 고정함으로써, FreeBridge는 생물학적으로 해석 가능한 섭동 역학을 생성하는 원칙적인 접근 방식을 제공합니다. 저자들은 이 정식화가 "가상 세포(virtual cell)" 시뮬레이션을 위한 견고한 기초를 제공하며, 종단 정렬 모델이 실제적인 중간 세포 상태를 포착하지 못하는 결정적인 격차를 해결한다고 주장합니다. 또한, 기하학적 제약이 없다면 생성 모델이 통계적으로는 유효한 종단을 생성하더라도 잠재 공간 내에서 생물학적으로 지지되지 않는 영역을 통과할 수 있음을 강조합니다.

한계점
저자들은 FreeBridge가 고정된 인코더에 의존한다는 점을 인정합니다. 즉, 인코더의 사전 학습에 포함되지 않은 형태학적 요소는 수송 학습만으로는 복구될 수 없습니다. 또한, 지원 거리(support-distance) 계산은 미니배치 크기에 따라 제곱($O(B^2)$)으로 스케일링되므로, 더 큰 지원 뱅크를 사용할 경우 근사 최근접 이웃 탐색이 필요할 수 있습니다.