Download lecture_7/csbm.py from ChatterjeeLab/CIS6270: direct link, hf CLI and curl.
- Browser
- Download file 13.4 kB
-
https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_7/csbm.py
- Command line
-
hf download hf://ChatterjeeLab/CIS6270/lecture_7/csbm.py
-
curl -L -o csbm.py https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_7/csbm.py
13.4 kB
| """Categorical Schrodinger Bridge Matching: D-IMF on a path law small enough to write out. | |
| With N intermediate times a path is N + 2 states and the whole path law is a finite table, so both | |
| projections have one-line expressions. The reciprocal projection holds the endpoint coupling and | |
| redraws the interior from the reference, and the Markov projection keeps the one-step conditionals | |
| and multiplies them into a chain. The bridge is the only law in both sets. On two states and three | |
| transitions the sixteen paths fit in a list, so the alternation runs exactly. | |
| Chapter 7, Section 7.25 of the notes. Runs on S = 2, four time points, the uniform reference with | |
| alpha = 0.35, the source law (0.5, 0.5) and the requested law (0.2, 0.8). Prints the uniform and | |
| the ordered reference row, the exact bridge coupling, the conditionals of the first Markov | |
| projection, the gap after each round, and how that gap moves with alpha and with the path length. | |
| python lecture_7/csbm.py # 5 rounds of D-IMF, the five the chapter reports | |
| python lecture_7/csbm.py --steps 1 # the single round the chapter writes out | |
| """ | |
| # %% 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 gap reaches 1e-11 in five rounds, well below float32. | |
| ALPHA = 0.35 # The probability of changing category in one step. | |
| TRANSITIONS = 3 # N + 1, so four time points and sixteen paths. | |
| PI_0 = torch.tensor([0.5, 0.5], dtype=DTYPE) # The source law. | |
| PI_1 = torch.tensor([0.2, 0.8], dtype=DTYPE) # The requested law at the last time. | |
| # %% 02. The reference over S categories | |
| # | |
| # The uniform reference holds the category with probability 1 - alpha and spreads the rest evenly | |
| # over the alternatives. When the categories have an order we can use a discretized Gaussian | |
| # centred on the current one, under which a neighbour is likelier than a distant category. | |
| # Section 7.25. | |
| def uniform_reference(alpha: float, size: int) -> torch.Tensor: | |
| """A_n(x, y) = 1 - alpha on the diagonal and alpha / (S - 1) off it. [S, S].""" | |
| off = alpha / (size - 1) # The share each alternative receives. | |
| return torch.full((size, size), off, dtype=DTYPE) + torch.eye(size, dtype=DTYPE) * ( | |
| 1.0 - alpha - off) | |
| def ordered_reference(width: float, size: int) -> torch.Tensor: | |
| """Rows proportional to exp(-(y - x)^2 / 2 width^2), under which a nearby category is likelier.""" | |
| offset = torch.arange(size, dtype=DTYPE)[None, :] - torch.arange(size, dtype=DTYPE)[:, None] | |
| weights = torch.exp(-offset ** 2 / (2.0 * width ** 2)) # [S, S] before normalizing. | |
| return weights / weights.sum(dim=1, keepdim=True) | |
| def sinkhorn(matrix: torch.Tensor, start: torch.Tensor, end: torch.Tensor, sweeps: int = 400 | |
| ) -> torch.Tensor: | |
| """Scale diag(p_0) A^N to the two marginals, and the result is the exact bridge coupling.""" | |
| f = torch.ones_like(start) | |
| g = torch.ones_like(end) | |
| for _ in range(sweeps): | |
| g = end / (matrix.T @ f) # Column sums become exactly p_1. | |
| f = start / (matrix @ g) # Row sums become exactly p_0. | |
| return f[:, None] * matrix * g[None, :] | |
| # %% 03. The two projections in discrete time | |
| # | |
| # A path law is a tensor with one axis per time. The reciprocal projection multiplies the endpoint | |
| # coupling by the reference bridge weight of every path, and the Markov projection reads off the | |
| # one-step conditionals and multiplies them back into a chain. Section 7.25.1. | |
| def reference_joint(matrix: torch.Tensor, steps: int) -> torch.Tensor: | |
| """The product of the transition probabilities along a path, as a tensor with one axis per time.""" | |
| joint = matrix.clone() # [x_0, x_1]. | |
| for _ in range(steps - 1): | |
| joint = joint.unsqueeze(-1) * matrix # Append one more transition on a new axis. | |
| return joint # [S] * (steps + 1). | |
| def bridge_table(matrix: torch.Tensor, steps: int) -> torch.Tensor: | |
| """R(x_interior | x_0, x_N), the reference bridge weight of every path given its two ends.""" | |
| whole = torch.linalg.matrix_power(matrix, steps) # [S, S] the reference over the interval. | |
| size = len(matrix) | |
| shape = [size] + [1] * (steps - 1) + [size] # Broadcast over the interior axes. | |
| return reference_joint(matrix, steps) / whole.reshape(shape) | |
| def reciprocal_projection(coupling: torch.Tensor, matrix: torch.Tensor, steps: int) -> torch.Tensor: | |
| """proj_R(P) = P(x_0, x_N) R(x_in | x_0, x_N), the endpoint coupling with reference bridges.""" | |
| size = len(matrix) | |
| shape = [size] + [1] * (steps - 1) + [size] | |
| return coupling.reshape(shape) * bridge_table(matrix, steps) | |
| def markov_projection(law: torch.Tensor) -> tuple[list[torch.Tensor], torch.Tensor]: | |
| """proj_M(P) = P(x_0) prod_n P(x_{t_n} | x_{t_{n-1}}), as the conditionals and the coupling.""" | |
| conditionals = [] | |
| for step in range(law.dim() - 1): | |
| axes = [d for d in range(law.dim()) if d not in (step, step + 1)] | |
| pair = law.sum(dim=axes) # [S, S] the joint of two consecutive states. | |
| conditionals.append(pair / pair.sum(dim=1, keepdim=True)) | |
| joint = torch.diag(law.sum(dim=list(range(1, law.dim())))) # [S, S] the law at time 0. | |
| for matrix in conditionals: | |
| joint = joint @ matrix | |
| return conditionals, joint | |
| def transition_loss(true: torch.Tensor, network: torch.Tensor) -> float: | |
| """KL(P(x_t | x_s) || m_theta(x_t | x_s)) on one transition out of one state, in nats.""" | |
| return float((true * (true / network).log()).sum()) | |
| # %% 04. D-IMF on sixteen paths | |
| # | |
| # We alternate the two projections, iterative Markovian fitting with discrete time. The | |
| # starting law is the independent coupling filled in with reference bridges, so both requested | |
| # marginals are already exact and only the joint law over paths moves. Section 7.25.2. | |
| def d_imf(coupling: torch.Tensor, matrix: torch.Tensor, steps: int, rounds: int, | |
| target: torch.Tensor) -> tuple[list[float], list[torch.Tensor]]: | |
| """Run the alternation, returning the gap after each round and the first round's conditionals.""" | |
| gaps: list[float] = [] | |
| first: list[torch.Tensor] = [] | |
| for index in range(rounds): | |
| law = reciprocal_projection(coupling, matrix, steps) # Reference bridges in the interior. | |
| conditionals, coupling = markov_projection(law) # The chain and its endpoint coupling. | |
| if index == 0: | |
| first = conditionals | |
| gaps.append(float((reciprocal_projection(coupling, matrix, steps) - target).abs().sum())) | |
| return gaps, first | |
| def starting_gap(matrix: torch.Tensor, steps: int, start: torch.Tensor, | |
| end: torch.Tensor) -> tuple[float, torch.Tensor, torch.Tensor]: | |
| """The gap of the independent coupling from the bridge, with both path laws.""" | |
| whole = torch.linalg.matrix_power(matrix, steps) | |
| plan = sinkhorn(torch.diag(start) @ whole, start, end) | |
| target = reciprocal_projection(plan, matrix, steps) | |
| initial = reciprocal_projection(start[:, None] * end[None, :], matrix, steps) | |
| return float((initial - target).abs().sum()), initial, target | |
| # %% 05. The report | |
| def main() -> None: | |
| parser = common_parser(__doc__.split("\n")[0]) | |
| parser.set_defaults(steps=5) # --steps is the number of D-IMF rounds. | |
| args = parser.parse_args() | |
| set_seed(args.seed) | |
| loud = not args.quiet | |
| rounds = max(args.steps, 1) | |
| A = uniform_reference(ALPHA, len(PI_0)) | |
| whole = torch.linalg.matrix_power(A, TRANSITIONS) | |
| ordered = ordered_reference(1.0, 5) # Five ordered categories, a width of one category. | |
| if loud: | |
| print(f"S = 2, {TRANSITIONS + 1} time points, alpha = {ALPHA}, " | |
| f"{len(PI_0) ** (TRANSITIONS + 1)} paths") | |
| print(f" A = {[[round(float(z), 2) for z in row] for row in A]}, " | |
| f"A^{TRANSITIONS} = {[[round(float(z), 4) for z in row] for row in whole]}") | |
| print(f" the ordered reference out of the middle of five categories: " | |
| f"{[round(float(z), 4) for z in ordered[2]]}") | |
| plan = sinkhorn(torch.diag(PI_0) @ whole, PI_0, PI_1) | |
| gap_0, initial, target = starting_gap(A, TRANSITIONS, PI_0, PI_1) | |
| interior = [initial.sum(dim=[d for d in range(TRANSITIONS + 1) if d != k]) for k in (1, 2)] | |
| independent = PI_0[:, None] * PI_1[None, :] # Where the alternation starts. | |
| if loud: | |
| print(f"\npi* = {[[round(float(z), 5) for z in row] for row in plan]}, against the " | |
| f"independent coupling " | |
| f"{[[round(float(z), 1) for z in row] for row in independent]}") | |
| print(f" the path 1,1,1,1 has reference bridge weight " | |
| f"{float(bridge_table(A, TRANSITIONS)[0, 0, 0, 0]):.4f}, so its probability is " | |
| f"{float(initial[0, 0, 0, 0]):.4f} at the start and {float(target[0, 0, 0, 0]):.4f} " | |
| f"under the bridge") | |
| print(f" the path 2,2,2,2 has {float(initial[1, 1, 1, 1]):.4f} and " | |
| f"{float(target[1, 1, 1, 1]):.4f}") | |
| print(f" the two laws agree at the interior times, " | |
| f"{[[round(float(z), 4) for z in m] for m in interior]}, and differ by " | |
| f"{gap_0:.4f} over the paths") | |
| gaps, conditionals = d_imf(independent, A, TRANSITIONS, rounds, target) | |
| # Eight rounds reach machine precision, past which a ratio of two gaps measures rounding, so | |
| # the report quotes the first one. | |
| ratios = [gaps[k] / gaps[k + 1] for k in range(len(gaps) - 1) if gaps[k + 1] > 1e-14] | |
| per_round = ratios[0] if ratios else gap_0 / gaps[0] # One round alone gives only the first. | |
| network = torch.tensor([0.25, 0.75], dtype=DTYPE) # A network that is wrong by a little. | |
| one_term = transition_loss(conditionals[-1][0], network) | |
| if loud: | |
| print(f"\nthe {len(conditionals)} conditionals of the first Markov projection") | |
| for matrix in conditionals: | |
| print(f" {[[round(float(z), 4) for z in row] for row in matrix]}") | |
| print(f" gap {gap_0:.4f} at the start, then {['%.1e' % x for x in gaps]}") | |
| print(f" ratios {[round(r, 1) for r in ratios]}") | |
| print(f" the true conditional {[round(float(z), 4) for z in conditionals[-1][0]]} against " | |
| f"a network at {[float(z) for z in network]} costs {one_term:.4f} nats") | |
| # A reference that mixes well and a longer chain both leave the bridge less to do. The starting | |
| # gap records that and the rate of convergence follows it. Two rounds separate the two ratios | |
| # from the float64 floor, and the sixty-four-path column reaches it after three. | |
| mixing = {} | |
| for alpha in (0.2, ALPHA, 0.4): | |
| matrix = uniform_reference(alpha, len(PI_0)) | |
| start, _, goal = starting_gap(matrix, TRANSITIONS, PI_0, PI_1) | |
| later, _ = d_imf(independent, matrix, TRANSITIONS, 2, goal) | |
| mixing[alpha] = (start, later[0] / later[1]) | |
| lengths = {} | |
| for steps in (2, TRANSITIONS, 5): | |
| start, _, goal = starting_gap(A, steps, PI_0, PI_1) | |
| later, _ = d_imf(independent, A, steps, 2, goal) | |
| lengths[steps] = (start, later[0] / later[1]) | |
| if loud: | |
| print("\nalpha starting gap gap divided per round") | |
| for alpha, (start, ratio) in mixing.items(): | |
| print(f" {alpha:<6g} {start:<14.5f} {ratio:.0f}") | |
| print("\npaths starting gap gap divided per round") | |
| for steps, (start, ratio) in lengths.items(): | |
| print(f" {len(PI_0) ** (steps + 1):<6d} {start:<14.5f} {ratio:.0f}") | |
| report = { | |
| "A^3": [round(float(z), 4) for z in whole.flatten()], | |
| "the ordered reference row": [round(float(z), 4) for z in ordered[2]], | |
| "pi*_11": float(plan[0, 0]), | |
| "pi*_12": float(plan[0, 1]), | |
| "reference bridge weight of 1,1,1,1": float(bridge_table(A, TRANSITIONS)[0, 0, 0, 0]), | |
| "P^0(1,1,1,1)": float(initial[0, 0, 0, 0]), | |
| "P*(1,1,1,1)": float(target[0, 0, 0, 0]), | |
| "P^0(2,2,2,2)": float(initial[1, 1, 1, 1]), | |
| "P*(2,2,2,2)": float(target[1, 1, 1, 1]), | |
| "marginal at the first interior time": [round(float(z), 4) for z in interior[0]], | |
| "marginal at the second": [round(float(z), 4) for z in interior[1]], | |
| "gap at the start": gap_0, | |
| "gap after one round": gaps[0], | |
| "gap after the last round": gaps[-1], | |
| "rounds run": rounds, | |
| "gap divided per round": per_round, | |
| "first conditional, out of state 1": [round(float(z), 4) for z in conditionals[0][0]], | |
| "last conditional, out of state 1": [round(float(z), 4) for z in conditionals[-1][0]], | |
| "one transition against a wrong network": one_term, | |
| "starting gap at alpha = 0.2": mixing[0.2][0], | |
| "rounds divide the gap by, at alpha = 0.2": mixing[0.2][1], | |
| "rounds divide the gap by, at alpha = 0.4": mixing[0.4][1], | |
| "starting gap on eight paths": lengths[2][0], | |
| "starting gap on sixty-four paths": lengths[5][0], | |
| "rounds divide the gap by, on eight paths": lengths[2][1], | |
| "rounds divide the gap by, on sixty-four paths": lengths[5][1], | |
| } | |
| print_report("Categorical Schrodinger Bridge Matching, Chapter 7 Section 7.25", report) | |
| save_run(args.out, config=vars(args), report=report) # Written only with --out. | |
| if __name__ == "__main__": | |
| main() | |