기술 요약: 카운트 데이터용 흐름 매칭 (count-FM)

문제 제기

단일 세포 RNA 시퀀싱 (scRNA-seq) 및 신경 스파이크 트레인 분석과 같은 분야에서 널리 발견되는 고차원 카운트 데이터는 생성 모델링에 고유한 도전을 제시합니다. 이러한 데이터 유형은 이산성, 희소성, 그리고 복잡한 상관 구조를 특징으로 합니다. 기존 심층 생성 접근법들은 종종 이 영역에서 어려움을 겪습니다:

• 범주형 상태 접근법: 카운트를 토큰과 같은 범주형 변수로 취급하는 방법들은 최대 카운트 범위에 비례하여 출력 차원을 확장해야 합니다. 이는 매개변수 비효율성을 초래하며, 인접한 카운트들을 관련 없는 상태로 취급하여 카운트 공간의 고유한 기하학을 무시합니다.
• 연속 변환: 양자화 제거 또는 잠재 표현을 통해 카운트를 연속 공간으로 매핑하는 접근법들은 이산 확률 질량을 연속 밀도로 대체하여 잠재적인 이산 구조를 흐리게 만들 수 있습니다.
• 기존 카운트 모델의 한계: 일부 카운트 전용 점프 모델 (예: Poisson-JUMP) 이 존재하지만, 이들은 종종 임의의 카운트 분포 간의 수송을 위해 설계되지 않았거나 현대적인 흐름 기반 방법들의 효율적인 학습 및 생성 특성을 갖추지 못했습니다.

본 논문은 카운트 공간에서 직접 작동하며, 데이터의 이산 기하학을 존중하고, 범주형 상태 모델의 매개변수 오버헤드 없이 임의의 소스 및 타겟 분포 간의 효율적인 수송을 가능하게 하는 생성 프레임워크의 필요성을 다룹니다.

방법론: count-FM

저자들은 국소 단위 점프를 기반으로 한 연속 시간 출생 - 사망 과정을 기반으로 한 카운트 데이터 전용 흐름 매칭 프레임워크인 count-FM을 제안합니다.

1. 연속 시간 출생 - 사망 과정

이 모델은 소스 분포 $p_0$와 타겟 분포 $p_1$ 간의 전이를 카운트 벡터 공간 $\mathbb{N}_0^d$ 위의 연속 시간 마르코프 점프 과정 (CTMC) 으로 정의합니다. 이 과정은 각 좌표 $i$에 대해 국소 단위 출생 ($+1$) 과 사망 ($-1$) 을 통해 진화합니다:

• 출생률: $\lambda_{t,i}(x)$
• 사망률: $\mu_{t,i}(x)$
• 동역학: 작은 구간 $h$에서 상태 $x$에서 $x + e_i$로 전이할 확률은 $h \lambda_{t,i}(x)$이며, $x - e_i$로 전이할 확률은 $h \mu_{t,i}(x)$입니다. 카운트가 0 일 때 사망은 허용되지 않습니다.

이 매개변수화는 매우 효율적입니다: 가능한 모든 카운트 레벨에 대한 전체 범주형 분포를 예측하는 대신 (이는 최대 카운트에 비례하여 확장됨), 모델은 카운트 범위와 관계없이 차원당 두 개의 비율만 예측합니다 (2d 개의 출력).

2. 조건부 이항 브리지 및 학습

모델을 학습하기 위해 저자들은 고정된 끝점 $(x_0, x_1)$ 사이의 조건부 브리지를 정의합니다. 각 좌표 $i$에 대해 브리지는 이항 분포를 사용하여 구성됩니다:
$$X_t^{(i)} = x_0^{(i)} + \text{sgn}(x_1^{(i)} - x_0^{(i)}) B_t^{(i)}, \quad B_t^{(i)} \sim \text{Binomial}(|x_1^{(i)} - x_0^{(i)}|, t)$$
이 브리지는 조건부 평균이 $x_0$에서 $x_1$로 선형적으로 이동하면서도 정수 값의 중간 상태를 유지하도록 보장합니다.

