Blame for lecture_6/consistency.py

"""Consistency models and consistency trajectory models, one time fixed and then both free.

A consistency model fixes the end time at the data end and learns the top edge of the triangle of
start and end times, with the identity axiom imposed at the corner by the boundary coefficients
c_skip and c_out. A consistency trajectory model puts the second time back and learns the whole
triangle, writing the map as a weighted average of the input and a head that predicts a
destination. Here we distill both from a teacher velocity available in closed form.

Chapter 6, Section 6.20 of the notes. Trained on the four Gaussian blobs of
``cis6270.data.gaussian_mixture_2d``. Prints the boundary parameterization on dx/dt = x, where the
model returns e^{0.5}, e^{0.1} and 1 at three start times, the teacher's error floor of
8.189e-05, and the destination head of a consistency trajectory model at two pairs of times.

    python lecture_6/consistency.py                 # two models, 1500 steps each
    python lecture_6/consistency.py --steps 20      # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math

import torch

from cis6270.data import gaussian_mixture_2d
from cis6270.metrics import mean_log_density, mode_report
from cis6270.nets import TwoTimeMLP, count_parameters
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train


torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
# than it saves. ``tests/conftest.py`` sets the same thing for the test session, and this line
# acts when a reader runs the script directly.

TEACHER_STEP = 0.02  # h, the length of the teacher's solver step inside the distillation loss.


# %% 02. The teacher velocity, in closed form on the empirical data
#
# The posterior over the data point given the interpolant is a softmax over the training set, so
# the exact marginal velocity of the four blobs needs no training. A trained flow matching model
# would serve in its place, with its own error on top of the solver's.
def teacher_denoiser(x, t, data):
    """D_t(x) = E[X_1 | X_t = x], the posterior mean over the training points. [B, 2]."""
    spread = (1.0 - t).clamp(min=1e-4)[:, None]  # [B, 1] the width of the conditional.
    squared = ((x[:, None, :] - t[:, None, None] * data[None]) ** 2).sum(-1)  # [B, N].
    weights = torch.softmax(-squared / (2.0 * spread ** 2), dim=1)  # [B, N] posterior over data.
    return weights @ data  # [B, 2] the posterior mean.


def teacher_velocity(x, t, data):
    """b(x,t) = (D_t(x) - x)/(1 - t) for the linear interpolant. [B, 2]."""
    return (teacher_denoiser(x, t, data) - x) / (1.0 - t).clamp(min=1e-4)[:, None]


def solver_step(x, t, h, data):
    """One Euler step of the teacher, written Phi_h in the chapter."""
    return x + h * teacher_velocity(x, t, data)  # [B, 2] the state at time t + h.


# %% 03. The boundary parameterization, checked on the scalar flow
#
# f_theta(x, s) = c_skip(s) x + c_out(s) F_theta(x, s) with c_skip = 1 and c_out = 1 - s. At the
# data end the output coefficient is zero and the model returns its input, so the loss spends no
# training signal on a condition we already know.
def scalar_consistency(x_value: float = 1.0):
    """The velocity head, the output coefficient and the model at three start times."""
    out = {}
    for s in (0.5, 0.9, 1.0):
        remaining = 1.0 - s  # c_out(s), zero at the data end.
        head = x_value * math.expm1(remaining) / remaining if remaining > 0 else x_value
        out[f"c_out({s})"] = remaining
        out[f"v_{s},1(1)"] = head
        out[f"f({s})"] = x_value + remaining * head  # Equals e^{1-s} on this flow.
    return out


def teacher_floor(x_value: float = 1.0, s: float = 0.5, h: float = 0.01):
    """The distillation residual at the exact map, the teacher's own Euler error."""
    exact = lambda x, t: x * math.exp(1.0 - t)  # f(x, t) = x e^{1-t} on dx/dt = x.
    moved = x_value + h * x_value  # One teacher Euler step, and it undershoots e^{h}.
    residual = exact(x_value, s) - exact(moved, s + h)
    local = math.exp(h) - 1.0 - h  # The local error of that step.
    return {"f(1, 0.5)": exact(x_value, s), "f(1.01, 0.51)": exact(moved, s + h),
            "distillation residual": residual, "local Euler error": local,
            "carried to the data end": local * math.exp(1.0 - s - h)}


