Blame for lecture_7/ot_flow_matching.py

"""Optimal transport inside flow matching: the same loss, trained on a different coupling.

The flow matching loss of Chapter 2 is defined for any joint law whose two marginals are p_0 and
p_1, and the independent draw of a source sample and a data sample is one such law. We replace it
with the optimal coupling of the batch. That changes the training pairs and leaves the loss, the
network and the sampler as they are. Under the optimal coupling no two paths cross, so the marginal
velocity at a point is the velocity of the single path through it and the trajectories come out
straight.

Chapter 7, Sections 7.8 and 7.8.1 of the notes. Trains two velocity networks on the four-Gaussian
target of `cis6270.data`, one on the independent coupling and one on the exact coupling of each
minibatch. Prints the coupling cost, the straightness of the learned paths and the one-step error.

    python lecture_7/ot_flow_matching.py              # 300 steps per model, under a minute
    python lecture_7/ot_flow_matching.py --steps 20   # a few seconds, for a first look
"""

# %% 01. Imports and configuration
from __future__ import annotations

import torch

from cis6270.data import gaussian_mixture_2d
from cis6270.metrics import mode_report
from cis6270.nets import VelocityMLP, count_parameters
from cis6270.runner import (
    common_parser, pick_device, print_report, save_run, set_seed, train,
)


# %% 02. The coupling of a batch
#
# A batch holds n source samples and m data samples with uniform masses 1/n and 1/m, and the
# transport problem on that batch has nm unknowns against n + m equality constraints. With n = m the
# vertex solution is a permutation, returned by the Hungarian algorithm in order n^3 operations.
# Section 7.8.1.