이 조건부 브리지에 질량 보존 (콜모고로프 전진 방정식) 을 강제함으로써, 저자들은 조건부 출생률과 사망률에 대한 폐쇄형 식을 유도합니다:
$$\lambda_{t,i}(x | x_0, x_1) = \frac{(x_1^{(i)} - x)^+}{1-t}, \quad \mu_{t,i}(x | x_0, x_1) = \frac{(x - x_1^{(i)})^+}{1-t}$$
여기서 $(\cdot)^+$는 양수 정류 (positive rectification) 를 나타냅니다.

3. 학습 목적 함수

모델은 조건부 브리지 과정과 모델이 학습한 과정 사이의 경로 공간 쿨백 - 라이블러 (KL) 발산을 최소화함으로써 학습됩니다. 이는 국소 비율 매칭 목적 함수로 분해됩니다:
$$\mathcal{L}_{train}(\theta) = \mathbb{E} \left[ \sum_{i=1}^d \ell(\lambda_{t,i}, \lambda_{\theta,i}) + \sum_{i=1}^d \ell(\mu_{t,i}, \mu_{\theta,i}) \right]$$
여기서 $\ell(u, v) = v - u \log v$는 일반화된 KL 항입니다. 이는 식별 가능한 조건부 비율을 학습함으로써 주변 전이 비율의 시뮬레이션 없는 학습을 가능하게 합니다.

4. 샘플링 및 커플링

• 샘플링: 샘플은 1 차 국소 점프 이산화를 사용하여 $t=\epsilon_t$에서 $t=1-\epsilon_t$까지 학습된 출생 - 사망 과정을 forward 로 시뮬레이션하여 생성됩니다.
• 끝점 커플링: 소스 및 타겟 샘플을 짝짓기 위한 두 가지 커플링 전략이 탐색됩니다:
  • 독립 커플링: 무작위 짝짓기 ($p_0(x_0)p_1(x_1)$).
  • 최적 수송 (OT) 커플링: 대칭 포아송 비용을 사용한 미니배치 OT 짝짓기. 이는 일반적으로 더 낮은 곡률을 가진 전이를 유도하여 해석 가능성과 샘플링 효율성을 향상시킵니다.

5. 조건부 생성

조건부 작업의 경우, 모델은 공변량 $y$를 포함하고 **클래시퍼 프리 가이드 (CFG)**를 활용합니다. 학습 중에는 조건 변수가 무작위로 드롭됩니다. 샘플링 시에는 조건부 및 무조건부 비율이 가이드 스케일 $w$를 통해 결합되어 조건의 강도를 조절합니다.

주요 기여

1. count-FM 프레임워크: 국소 출생 - 사망 동역학에 기반한 카운트 데이터용 흐름 매칭 프레임워크로, 중간 상태가 카운트 값을 유지하면서 카운트 공간에서 직접 수송을 가능하게 합니다.
2. 실용적인 학습 체계: 조건부 이항 브리지를 통해 주변 전이 비율을 학습하는 방법으로, 범주형 상태 모델링의 복잡성을 피하면서도 흐름 매칭의 효율성을 유지합니다.
3. 매개변수 효율성: 전체 범주형 분포 대신 국소 출생 및 사망 비율만 예측하는 ($2d$개의 출력) 공식으로, 큰 카운트 범위를 가진 고차원 데이터에 적합합니다.

실험 결과

1. 시뮬레이션 (2D 감마 - 포아송 혼합)

