M6.3 五萬個倉庫怎麼算：Sinkhorn 與 Minibatch 近似

本篇重用3 篇

二十五億格的表

三個倉庫對三家店的運輸表只有九格；五萬對五萬卻有 $2.5\times10^9$ 格。若每格用 8 bytes，單是成本矩陣就約二十 GB，還沒算 solver 的其他資料。精確求解的成本依演算法與結構而定，不能把所有 transport LP 都寫成同一個 $n^3$ 複雜度。

而「五萬個」在現實裡很常見：五萬張圖片的特徵、五萬個顧客、五萬個粒子。你需要 $W_2$ 或最佳配對，但不可能精確算。

如果整張表太大，可以改問題，也可以縮小一次處理的範圍。前者允許較分散、容易計算的配對；後者每次只替一小批貨配對。這兩種近似各改了什麼？算出的成本還能跟原本的 OT 比較嗎？

課堂提問Q1

翻成數學。「差不多省、但好算」可以怎麼寫成一個修改過的最佳化問題？「在小組裡配」估的是什麼量——它和 $W_2(\mu,\nu)$ 是同一個東西嗎？

先想一想，再展開看整理後的答案

課堂上第一條路的候選很多（貪婪配對、只看最近的幾家店…），但有一個修改特別乾淨：在成本上加一個懲罰「太集中」的項。運輸表 $\pi$ 越集中（每列只有一格不為零）越像精確解、越難算；越分散（接近獨立 coupling）越平滑。用 entropy $H(\pi)=-\sum\pi_{ij}\log\pi_{ij}$ 量「分散程度」，解

\min_{\pi\in\Pi(a,b)}\ \sum_{ij}\pi_{ij}c_{ij}\;-\;\varepsilon\,H(\pi),\qquad\varepsilon>0 .

這叫 entropic optimal transport。 $\varepsilon$ 是「願意為了好算而多付多少油錢」的兌換率。它的最佳解會比精確解分散、成本偏高——這是第一條路的偏差方向，下面會證明。

第二條路是抽 $B$ 個來源與 $B$ 個終點，形成 empirical distributions $\hat\mu_B,\hat\nu_B$ ，再解組內 OT。下面比較的是 $\mathbb E[W_2^2(\hat\mu_B,\hat\nu_B)]$ 與 $W_2^2(\mu,\nu)$ ：平方成本的期望不低於 population 值，卻不代表每個抽出的 batch 都比較高，也不能直接把開根號後的期望套進同一證明。

把稜角磨圓，或只看一小塊

有限 transport LP 至少能選到一個 vertex 最佳解，它至多有 $m+n-1$ 個非零 entries；不是每個最佳解都一定是 vertex。這種稀疏結構有利於精確求解，但大矩陣仍可能難處理。

Entropic OT 把稜角磨圓。 加了 $-\varepsilon H(\pi)$ 之後，目標是 strictly convex 的，最佳解在多面體內部、每一格都嚴格大於零，而且有一個極簡的形狀：一個固定的矩陣 $K_{ij}=e^{-c_{ij}/\varepsilon}$ ，左右各乘一個對角縮放。要找到它，只要輪流把列和調成 $a$ 、把欄和調成 $b$ ——每一輪是一次矩陣乘向量。這就是 Sinkhorn。

Minibatch 只看一小塊。 五萬對五萬看不完，就看兩百對兩百。組內的問題小到能精確解。代價是每一組都在「以偏概全」：組裡剛好沒有某個區域的店，那個區域的倉庫就被迫送遠。

大規模 OT：entropic 與 minibatch 兩種近似

圖：Sinkhorn 用 ε 磨圓稜角；minibatch 用 B 限制視野，兩者造成不同方向的偏差。

Sinkhorn

解的形狀。 對 entropic 問題寫 Lagrangian，對每個 $\pi_{ij}$ 微分令其為零：

c_{ij}+\varepsilon(\log\pi_{ij}+1)-f_i-g_j=0 \quad\Longrightarrow\quad \pi_{ij}=\underbrace{e^{(f_i-\varepsilon)/\varepsilon}}_{u_i}\;\underbrace{e^{-c_{ij}/\varepsilon}}_{K_{ij}}\;\underbrace{e^{g_j/\varepsilon}}_{v_j},