def batch_cost_matrix(x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor:
    """C_ij = ||x0_i - x1_j||^2 between every source sample and every data sample. [n, m]."""
    return torch.cdist(x0, x1) ** 2  # [n, m] squared Euclidean cost.


def optimal_pairing(x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor:
    """The permutation of the data samples that minimizes the total squared cost of the batch."""
    from scipy.optimize import linear_sum_assignment  # The assignment problem of Section 7.2.

    C = batch_cost_matrix(x0, x1).detach().cpu().numpy()  # [n, n].
    _, columns = linear_sum_assignment(C)  # The cheapest permutation, in order n^3 operations.
    return torch.as_tensor(columns, device=x0.device, dtype=torch.long)  # [n].


def couple(x0: torch.Tensor, x1: torch.Tensor, kind: str) -> torch.Tensor:
    """Reorder the data samples so that row i of the batch is the pair the coupling assigns to x0_i."""
    if kind == "independent":
        return x1  # Two independent draws, the coupling of Chapters 2 through 6.
    return x1[optimal_pairing(x0, x1)]  # The optimal coupling of this batch.


def vertex_count(n: int, m: int) -> dict[str, int]:
    """The size of the batch problem: nm unknowns, n + m constraints, at most n + m - 1 loaded."""
    return {"unknowns": n * m, "constraints": n + m, "loaded_routes": n + m - 1}


# %% 03. The loss
#
# The network still predicts a velocity, the target is still the straight-line displacement, and the
# time is still uniform on [0, 1]. Section 7.8.


def flow_matching_loss(model: VelocityMLP, x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor:
    """E_t ||v_theta((1-t) x_0 + t x_1) - (x_1 - x_0)||^2 over the pairs the coupling supplies."""
    t = torch.rand(x0.shape[0], device=x0.device)  # [B] one time per pair.
    xt = (1.0 - t[:, None]) * x0 + t[:, None] * x1  # [B, 2] the point on the straight path.
    ut = x1 - x0  # [B, 2] the conditional velocity, constant along the path.
    return ((model(xt, t) - ut) ** 2).sum(dim=1).mean()  # Sum coordinates, average pairs.


def conditional_velocity_energy(x0: torch.Tensor, x1: torch.Tensor) -> float:
    """The average squared conditional velocity, the kinetic action by Proposition 7.5."""
    return float(((x1 - x0) ** 2).sum(dim=1).mean())


# %% 04. Sampling, and two measurements of the learned field
#
# One Euler step integrates a straight trajectory exactly, so the gap between a one-step sample and
# a many-step sample measures how far the learned paths are from straight. Second, we compare the
# velocity along a trajectory with the net displacement of that trajectory. Section 7.8.


@torch.no_grad()
def integrate(model: VelocityMLP, x0: torch.Tensor, steps: int) -> tuple[torch.Tensor, float]:
    """Euler integration from t = 0 to t = 1, returning the samples and the path straightness.

    The straightness is the average of ||v_t(x_t) - (x_1 - x_0)||^2 along the path, zero exactly
    when the velocity never changes along a trajectory.
    """
    x = x0.clone()  # [B, 2].
    dt = 1.0 / steps
    velocities = []
    for step in range(steps):
        t = torch.full((x.shape[0],), step * dt, device=x.device)  # [B] shared time.
        v = model(x, t)  # [B, 2] the learned velocity at the current state.
        velocities.append(v)
        x = x + dt * v  # One Euler step.
    net = (x - x0)[None, :, :]  # [1, B, 2] the net displacement of each trajectory.
    deviation = ((torch.stack(velocities) - net) ** 2).sum(dim=2).mean()  # Scalar.
    return x, float(deviation)


@torch.no_grad()
def one_step_error(model: VelocityMLP, x0: torch.Tensor, steps: int) -> float:
    """The distance between the one-step sample x_0 + v_0(x_0) and the many-step sample."""
    fine, _ = integrate(model, x0, steps)  # [B, 2] the reference trajectory endpoint.
    t0 = torch.zeros(x0.shape[0], device=x0.device)
    coarse = x0 + model(x0, t0)  # [B, 2] a single Euler step of size 1.
    return float((coarse - fine).norm(dim=1).mean())


# %% 05. Training one model
def fit(kind: str, args, device: torch.device, data: torch.Tensor) -> dict:
    """Train a velocity network on the coupling named by ``kind`` and measure what it learned."""
    set_seed(args.seed)  # The same initialization and the same batches for both couplings.
    model = VelocityMLP(dim=2, width=args.width).to(device)
    costs: list[float] = []

    def loss_fn(step: int) -> torch.Tensor:
        x1 = data[torch.randint(0, len(data), (args.batch_size,), device=device)]  # [B, 2].
        x0 = torch.randn_like(x1)  # [B, 2] the source sample.
        paired = couple(x0, x1, kind)  # [B, 2] reordered by the coupling.
        costs.append(conditional_velocity_energy(x0, paired))
        return flow_matching_loss(model, x0, paired)

    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=True)
    torch.manual_seed(args.seed + 1)  # The same starting noise for both samplers.
    x0 = torch.randn(args.samples, 2, device=device)
    samples, straightness = integrate(model, x0, args.sample_steps)
    report = mode_report(samples.cpu())
    return {
        "kind": kind,
        "loss": sum(losses[-50:]) / len(losses[-50:]),
        "coupling_cost": sum(costs) / len(costs),
        "straightness": straightness,
        "one_step_error": one_step_error(model, x0, args.sample_steps),
        "modes_found": report["modes_found"],
        "nearest_distance": report["nearest_distance"],
    }


# %% 06. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=300, batch_size=128, width=64, samples=512, sample_steps=50)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)
    loud = not args.quiet
    data = gaussian_mixture_2d(num_points=4096, seed=args.seed).to(device)  # [4096, 2].

    # The two-point batch of Section 7.8.1, where both couplings can be written out.
    x = torch.tensor([[0.0], [2.0]])  # Two source samples on the line.
    y = torch.tensor([[1.0], [3.0]])  # Two data samples.
    C = batch_cost_matrix(x, y)  # [[1, 9], [1, 1]].
    independent_energy = float(C.mean())  # Weight 1/4 on each of the four pairs.
    optimal_energy = float(C[torch.arange(2), optimal_pairing(x, y)].mean())  # Weight 1/2 each.
    sizes = vertex_count(256, 256)  # The counting the chapter does, at n = m = 256.
    if loud:
        print("Two source samples at 0 and 2, two data samples at 1 and 3")
        print(f"  independent coupling, average squared velocity {independent_energy:.0f}")
        print(f"  optimal coupling, average squared velocity {optimal_energy:.0f}")
        print(f"\nA batch of 256 and 256: {sizes['unknowns']:,} unknowns, "
              f"{sizes['constraints']} constraints, at most {sizes['loaded_routes']} loaded routes")
        print(f"  an exact solve costs about {256 ** 3:.3g} operations, repeated every batch")
        print(f"  training batch {args.batch_size}, device={device}, parameters="
              f"{count_parameters(VelocityMLP(dim=2, width=args.width)):,}")

    results = {}
    for kind in ("independent", "optimal"):
        if loud:
            print(f"\ntraining on the {kind} coupling")
        results[kind] = fit(kind, args, device, data)
        if loud:
            r = results[kind]
            print(f"  loss {r['loss']:.4f}   batch coupling cost {r['coupling_cost']:.3f}")
            print(f"  straightness {r['straightness']:.4f}   "
                  f"one-step error {r['one_step_error']:.4f}")

    ratio = results["independent"]["coupling_cost"] / max(results["optimal"]["coupling_cost"], 1e-12)
    report = {
        "two-point batch, independent": independent_energy,
        "two-point batch, optimal": optimal_energy,
        "batch plan unknowns": sizes["unknowns"],
        "batch plan constraints": sizes["constraints"],
        "loaded routes at a vertex": sizes["loaded_routes"],
        "coupling cost, independent": results["independent"]["coupling_cost"],
        "coupling cost, optimal": results["optimal"]["coupling_cost"],
        "cost ratio": ratio,
        "training loss, independent": results["independent"]["loss"],
        "training loss, optimal": results["optimal"]["loss"],
        "straightness, independent": results["independent"]["straightness"],
        "straightness, optimal": results["optimal"]["straightness"],
        "one-step error, independent": results["independent"]["one_step_error"],
        "one-step error, optimal": results["optimal"]["one_step_error"],
        "modes found, independent": results["independent"]["modes_found"],
        "modes found, optimal": results["optimal"]["modes_found"],
        "nearest mode, independent": results["independent"]["nearest_distance"],
        "nearest mode, optimal": results["optimal"]["nearest_distance"],
    }
    print_report("Optimal transport inside flow matching, Chapter 7 Section 7.8", report)
    save_run(args.out, config=vars(args), report=report)  # Written only with --out.


if __name__ == "__main__":
    main()