個体を追跡できなくても、集団の分布は比較できます。最適輸送（optimal transport、OT）は、ある分布を別の分布へ移す際に移動コストが最小になる配分を探します。得られるのは制約に整合するカップリングであり、実際に誰が誰へ変化したかを復元した名簿ではありません。

前提と時間配分。 確率分布、行列、距離、PCAの状態空間を復習します。90分を対応付けの問題15分、制約とエントロピー25分、Python30分、解釈20分に分けます。balanced OTの定式化、周辺分布の検証、正則化の効果の説明が目標です。共通環境、完全なコード、実習一式を使います。追加パッケージはPOTで、otとしてインポートします。

個体IDのない2枚の写真

初期の細胞は状態地図の中央に集まり、後期には二つの枝に分かれています。各時刻で測定した細胞は異なります。最近傍だけで結ぶと、全出発点が同じ便利な到着点を選び、残りの到着集団を無視してしまう場合があります。

小さな2D例から始めましょう。出発点を$(0,0),(1,0)$、到着点を$(0,1),(1,1)$とし、それぞれに質量$1/2$を与えます。二乗ユークリッド距離なら垂直移動のコストは1、交差する移動は2です。最適計画は各出発点から半分の質量を垂直に送ります。合計コストは2ではなく1です。それぞれの移動が運ぶのは全体の半分だからです。

到着点が同じ場所なら、複数の個体対応が同じコストを持つかもしれません。ソルバーは未観測の証拠を作れません。この曖昧さは、大きく美しい図を描いても残ります。

必要なのは質量保存

出発側の質量を$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. $$

行制約は各出発点が持つ質量を過不足なく送り、列制約は各到着点が必要な質量を過不足なく受け取ることを表します。1点の質量を複数の到着点へ分割できるため、1対1の対応である必要はありません。

出発側が50点、到着側が80点でも、各点に$1/50$と$1/80$を割り当てれば全体はどちらも1です。これは実世界の細胞数やセッション数が等しいという意味ではありません。正規化により絶対量の情報を捨てています。

コスト関数には科学的な仮定が入ります。ユークリッド距離は表現空間での直線的な移動を好みます。測地距離は空白を横切る移動を高価にできますが、近傍グラフという仮定が増えます。特徴のスケール、バッチ効果、欠落した状態も「安い移動」を変えます。

Wasserstein距離と根の関係

距離$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距離と呼ぶと、指数と距離の意味が曖昧になります。

実習では小さな線形計画問題で正確な$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) $$

を反復します。除算は要素ごとです。行と列を交互に調整して二つの周辺分布を満たします。対数領域の実装は小さなカーネル値がゼロへアンダーフローする問題を抑えます。この計算上のつながりはCuturi, 2013で扱われています。

正則化された目的値は、分布をそれ自身へ移す場合も一般にはゼロになりません。バイアスを補正する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公式ドキュメントで確認できます。

3種類のエントロピー強度による輸送計画

指定シードの実行では、正確な$W_2$はPCA座標の単位で約7.14です。コストを中央値で割り、$\varepsilon=0.05,0.2,1$を比較します。周辺分布の誤差は$10^{-6}$未満を要求します。上段は各出発点の大きい接続2本を示します。下段は両軸をPC 2で並べ、共通の色尺度ですべての条件付き質量を表示します。保存の検証には行列全体を使います。

エントロピーが大きいほど線が広く分散する傾向があります。密な図は生物学的な遷移を多く発見した証拠ではありません。不確実性を分散させる最適化上の選好が変わったのです。

生物とキャラクターチャットへの適用

細胞は増殖、分裂、死亡します。正規化したbalanced OTは形の変化を記述しても増殖を隠すことがあります。Unbalanced OTでは保存則を緩める罰則や成長モデルを導入します。Waddington-OTは、異なる時点の細胞集団から発生経路を分析する応用例です。Schiebinger et al., 2019

会話では初期と後期のターンや、コホート間の状態分布を比較できます。変化には物語の進行だけでなく、ユーザーの選択、打ち切り、キャラクター構成の違いも含まれます。セッションIDがあるなら、実際の遷移と推定された結合を比較して追加情報を使いましょう。製品バージョン間の結合だけではプロンプトの因果効果は特定できません。

演習

  1. 計画を検証する。 $a=b=(1/2,1/2)$で$\gamma=\left[\begin{smallmatrix}.5&0\\.5&0\end{smallmatrix}\right]$は可能ですか。答え: 行は正しい一方、列和が$(1,0)$なので不可能です。
  2. 根を確認する。 二乗距離の正則化なし最適コストが9なら$W_2$は？ 答え: 3。正則化した計画の移動コストが9でも、正則化なしの最適値が9とは限りません。
  3. 幾何を変える。 一つの特徴を10倍にして再計算します。ヒント: 二乗距離ではその寄与が100倍になります。前処理が輸送モデルの一部である理由を説明してください。

次は、分布を二つの端点の間で動かす連続ベクトル場を学びます。