# %% 04. Consistency trajectory models, with both times free
#
# X_theta(x, s, t) = a_{s,t} x + (1 - a_{s,t}) g_theta(x, s, t) with a_{s,t} = (1-t)/(1-s), equal
# to one on the diagonal, so the map returns its input there for any head whatever.
def ctm_weight(s, t):
    """a_{s,t} = (1-t)/(1-s), the share of the input kept by the step."""
    return (1.0 - t) / (1.0 - s).clamp(min=1e-4)


def ctm_map(model, x, s, t):
    """The weighted average of the input and the destination head."""
    a = ctm_weight(s, t)[:, None]  # [B, 1].
    return a * x + (1.0 - a) * model(x, s, t)  # [B, 2].


def scalar_ctm(x_value: float = 1.0):
    """The head solved from the exact map at two pairs of times, and the two-route sampler."""
    out = {}
    for s, t in ((0.0, 0.5), (0.0, 1.0)):
        a = (1.0 - t) / (1.0 - s)
        reached = x_value * math.exp(t - s)  # The exact map of dx/dt = x.
        out[f"a_{s},{t}"] = a
        out[f"g_{s},{t}"] = (reached - a * x_value) / (1.0 - a)  # Solve the weighted average.
    short = x_value * math.exp(0.5)  # X_{0,0.5}(1), the direct short step.
    long_then_back = (x_value * math.exp(1.0)) * math.exp(-0.5)  # X_{1,0.5}(X_{0,1}(1)).
    out["short step"], out["long jump and return"] = short, long_then_back
    out["two routes apart"] = abs(short - long_then_back)
    return out


# %% 05. The two losses
#
# Both compare the model at one time against the model at a state the teacher has advanced by h,
# with the target detached. That detachment is the frozen copy of the weights the chapter writes as
# theta minus. The consistency model holds the end time at one; the trajectory model draws it.
def make_consistency_loss(model, data, batch_size, device):
    ones = torch.ones(batch_size, device=device)

    def loss_fn(step: int) -> torch.Tensor:
        index = torch.randint(0, len(data), (batch_size,), device=device)
        x1, x0 = data[index], torch.randn(batch_size, 2, device=device)
        s = (1.0 - TEACHER_STEP) * torch.rand(batch_size, device=device)  # [B] with s + h <= 1.
        xs = (1 - s[:, None]) * x0 + s[:, None] * x1  # [B, 2] the interpolant at the start time.
        with torch.no_grad():  # The second leg is the teacher's, and no gradient reaches it.
            advanced = solver_step(xs, s, TEACHER_STEP, data)
            target = model.flow_map(advanced, s + TEACHER_STEP, ones)  # f_{theta-}(Phi_h(x_s), s+h).
        return ((model.flow_map(xs, s, ones) - target) ** 2).mean()

    return loss_fn


def make_trajectory_loss(model, data, batch_size, device):

    def loss_fn(step: int) -> torch.Tensor:
        index = torch.randint(0, len(data), (batch_size,), device=device)
        x1, x0 = data[index], torch.randn(batch_size, 2, device=device)
        s = (1.0 - TEACHER_STEP) * torch.rand(batch_size, device=device)  # [B] start time.
        t = s + TEACHER_STEP + (1.0 - s - TEACHER_STEP) * torch.rand(batch_size, device=device)
        xs = (1 - s[:, None]) * x0 + s[:, None] * x1  # [B, 2].
        with torch.no_grad():
            advanced = solver_step(xs, s, TEACHER_STEP, data)
            target = ctm_map(model, advanced, s + TEACHER_STEP, t)  # The same end time, one h later.
        return ((ctm_map(model, xs, s, t) - target) ** 2).mean()

    return loss_fn


