Download lecture_7/sinkhorn.py from ChatterjeeLab/CIS6270: direct link, hf CLI and curl.
- Browser
- Download file 11.4 kB
-
https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_7/sinkhorn.py
- Command line
-
hf download hf://ChatterjeeLab/CIS6270/lecture_7/sinkhorn.py
-
curl -L -o sinkhorn.py https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_7/sinkhorn.py
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() | |