개체를 계속 추적하지 못해도 집단의 분포는 비교할 수 있습니다. 최적 수송(optimal transport, OT)은 한 분포를 다른 분포로 옮길 때 이동 비용이 가장 작은 배분을 찾습니다. 결과는 주어진 제약을 만족하는 결합(coupling)이며, 실제로 어느 개체가 누구로 변했는지를 복원한 명단은 아닙니다.
선수 지식과 수업 계획. 확률분포, 행렬, 거리, PCA 상태 공간을 복습합니다. 90분을 매칭 문제 15분, 제약과 엔트로피 25분, Python 30분, 해석 20분으로 나눕니다. 목표는 balanced OT의 목적식 작성, 주변분포 제약 검증, 정규화 효과 설명입니다. 공통 환경, 전체 코드, 전체 ZIP을 사용합니다. 추가 패키지는 POT이며 Python에서는 ot로 불러옵니다.
개체 ID가 없는 두 장의 사진
이른 시점의 세포들은 상태 지도의 중앙에 모여 있고, 늦은 시점에는 두 갈래로 나뉘어 있습니다. 두 시점에서 측정한 세포는 서로 다릅니다. 가장 가까운 점을 연결하고 싶지만, 모든 출발점이 같은 도착점만 선택하면 실제로 관측한 다른 도착 상태를 설명하지 못합니다.
더 작은 2D 예를 봅시다. 출발점은 $(0,0),(1,0)$이고 도착점은 $(0,1),(1,1)$입니다. 각 점의 질량을 $1/2$로 둡니다. 제곱 유클리드 거리에서 수직 이동의 비용은 1, 교차 이동의 비용은 2입니다. 최적 계획은 각 출발점의 절반 질량을 수직으로 보냅니다. 총비용은 2가 아니라 1입니다. 각 이동이 전체 집단의 절반만 운반하기 때문입니다.
도착점 둘이 같은 위치라면 개체 ID를 여러 방식으로 연결해도 동일한 비용이 나올 수 있습니다. 최적화는 관측하지 않은 증거를 만들어주지 않습니다. 복잡하고 예쁜 그림에서도 이 모호성은 남습니다.
핵심은 질량 보존 제약이다
출발 질량을 $a_i$, 도착 질량을 $b_j$라 하고 각각의 합을 1로 둡니다. $\gamma_{ij}$는 출발점 $i$에서 도착점 $j$로 보내는 질량입니다.
$$ \min_{\gamma\geq0}\sum_{i,j}C_{ij}\gamma_{ij},\qquad \gamma\mathbf 1=a,\quad \gamma^\top\mathbf 1=b. $$행 제약은 각 출발점이 가진 양만큼 정확히 보내라는 뜻입니다. 열 제약은 각 도착점이 필요한 양을 정확히 받으라는 뜻입니다. 한 출발점의 질량을 여러 도착점으로 나눌 수 있으므로 일대일 대응일 필요가 없습니다.
출발 관측이 50개, 도착 관측이 80개여도 각각 $1/50$, $1/80$의 가중치를 주면 총질량은 모두 1입니다. 이 정규화는 실제 세포 수나 사용자 수가 같다는 뜻이 아닙니다. 절대 개체 수 정보는 사라집니다.
비용 함수에는 분석자의 가정이 들어갑니다. 유클리드 거리는 상태 공간의 직선 이동을 선호합니다. 측지 거리는 빈 공간을 가로지르는 연결을 비싸게 만들 수 있지만, 이웃 그래프라는 가정이 추가됩니다. 특징 스케일, 배치 효과, 빠진 상태도 어떤 이동이 저렴한지 결정합니다.
Wasserstein distance에는 지수가 있다
거리 $d$와 $p\geq1$에 대해
$$ W_p(\mu,\nu)=\left[\min_{\gamma\in\Pi(a,b)} \sum_{ij}d(x_i,y_j)^p\gamma_{ij}\right]^{1/p} $$입니다. $C_{ij}=\|x_i-y_j\|^2$이면 정규화하지 않은 최적 비용이 $W_2^2$이고, 제곱근을 취해야 $W_2$입니다. $C=d$이면 최적 비용은 $W_1$입니다. 모든 수송 목적값을 뭉뚱그려 Wasserstein distance라고 부르면 지수와 거리의 의미를 놓칩니다.
실습에서는 작은 선형계획 문제로 정확한 $W_2$를 계산합니다. 엔트로피 정규화를 적용한 수송 계획의 이동 비용은 별도 항목으로 보고합니다. 값이 비슷해도 동일한 양은 아닙니다.
엔트로피는 질량을 분산시킨다
다음 정규화 항을 더합니다.
$$ \min_{\gamma\in\Pi(a,b)}\langle C,\gamma\rangle +\varepsilon\sum_{ij}\gamma_{ij}(\log\gamma_{ij}-1). $$$\varepsilon$이 크면 질량을 여러 연결에 나누는 계획을 선호합니다. 작으면 이동 비용을 더 중시하지만 수치 계산이 어려워질 수 있습니다. $\varepsilon$은 비용과 같은 단위를 가지므로 비용 스케일을 모른 채 0.2라는 숫자만 비교하면 의미가 없습니다.
$K_{ij}=\exp(-C_{ij}/\varepsilon)$로 두면 Sinkhorn 반복은
$$ \nu\leftarrow a/(K\omega),\qquad \omega\leftarrow b/(K^\top\nu),\qquad \gamma=\operatorname{diag}(\nu)K\operatorname{diag}(\omega) $$입니다. 나눗셈은 원소별 연산입니다. 행과 열의 크기를 번갈아 조정해 두 주변분포를 맞춥니다. 로그 공간 구현은 매우 작은 커널 값이 0으로 언더플로되는 문제를 줄입니다. 엔트로피 수송과 효율적인 계산의 연결은 Cuturi, 2013에서 다룹니다.
정규화된 목적값은 일반적으로 자기 자신과 비교해도 0이 아닐 수 있습니다. 편향을 제거한 Sinkhorn divergence는 별도 정의이며 이번 실습에서 계산하지 않습니다.
Python: 그림보다 먼저 보존을 검증한다
전체 코드는 서로 다른 학습 세션 집합에서 이른 시점과 늦은 시점의 관측을 선택해 개체 대응을 의도적으로 없앱니다. PCA도 학습 세션으로 적합합니다. A, B가 두 시점의 투영 좌표일 때 핵심 계산은 다음과 같습니다.
import numpy as np
import ot
from common import make_dataset, prepare_embedding
from lesson10 import population_snapshots
data = make_dataset()
_, Z, _, _ = prepare_embedding(data)
A, B = population_snapshots(data, Z)
a = np.full(len(A), 1 / len(A))
b = np.full(len(B), 1 / len(B))
cost = ot.dist(A, B, metric="sqeuclidean")
exact_plan = ot.emd(a, b, cost)
exact_w2 = np.sqrt(np.sum(exact_plan * cost))
normalized_cost = cost / np.median(cost)
plan = ot.sinkhorn(a, b, normalized_cost, reg=.2,
method="sinkhorn_log", numItermax=5000,
stopThr=1e-10)
assert np.max(abs(plan.sum(axis=1) - a)) < 1e-6
assert np.max(abs(plan.sum(axis=0) - b)) < 1e-6
공통 환경에서 python lesson10.py --output outputs로 실행합니다. 전체 스크립트가 위 코드에 필요한 A, B를 구성합니다. 각 함수의 인터페이스는 POT 공식 문서에 설명되어 있습니다.
제공한 시드에서 정확한 $W_2$는 PCA 좌표 단위로 약 7.14입니다. 비용을 중앙값으로 나눈 뒤 $\varepsilon=0.05,0.2,1$을 비교합니다. 주변분포 오차는 반드시 $10^{-6}$ 미만이어야 합니다. 위 패널은 출발점마다 가장 큰 연결 두 개를 표시합니다. 아래 히트맵은 두 축을 PC 2 순서로 정렬하고 공통 색상 범위로 모든 조건부 질량을 보여줍니다. 보존 검증에는 전체 행렬을 사용합니다.
엔트로피가 커지면 연결이 넓게 퍼지는 경향이 있습니다. 선이 많아졌다고 생물학적 전이를 더 많이 발견한 것은 아닙니다. 최적화에서 불확실성을 분산시키는 선호가 바뀐 것입니다.
바이오와 캐릭터 채팅에서의 해석
세포는 증식하고 분열하며 죽습니다. 정규화된 balanced OT는 분포의 모양 변화는 설명해도 증식을 숨길 수 있습니다. Unbalanced OT는 질량 보존을 완화하거나 성장 모델을 도입하는 방향입니다. Waddington-OT는 서로 다른 시점의 세포 집단에서 발달 경로를 분석하는 응용 예입니다. Schiebinger et al., 2019
캐릭터 채팅에서는 초기와 후기 턴의 상태 분포, 혹은 서로 다른 코호트의 상태 분포를 비교할 수 있습니다. 차이는 이야기 진행뿐 아니라 사용자 선택, 중도절단, 캐릭터 구성 변화에서도 생깁니다. 세션 ID가 있다면 실제 전이와 OT 결과를 비교해 추가 증거를 활용해야 합니다. 제품 버전 사이의 수송만으로 프롬프트의 인과 효과를 알 수는 없습니다.
연습문제
- 제약 확인. $a=b=(1/2,1/2)$일 때 $\gamma=\left[\begin{smallmatrix}.5&0\\.5&0\end{smallmatrix}\right]$는 가능한가요? 답: 행 합은 맞지만 열 합이 $(1,0)$이므로 불가능합니다.
- 제곱근 확인. 제곱거리의 비정규화 최적 비용이 9라면 $W_2$는? 답: 3입니다. 정규화된 계획의 이동 비용이 9인 것만으로 비정규화 최적 비용도 9라고 할 수는 없습니다.
- 기하 바꾸기. 한 특징을 10배로 늘려 비용과 계획을 다시 계산하세요. 힌트: 제곱거리에서 해당 특징의 기여는 100배입니다. 전처리가 수송 모델의 일부인 이유를 설명해 보세요.
다음 강의에서는 분포를 두 끝점 사이로 움직이는 연속 벡터장을 학습합니다.