CIS6270 / lecture_7 /sinkhorn.py
pranamanam's picture
Upload 87 files
0ba9d09 verified
Raw History Blame Contribute Delete
11.4 kB
"""Entropic transport and Sinkhorn's algorithm on the chapter's two-by-two example.
We subtract epsilon times the entropy from the transport cost. The objective becomes strictly
convex. Its minimizer is the Gibbs kernel exp(-C/epsilon) with one scaling per row and one per
column, and we fix those scalings from the two marginal constraints by rescaling alternately, one
relative entropy projection per half-step.
Chapter 7, Sections 7.9 and 7.10 of the notes. Runs on a = (0.6, 0.4), b = (0.3, 0.7) and the cost
matrix with entries 1, 9, 1, 1. Prints the kernel, the first sweep by hand, the marginal error after
each sweep, the converged plan and its multipliers, and the plan at four values of epsilon.
python lecture_7/sinkhorn.py # 40 sweeps, every worked number of Section 7.10
python lecture_7/sinkhorn.py --steps 1 # the first sweep alone, computed by hand
"""
# %% 01. Imports and configuration
from __future__ import annotations
import torch
from cis6270.runner import common_parser, print_report, save_run, set_seed
DTYPE = torch.float64 # The chapter quotes six figures of the kernel and of the converged plan.
COST = torch.tensor([[1.0, 9.0], [1.0, 1.0]], dtype=DTYPE) # C of Section 7.3.1.
SOURCE_A = torch.tensor([0.6, 0.4], dtype=DTYPE) # Row sums.
TARGET_B = torch.tensor([0.3, 0.7], dtype=DTYPE) # Column sums.
EXACT_PLAN = torch.tensor([[0.3, 0.3], [0.0, 0.4]], dtype=DTYPE) # pi* of Section 7.3.2.
EPSILON = 2.0 # The strength the chapter runs by hand.
# %% 02. Entropy of a plan, and the entropic objective
#
# Entropy is larger when the mass spreads over many routes and smaller when it concentrates, with
# the convention 0 log 0 = 0. We subtract epsilon H(pi) from the cost. This makes the objective
# strictly convex, so the minimizer is unique and every entry of it is positive. Section 7.9.
def entropy(pi: torch.Tensor) -> float:
"""H(pi) = -sum_ij pi_ij log pi_ij, in nats, with an empty route contributing nothing."""
positive = pi[pi > 0] # Drop the zero entries, and each contributes 0 log 0 = 0.
return float(-(positive * positive.log()).sum())
def entropic_objective(pi: torch.Tensor, C: torch.Tensor, epsilon: float) -> float:
"""<C, pi> - epsilon H(pi), the quantity the entropic problem minimizes over Pi(a, b)."""
return float((C * pi).sum()) - epsilon * entropy(pi)
# %% 03. Relative entropy from a positive matrix
#
# The definition still holds when xi is any strictly positive matrix. Take xi to be the independent
# coupling and the logarithm splits into three pieces, giving -H(pi) + H(a) + H(b), so the entropy
# term and a relative entropy from independence differ by a constant on Pi(a, b).
# Section 7.9.1.
def relative_entropy(pi: torch.Tensor, xi: torch.Tensor) -> float:
"""KL(pi || xi) = sum_ij pi_ij log(pi_ij / xi_ij) for a strictly positive xi."""
mask = pi > 0 # Routes with no mass contribute nothing.
return float((pi[mask] * (pi[mask] / xi[mask]).log()).sum())
def independent_coupling(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""a_i b_j, the coupling of largest entropy and the one Chapters 2 through 6 used."""
return a[:, None] * b[None, :] # [n, m].
# %% 04. The Lagrangian and the scaling form
#
# We set the derivative of the Lagrangian in a single entry to zero and get
# C_ij + epsilon (log pi_ij + 1) - f_i - g_j = 0. Solving that for the logarithm, we split the
# minimizer into a row factor, a kernel entry and a column factor. Section 7.9.2.
def gibbs_kernel(C: torch.Tensor, epsilon: float) -> torch.Tensor:
"""K = exp(-C / epsilon), taken entrywise, with larger entries on the cheaper routes."""
return torch.exp(-C / epsilon) # [n, m], every entry strictly positive.
def scaling_form(u: torch.Tensor, K: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
"""diag(u) K diag(v), the plan with row i scaled by u_i and column j scaled by v_j."""
return u[:, None] * K * v[None, :] # [n, m].
def multipliers_from_scalings(u: torch.Tensor, epsilon: float) -> torch.Tensor:
"""Invert u_i = exp(f_i / epsilon - 1/2) for f_i = epsilon (log u_i + 1/2)."""
return epsilon * (u.log() + 0.5) # [n] the additive price behind a multiplicative scaling.
# %% 05. The two updates
#
# The row sums of diag(u) K diag(v) are u * (K v) and the column sums are v * (K^T u). We solve
# each equation for its own scaling with the other held fixed, and the two updates follow. One sweep
# is the pair of them, at two matrix-vector products. Section 7.10.
def sinkhorn_sweeps(
K: torch.Tensor, a: torch.Tensor, b: torch.Tensor, sweeps: int,
) -> tuple[torch.Tensor, torch.Tensor, list[float]]:
"""Alternate u <- a / (K v) and v <- b / (K^T u), returning the scalings and the sweep errors."""
u = torch.ones_like(a)
v = torch.ones_like(b)
errors: list[float] = []
for _ in range(sweeps):
u = a / (K @ v) # Row sums become exactly a.
v = b / (K.T @ u) # Column sums become exactly b, and the row sums move away again.
errors.append(marginal_error(scaling_form(u, K, v), a, b))
return u, v, errors
def marginal_error(pi: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> float:
"""Total absolute error of the row sums against a plus that of the column sums against b."""
return float((pi.sum(dim=1) - a).abs().sum() + (pi.sum(dim=0) - b).abs().sum())
def sweeps_to_tolerance(K: torch.Tensor, a: torch.Tensor, b: torch.Tensor,
tolerance: float = 1e-9, limit: int = 2000) -> int:
"""The number of sweeps that brings the total marginal error below the tolerance."""
_, _, errors = sinkhorn_sweeps(K, a, b, limit)
for index, error in enumerate(errors):
if error < tolerance:
return index + 1
return limit
# %% 06. Each update as a relative entropy projection
#
# We minimize KL(pi || xi) over the matrices with row sums a, with no condition on the columns, and
# get pi_ij = xi_ij exp(lambda_i - 1), or xi with each row rescaled. The row sums fix the
# factor at a_i / sum_j xi_ij, and the u update performs exactly that rescaling. Section 7.10.1.
def project_rows(xi: torch.Tensor, a: torch.Tensor) -> torch.Tensor:
"""The matrix with row sums a closest to xi in relative entropy."""
return xi * (a / xi.sum(dim=1))[:, None] # One factor per row.
def project_columns(xi: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""The matrix with column sums b closest to xi in relative entropy."""
return xi * (b / xi.sum(dim=0))[None, :] # One factor per column.
# %% 07. The report
def main() -> None:
parser = common_parser(__doc__.split("\n")[0])
parser.add_argument("--epsilon", type=float, default=EPSILON, help="regularization strength")
parser.set_defaults(steps=40) # --steps is the number of Sinkhorn sweeps.
args = parser.parse_args()
set_seed(args.seed)
loud = not args.quiet
epsilon = args.epsilon
K = gibbs_kernel(COST, epsilon)
independent = independent_coupling(SOURCE_A, TARGET_B)
if loud:
print(f"H(pi*) = {entropy(EXACT_PLAN):.3f} H(a x b) = {entropy(independent):.3f} "
f"H(a) + H(b) = {entropy(SOURCE_A) + entropy(TARGET_B):.3f}")
print(f"KL(pi* || a x b) = {relative_entropy(EXACT_PLAN, independent):.3f} nats")
print(f"\nK = exp(-C/{epsilon:g}) = {[[round(float(k), 6) for k in row] for row in K]}")
print(f" kernel ratio between a cheap route and the expensive one "
f"{float(K[0, 0] / K[0, 1]):.1f}")
# The first sweep, written out the way Section 7.10.2 writes it.
v = torch.ones_like(TARGET_B)
Kv = K @ v # Row sums of K, since v = (1, 1).
u = SOURCE_A / Kv
Ktu = K.T @ u
v = TARGET_B / Ktu
first = scaling_form(u, K, v)
if loud:
print(f"\nsweep 1 K v = {[round(float(z), 6) for z in Kv]}"
f" u = {[round(float(z), 6) for z in u]}")
print(f" K^T u = {[round(float(z), 6) for z in Ktu]}"
f" v = {[round(float(z), 5) for z in v]}")
print(f" diag(u) K diag(v) = {[[round(float(z), 4) for z in row] for row in first]}")
print(f" column sums {[round(float(z), 4) for z in first.sum(0)]} "
f"row sums {[round(float(z), 4) for z in first.sum(1)]}")
after_u = project_rows(first, SOURCE_A) # The next update of u, as a row projection.
print(f" after the next u update: rows {[round(float(z), 3) for z in after_u.sum(1)]} "
f"columns {[round(float(z), 3) for z in after_u.sum(0)]}")
u, v, errors = sinkhorn_sweeps(K, SOURCE_A, TARGET_B, max(args.steps, 1))
plan = scaling_form(u, K, v)
f = multipliers_from_scalings(u, epsilon)
g = multipliers_from_scalings(v, epsilon)
sums = f[:, None] + g[None, :] # The entropic counterparts of the dual prices.
if loud:
print(f"\nmarginal error by sweep {[round(e, 3) for e in errors[:5]]}")
print(f" successive ratios "
f"{[round(errors[i + 1] / errors[i], 2) for i in range(min(4, len(errors) - 1))]}")
print(f" pi^eps = {[[round(float(z), 6) for z in row] for row in plan]}")
print(f" u = {[round(float(z), 4) for z in u]} v = {[round(float(z), 4) for z in v]}")
print(f" f = {[round(float(z), 3) for z in f]} g = {[round(float(z), 3) for z in g]}")
print(f" f_i + g_j = {[round(float(z), 3) for z in sums.flatten()]} "
f"against C = {[round(float(z), 0) for z in COST.flatten()]}")
print(f" the route (2,2) carries {float(plan[1, 1]):.3f} and overshoots its cost by "
f"{float(sums[1, 1] - COST[1, 1]):.3f}")
print("\neps <C, pi^eps> pi^eps_21 sweeps to 1e-9")
for value in (10.0, 2.0, 0.5, 0.1):
kernel = gibbs_kernel(COST, value)
su, sv, _ = sinkhorn_sweeps(kernel, SOURCE_A, TARGET_B, 400)
p = scaling_form(su, kernel, sv)
print(f"{value:<7g} {float((COST * p).sum()):<14.7f} {float(p[1, 0]):<16.3e} "
f"{sweeps_to_tolerance(kernel, SOURCE_A, TARGET_B)}")
report = {
"H(pi*)": entropy(EXACT_PLAN),
"H(a x b) = H(a) + H(b)": entropy(independent),
"KL(pi* || a x b)": relative_entropy(EXACT_PLAN, independent),
"K_11": float(K[0, 0]),
"K_12": float(K[0, 1]),
"sweep 1 marginal error": errors[0],
"sweep 5 marginal error": errors[min(4, len(errors) - 1)],
"sweeps to 1e-9": sweeps_to_tolerance(K, SOURCE_A, TARGET_B),
"pi^eps_11": float(plan[0, 0]),
"pi^eps_21": float(plan[1, 0]),
"<C, pi^eps>": float((COST * plan).sum()),
"<C, pi*>": float((COST * EXACT_PLAN).sum()),
"H(pi^eps)": entropy(plan),
"objective at pi^eps": entropic_objective(plan, COST, epsilon),
"objective at pi*": entropic_objective(EXACT_PLAN, COST, epsilon),
"objective at a x b": entropic_objective(independent, COST, epsilon),
"KL(pi^eps || a x b)": relative_entropy(plan, independent),
"f": [round(float(z), 3) for z in f],
"g": [round(float(z), 3) for z in g],
}
print_report("Entropic transport and Sinkhorn, Chapter 7 Sections 7.9 and 7.10", report)
save_run(args.out, config=vars(args), report=report) # Written only with --out.
if __name__ == "__main__":
main()