$f,g$ 是列和、欄和兩組約束的 Lagrange multiplier。所以最佳解一定是 $\pi=\mathrm{diag}(u)\,K\,\mathrm{diag}(v)$ ： $K$ 由成本與 $\varepsilon$ 決定、只算一次；未知的只有 $m+n$ 個縮放因子 $u,v$ ，不是 $mn$ 個格子。

Sinkhorn 演算法。 把 $u,v$ 調到列和是 $a$ 、欄和是 $b$ 。兩個約束分開看各自是一行：

\text{列和}=a:\ \ u_i=\frac{a_i}{(Kv)_i},\qquad \text{欄和}=b:\ \ v_j=\frac{b_j}{(K^{\!\top}u)_j}.

固定 v 更新 u，再固定 u 更新 v，就輪流滿足列與欄的邊際。對正的 K 與正 marginals，Sinkhorn–Knopp [3] 的結果給出縮放後矩陣的收斂；u、v 本身有互相抵消的比例自由度。Cuturi [1] 將它用於 entropic OT。每輪是 $O(mn)$ 的矩陣運算，但所需輪數依 $\varepsilon$ 、矩陣與容許誤差而定，不能保證幾十輪完成。

$\varepsilon$ 兩端的極限。 $\varepsilon\to\infty$ ： $K\to$ 全 1 的矩陣， $\pi\to$ 獨立 coupling $ab^\top$ ——完全不看成本。 $\varepsilon\to0$ ： $-\varepsilon H$ 消失，回到精確 OT（但 $K=e^{-c/\varepsilon}$ 的元素會下溢到 0，數值上要在 $\log$ 域做——展開框）。中間的 $\varepsilon$ 是取捨：越小越準、越大越快收斂、越穩定。

偏差的方向。 精確解 $\pi^*$ 的成本 $\langle\pi^*,c\rangle$ 是所有 coupling 裡最小的；entropic 解 $\pi_\varepsilon$ 是另一個 coupling，所以 $\langle\pi_\varepsilon,c\rangle\ge\langle\pi^*,c\rangle$ ——entropic 解的運輸成本永遠不低於真實 OT。它多付的成本換來的是 $H(\pi_\varepsilon)\ge H(\pi^*)$ ：拆得更散。

展開細節log 域的 Sinkhorn、Sinkhorn divergence、以及 ε 的偏差量級

數值上 K = e^{−c/ε} 在 ε 小時會下溢（c/ε ≈ 700 就變成 0）。改在 log 域做：令 f = ε log u、g = ε log v，更新變成 f_i = ε log a_i − ε log Σ_j exp((g_j − c_{ij})/ε)，用 log-sum-exp 穩定計算。這是實務上的標準寫法。

對固定有限、總質量為 1 的 $m\times n$ 問題，最優性給出 $0\le\langle\pi_\varepsilon,c\rangle-\langle\pi^*,c\rangle\le\varepsilon(H(\pi_\varepsilon)-H(\pi^*))\le\varepsilon\log(mn)$ 。這是有限問題的界，不能直接當成所有連續問題的速率。

Genevay 等人 [5] 使用的 Sinkhorn divergence 是 $S_\varepsilon(\mu,\nu)=\mathrm{OT}_\varepsilon(\mu,\nu)-\frac12\mathrm{OT}_\varepsilon(\mu,\mu)-\frac12\mathrm{OT}_\varepsilon(\nu,\nu)$ ，這裡 $\mathrm{OT}_\varepsilon$ 指包含 entropy regularization 的最優目標值，不是只取 $\langle\pi_\varepsilon,c\rangle$ 。扣除自我項讓相同分布的值為零；其他距離性質仍需要成本與 kernel 的條件。

Minibatch 的偏差

抽 $B$ 個 $\mu$ 的樣本、 $B$ 個 $\nu$ 的樣本，各成經驗分佈 $\hat\mu_B,\hat\nu_B$ ，算 $W_2(\hat\mu_B,\hat\nu_B)$ ，對抽樣取期望。兩件事可以說清楚。

距離偏高。 這可以證明： $\mathrm{OT}(\mu,\nu)=\min_\pi\langle\pi,c\rangle$ 對 $(\mu,\nu)$ 這一對是jointly convex 的（兩個 coupling 的凸組合仍是對應邊際凸組合的 coupling，所以最小值不會高於凸組合），而 $\mathbb E[\hat\mu_B]=\mu$ 、 $\mathbb E[\hat\nu_B]=\nu$ ，M5.0 立刻給