• 성능: count-FM 과 count-FM-OT 는 D3PM, discrete-FM, tauLDR, SEDD, Poisson-JUMP 를 포함한 베이스라인 대비 가장 우수한 샘플 품질 (최저 2-워asserstein 거리 및 MMD$^2$) 을 달성했습니다.
• 효율성: count-FM 은 범주형 상태 베이스라인 (14,000~67,000+) 에 비해 훨씬 적은 학습 가능 매개변수 (2,372 개) 로 이러한 결과를 달성했습니다.
• 기하학: 중간 상태는 카운트 공간에서 매끄럽게 진화했습니다. OT 커플링 버전은 더 직선적인 전이 경로를 따랐으며, 더 적은 함수 평가 (NFEs) 로 타겟 품질에 도달했습니다.

2. 단일 세포 RNA-seq (치상회)

• 무조건부 생성: 2,930 개의 세포와 13,913 개의 유전자로 구성된 데이터셋에서 count-FM 은 샘플 품질 (W2 및 MMD$^2$) 에서 잠재 변수 베이스라인 (scVI) 과 잠재 확산 모델 (scDiffusion, scLDM) 을 능가했으며, 경쟁사 (약 29M~39M) 에 비해 적은 매개변수 (약 9.8M) 를 사용했습니다.
• 수송 (P12 에서 P35 로): 이 모델은 출생 후 12 일에서 35 일로의 발달 수송을 성공적으로 모델링했습니다. 계통 제한 OT 커플링을 사용하여 수송 경로는 알려진 발달 구조 (예: 신경아세포에서 과립세포로의 진행) 를 존중했으며, 카운트 공간에 머무르는 해석 가능한 중간 분포를 생성했습니다.

3. 신경 스파이크 트레인 (해마/내후각피질)

• 조건부 생성: 이 모델은 선형화된 위치와 주행 방향에 조건부인 스파이크 카운트를 생성하는 데 적용되었습니다.
• 비교: 가이드 스케일 $w=1$인 count-FM 은 결정론적 MLP 회귀 모델 및 포아송 MLP 베이스라인보다 우수한 성능을 보였습니다.
• 상관 구조: 포아송 MLP 가 인구 상관 관계를 과소평가한 것과 달리, count-FM 은 뉴런 간 의존 구조를 더 잘 보존했습니다.
• 가이드 트레이드오프: $w=1$은 공간 튜닝과 분포 보정 사이의 가장 좋은 균형을 제공했습니다. $w=2$는 플레이스 필드를 날카롭게 만들었지만 진폭 보정을 왜곡했습니다.

중요성 및 주장

본 논문은 count-FM 이 카운트 공간에서 직접 생성 및 수송을 가능하게 함으로써 고차원 카운트 데이터 모델링을 위한 효과적인 프레임워크를 제공한다고 주장합니다. 그 주요 중요성은 다음과 같습니다:

• 기하학적 충실도: 국소 단위 점프를 통한 수송 모델링으로 이산 기하학을 존중하여 연속 근사의 흐려지는 효과를 방지합니다.
• 확장성: 매개변수 효율성은 범주형 상태 방법이 계산적으로 불가능해지는 고차원 응용 (예: 수천 개의 유전자) 에 실용적으로 만듭니다.
• 해석 가능성: 수송 경로에서의 중간 상태는 유효한 카운트 샘플이므로, 발달 생물학과 같은 응용 분야에서 의미 있는 전이 분포를 검사할 수 있습니다.
• 다용도성: 이 프레임워크는 무조건부 생성, 임의 분포 간의 수송, 그리고 조절 가능한 가이드를 통한 조건부 생성을 지원합니다.

저자들은 조건부 이항 브리지가 상대적으로 단순하여 모든 가능한 중간 경로를 탐색하지 못할 수 있으며, 현재 1 차 이산화가 더 정확한 시뮬레이션 체계로 개선될 수 있음을 한계로 인정합니다. 그러나 시뮬레이션 및 생물의학 데이터셋에 걸친 경험적 결과는 count-FM 이 경쟁 접근법보다 더 높은 모델링 효율성으로 높은 샘플 품질을 달성함을 보여줍니다.