# %% 06. Sampling
@torch.no_grad()
def consistency_sample(model, x0):
    """One call of the edge map, from the source straight to the data end."""
    zeros = torch.zeros(len(x0), device=x0.device)
    return model.flow_map(x0, zeros, torch.ones_like(zeros))  # [N, 2].


@torch.no_grad()
def trajectory_sample(model, x0, steps: int):
    """Cross the triangle in equal pieces, one evaluation of the weighted average per piece."""
    x = x0
    grid = torch.linspace(0.0, 1.0, steps + 1, device=x0.device)  # [steps+1] the time grid.
    for s, t in zip(grid[:-1], grid[1:]):
        x = ctm_map(model, x, s.expand(len(x)), t.expand(len(x)))
    return x  # [N, 2].


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=1500, batch_size=128, samples=512)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)
    data = gaussian_mixture_2d(num_points=1024, seed=args.seed).to(device)  # [1024, 2].

    report = scalar_consistency()
    report.update(teacher_floor())
    report.update(scalar_ctm())
    if not args.quiet:
        print(f"boundary parameterization: f(0.5) = {report['f(0.5)']:.6f}, "
              f"f(0.9) = {report['f(0.9)']:.6f}, f(1.0) = {report['f(1.0)']:.6f}")
        print(f"teacher floor {report['distillation residual']:.6e} against the carried local "
              f"error {report['carried to the data end']:.6e}")

    samples: list[str] = []
    set_seed(args.seed)
    consistency = TwoTimeMLP(dim=2, width=args.width).to(device)
    if not args.quiet:
        print(f"\ntraining the consistency model  parameters={count_parameters(consistency):,}")
    edge_losses = train(consistency,
                        make_consistency_loss(consistency, data, args.batch_size, device),
                        steps=args.steps, lr=args.lr, quiet=args.quiet)
    report["consistency final loss"] = sum(edge_losses[-20:]) / len(edge_losses[-20:])
    torch.manual_seed(args.seed + 1)
    x0 = torch.randn(args.samples, 2, device=device)  # [N, 2] the shared source points.
    points = consistency_sample(consistency, x0)
    stats = mode_report(points)
    report["consistency one-step nearest"] = stats["nearest_distance"]
    report["consistency one-step log density"] = mean_log_density(points.cpu())
    samples.append(f"consistency, 1 step: nearest {stats['nearest_distance']:.3f}  "
                   f"modes {stats['modes_found']}  entropy {stats['mode_entropy']:.3f}")

    set_seed(args.seed)
    trajectory = TwoTimeMLP(dim=2, width=args.width).to(device)
    if not args.quiet:
        print(f"\ntraining the trajectory model  parameters={count_parameters(trajectory):,}")
    losses = train(trajectory, make_trajectory_loss(trajectory, data, args.batch_size, device),
                   steps=args.steps, lr=args.lr, quiet=args.quiet)
    for steps in (1, 2, 4):
        points = trajectory_sample(trajectory, x0, steps)
        stats = mode_report(points)
        report[f"trajectory {steps}-step nearest"] = stats["nearest_distance"]
        samples.append(f"trajectory, {steps} step(s): nearest {stats['nearest_distance']:.3f}  "
                       f"modes {stats['modes_found']}  entropy {stats['mode_entropy']:.3f}")
    report["trajectory one step against four"] = float(
        (trajectory_sample(trajectory, x0, 1) - trajectory_sample(trajectory, x0, 4)).norm(dim=1)
        .mean())
    report["trajectory final loss"] = sum(losses[-20:]) / len(losses[-20:])

    print_report("Consistency and consistency trajectory models, Chapter 6", report, samples)
    save_run(args.out, model=trajectory, config=vars(args), report=report, losses=losses,
             samples=samples)


if __name__ == "__main__":
    main()