\mathbb E\big[W_2^2(\hat\mu_B,\hat\nu_B)\big]\;\ge\;W_2^2\big(\mathbb E\hat\mu_B,\mathbb E\hat\nu_B\big)=W_2^2(\mu,\nu).

例如兩端是同一個連續分布，population $W_2=0$ ，兩份獨立抽出的有限點集卻幾乎必然不同，empirical 距離因而為正。收斂速率依維度、矩條件與分布規則性而變，不能一律寫成 $B^{-1/d}$ [2,4]。

整體配對。 把每批最佳 coupling 平均，得到 $\bar\pi_B=\mathbb E[\pi_B^*]$ 。它的兩端 marginals 仍是 $\mu,\nu$ ，因此是合法 coupling，成本不低於 population OT。但有限 batch 的配對通常不是 deterministic map； $B=1$ 時更直接得到 independent coupling。適當矩條件下，大 batch 的成本可趨近 population OT；要宣稱 coupling 趨近某個指定 $\pi^*$ ，還需要唯一性等條件。

兩種近似殊途同歸：都給出一個比精確解更分散的 coupling、成本偏高；一個用 $\varepsilon$ 控制、一個用 $1/B$ 控制。

該用哪個

回答起點問題。第一條路（entropic）成本偏高、配對偏散，偏差由 $\varepsilon$ 控制，代價是 $O(mn)$ 的矩陣乘向量——五萬對五萬在 GPU 上可行，但 $K$ 要 $2.5\times10^9$ 格的記憶體（可以分塊）。第二條路（minibatch）成本偏高、配對偏亂，偏差由 $B$ 控制，記憶體只要 $B^2$ ，但每組的估計有隨機性、要多抽幾組平均。

選擇前先問：需要距離數值，還是訓練時的配對？Sinkhorn 修改目標，minibatch 限制這次能看到的端點，兩者也可以合用。需要依記憶體、精度與用途選擇，不能說其中一種永遠最好或是唯一可行方法。

有限成本與 $\varepsilon>0$ 讓理論 K 嚴格為正；計算時需避免 underflow，並確認兩側 marginal residual。只要求「沒有全零列欄」對一般帶零的矩陣還不夠。Minibatch 的期望平方成本比較則要用 iid sampling：上界方向成立在期望下，不是每一批都一定較高；entropic raw transport cost 也要和含 entropy 的 objective 分清楚。

縮小成兩千個，把三個數並排

import numpy as np
from scipy.special import logsumexp
from scipy.optimize import linear_sum_assignment
rng = np.random.default_rng(0)
n, d = 2000, 2
X = rng.normal(0, 1, (n, d)); Y = rng.normal(0, 1, (n, d)) @ np.diag([2, 0.5]) + np.array([3, 0])
print(9 + (1-2)**2 + (1-0.5)**2)                                            # closed form W₂² = 10.25（μ=N(0,I)、ν=N((3,0), diag(4,¼))）
C = ((X[:, None, :] - Y[None, :, :])**2).sum(-1)                            # 2000×2000 成本矩陣
r, c = linear_sum_assignment(C); print(C[r, c].mean())                      # 2000 個樣本的精確 OT ≈ 10.73（抽樣誤差在上方）
def sinkhorn_log(C, eps, iters=1000):                                        # log 域的 Sinkhorn：輪流修列和、欄和
    m = C.shape[0]; loga = -np.log(m) * np.ones(m); f = np.zeros(m); g = np.zeros(m)
    for _ in range(iters):
        f = eps * (loga - logsumexp((g[None, :] - C) / eps, axis=1))       # 修列和
        g = eps * (loga - logsumexp((f[:, None] - C) / eps, axis=0))       # 修欄和
    logP = (f[:, None] + g[None, :] - C) / eps; P = np.exp(logP)
    return (P * C).sum(), -(P * logP).sum()
for eps in [0.05, 0.2, 1.0, 5.0]:
    cost, H = sinkhorn_log(C, eps); print(eps, round(cost, 2), round(H, 2))  # 成本 10.77 → 13.31 單調上升；entropy 上升
def minibatch_w2(C, B, reps=200):                                           # 小組內用 assignment 精確解
    vals = []
    for _ in range(reps):
        i = rng.choice(n, B, replace=False); j = rng.choice(n, B, replace=False)
        S = C[np.ix_(i, j)]; r, c = linear_sum_assignment(S); vals.append(S[r, c].mean())
    return np.mean(vals)
for B in [8, 32, 128, 512]: print(B, round(minibatch_w2(C, B), 2))         # 11.91 → 10.77：從上方逼近 10.73
# 破壞：μ = ν（真值 0）
Y2 = rng.normal(0, 1, (n, d)); C2 = ((X[:, None, :] - Y2[None, :, :])**2).sum(-1)
for B in [8, 512]: print('same', B, round(minibatch_w2(C2, B, 50), 3))     # 1.33、0.064：相同分佈被判成有差

三組數字並排。closed form 真值 $10.25$ ；用全部 2000 個樣本精確解 assignment 得 $10.73$ （有限樣本的估計本身就偏高一點）。Sinkhorn 成本隨 $\varepsilon$ 從 $10.77$ 單調升到 $13.31$ 、entropy 從 $11.6$ 升到 $15.0$ ——偏高、偏散，而且 $\varepsilon=0.05$ 就已經在 $K=e^{-C/\varepsilon}$ 下溢的邊緣，所以程式直接寫成 $\log$ 域。minibatch 從 $B=8$ 的 $11.91$ 降到 $B=512$ 的 $10.77$ ，從上方逼近 $10.73$ 。

這裡有一個實作上很會騙人的地方，最好自己驗一次。把 iters 從 1000 降回 300， $\varepsilon=0.05$ 會給出 $10.49$ ——比精確解 $10.73$ 還低，直接違反上面那個「entropic 解的成本不低於真實 OT」。原因不是定理錯了，是那個 $P$ 還不是合法的 coupling：300 輪之後列邊際的 $L_1$ 誤差還有 $0.038$ ，質量沒有守好，成本自然可以被壓到下界以下。 $\varepsilon$ 越小收斂越慢，所以看到「比精確解還低」的第一件事是去檢查邊際，不是去懷疑定理。

刻意違反一個假設。 把 $Y$ 換成與 $X$ 同分佈的另一組樣本（真值 $W_2=0$ ）：2000 對 2000 的精確解給 $0.02$ ，minibatch 在 $B=8$ 時給 $1.33$ 、 $B=512$ 時給 $0.067$ ——它會把「兩個相同的分佈」判成「有差」， $B$ 越小越嚴重。再把 $\varepsilon$ 壓到 $10^{-3}$ 且改回普通域的 $K=e^{-C/\varepsilon}$ ：矩陣大量下溢成 0、K @ v 出現零、除法變成 inf——這就是為什麼要在 $\log$ 域做。

互動 demo：切換兩種近似並逐步更新，觀察 coupling、成本、entropy 與邊際誤差。

檢查理解

想一想

你用 Sinkhorn 算兩個點雲的距離，發現把同一個點雲跟自己比，算出來的成本不是零。這是：

參考文獻

Cuturi, M. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. NeurIPS 2013.（把 entropic 正則化與 Sinkhorn 迭代帶進大規模 OT。）
Peyré, G., Cuturi, M. Computational Optimal Transport. Foundations and Trends in Machine Learning 11(5–6), 2019.（第 4 章：entropic OT、Sinkhorn 的推導與收斂、log 域實作、 $\varepsilon$ 的偏差量級；第 8.5 節：經驗分佈的收斂速率。）
Sinkhorn, R., Knopp, P. Concerning Nonnegative Matrices and Doubly Stochastic Matrices. Pacific Journal of Mathematics 21(2), 1967.（交替正規化收斂到唯一的對角縮放。）
Fatras, K., Zine, Y., Flamary, R., Gribonval, R., Courty, N. Learning with Minibatch Wasserstein: Asymptotic and Gradient Properties. AISTATS 2020.（minibatch OT 的偏差、 $\mu=\nu$ 時不為零、隨 $B$ 的收斂。）
Genevay, A., Peyré, G., Cuturi, M. Learning Generative Models with Sinkhorn Divergences. AISTATS 2018.（Sinkhorn divergence：扣掉自我運輸成本。）