Blame for lecture_6/shortcut.py

"""Shortcut models, the composition identity indexed by step size.

The second input of the network is the length of the step and not the time at which it ends, and we
enforce the identity only between a step and the step twice as long. One doubling constraint
suffices because the ladder chains. With a finest step of 1/128 and seven doublings, the one-step
map at level seven rests on seven relations and one regression against the data.

Chapter 6, Section 6.20 of the notes. Trained on the four Gaussian blobs of
``cis6270.data.gaussian_mixture_2d``. Prints the doubling worked on dx/dt = x, where two steps of
0.25 average to 1.297443 and the direct step of 0.5 returns the same number, then the dyadic ladder
and the trained model's doubling residual at every rung.

    python lecture_6/shortcut.py                 # 2000 steps, about twenty seconds
    python lecture_6/shortcut.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.

FINEST = 1.0 / 128.0  # d_0, the finest step of the ladder.
LEVELS = 7  # log2(128) doublings, so the step sizes are d_k = 2^k d_0 for k = 0 to 7.
ZERO_FRACTION = 0.75  # Three quarters of each batch sits at step size zero.


# %% 02. The network in step-size coordinates
#
# The substitution s_theta(x, t, d) = v_{t, t+d}(x) turns the published network into the two-time
# average velocity of this chapter, so the same weights serve both readings.
def shortcut(model, x, t, d):
    """s_theta(x, t, d), the average velocity over the step of length d starting at time t."""
    return model(x, t, t + d)  # [B, 2] = v_{t, t+d}(x).


def take_step(model, x, t, d):
    """One step of length d, the flow map X_{t, t+d} in these coordinates."""
    return x + d[:, None] * shortcut(model, x, t, d)  # [B, 2].


# %% 03. The doubling identity, worked on the scalar flow
#
# s(x, t, 2d) = half s(x, t, d) + half s(x', t+d, d) with x' = x + d s(x, t, d). On dx/dt = x the
# average velocity is x (e^d - 1)/d, and the identity holds at every step length at once.
def scalar_doubling(x_value: float = 1.0, d: float = 0.25):
    """Every term of the doubling identity on dx/dt = x, each one computed here."""
    first = x_value * math.expm1(d) / d  # s(1, 0, 0.25).
    reached = x_value + d * first  # x', or e^{0.25} to six figures.
    second = reached * math.expm1(d) / d  # s(x', 0.25, 0.25), faster because the state has grown.
    doubled = x_value * math.expm1(2.0 * d) / (2.0 * d)  # s(1, 0, 0.5), computed directly.
    return {
        "s(1, 0, d)": first,
        "state after the first step": reached,
        "s(x', d, d)": second,
        "average of the two": 0.5 * (first + second),
        "s(1, 0, 2d) direct": doubled,
        "doubling gap": abs(0.5 * (first + second) - doubled),
        "state after the doubled step": x_value + 2.0 * d * doubled,
        "exp(2d)": math.exp(2.0 * d),
    }


# %% 04. The dyadic ladder
#
# Eight step sizes from 1/128 up to 1, each training pair tying level k to level k+1. A state at
# level zero is already trained against the diagonal by the anchor, so the correctness of the
# one-step map rests on a chain of seven relations and one regression.
def ladder():
    """The step sizes d_k = 2^k d_0, from the finest to the whole interval."""
    return [FINEST * 2.0 ** k for k in range(LEVELS + 1)]  # [1/128, ..., 1.0].


# %% 05. The loss
#
# At step size zero the target is the flow matching target of Chapter 2, and above it the batch
# draws a level of the ladder, forms the two short steps and regresses the doubled step onto their
# average, with the target detached. The chapter puts three quarters of the batch on the zero step
# and the rest on the doubling. Here we evaluate both terms on the whole batch and weight them in
# that proportion. The expectation is the same at one network call more per step.
def make_loss(model, data, batch_size, device):
    sizes = torch.tensor(ladder()[:LEVELS], device=device)  # Levels whose doubling still fits in [0,1].

    def loss_fn(step: int) -> torch.Tensor:
        index = torch.randint(0, len(data), (batch_size,), device=device)
        x1 = data[index]  # [B, 2] the data draw.
        x0 = torch.randn_like(x1)  # [B, 2] the source draw.
        t = torch.rand(batch_size, device=device)  # [B] a time for the zero-step rows.
        xt = (1 - t[:, None]) * x0 + t[:, None] * x1  # [B, 2] the interpolant.
        zero = torch.zeros(batch_size, device=device)
        anchor = ((shortcut(model, xt, t, zero) - (x1 - x0)) ** 2).mean()  # Step size zero.

        d = sizes[torch.randint(0, LEVELS, (batch_size,), device=device)]  # [B] one rung each.
        start = (1.0 - 2.0 * d) * torch.rand(batch_size, device=device)  # [B] t with t + 2d <= 1.
        xs = (1 - start[:, None]) * x0 + start[:, None] * x1  # [B, 2] the interpolant at t.
        with torch.no_grad():  # The two legs are frozen; only the doubled step is corrected.
            first = shortcut(model, xs, start, d)
            reached = xs + d[:, None] * first
            second = shortcut(model, reached, start + d, d)
            target = 0.5 * (first + second)
        doubled = ((shortcut(model, xs, start, 2.0 * d) - target) ** 2).mean()
        return ZERO_FRACTION * anchor + (1.0 - ZERO_FRACTION) * doubled

    return loss_fn


# %% 06. Sampling, and the residual at every rung
@torch.no_grad()
def sample(model, x0, steps: int):
    """Cross [0, 1] in equal steps of length 1/steps, one network evaluation per step."""
    x, d = x0, torch.full((len(x0),), 1.0 / steps, device=x0.device)
    for k in range(steps):
        t = torch.full((len(x0),), k / steps, device=x0.device)
        x = take_step(model, x, t, d)
    return x  # [N, 2].


@torch.no_grad()
def rung_residuals(model, x0):
    """The squared doubling gap at each level of the ladder, measured on source points."""
    out = {}
    for k, d in enumerate(ladder()[:LEVELS]):
        step = torch.full((len(x0),), d, device=x0.device)
        start = torch.zeros(len(x0), device=x0.device)
        first = shortcut(model, x0, start, step)
        second = shortcut(model, x0 + step[:, None] * first, start + step, step)
        gap = shortcut(model, x0, start, 2.0 * step) - 0.5 * (first + second)
        out[f"rung {k}: d = 1/{int(round(1.0 / d))}"] = float((gap ** 2).mean())
    return out


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

    report = scalar_doubling()
    report["ladder levels"] = LEVELS
    report["finest step d_0"] = FINEST
    if not args.quiet:
        print(f"on dx/dt = x: two steps of 0.25 average to {report['average of the two']:.6f}, "
              f"one step of 0.5 returns {report['s(1, 0, 2d) direct']:.6f}")
        print(f"ladder {[round(d, 5) for d in ladder()]}")

    model = TwoTimeMLP(dim=2, width=args.width).to(device)
    if not args.quiet:
        print(f"\ntraining the shortcut model  parameters={count_parameters(model):,}")
    losses = train(model, make_loss(model, data, args.batch_size, device),
                   steps=args.steps, lr=args.lr, quiet=args.quiet)

    torch.manual_seed(args.seed + 1)
    x0 = torch.randn(args.samples, 2, device=device)  # [N, 2] the shared source points.
    report.update(rung_residuals(model, x0))
    samples: list[str] = []
    for steps in (1, 2, 4, 8):
        points = sample(model, x0, steps)
        stats = mode_report(points)
        report[f"{steps}-step nearest distance"] = stats["nearest_distance"]
        report[f"{steps}-step log density"] = mean_log_density(points.cpu())
        samples.append(f"{steps} step(s): nearest {stats['nearest_distance']:.3f}  "
                       f"modes {stats['modes_found']}  entropy {stats['mode_entropy']:.3f}")
    report["one step against eight"] = float(
        (sample(model, x0, 1) - sample(model, x0, 8)).norm(dim=1).mean())
    report["final loss"] = sum(losses[-20:]) / len(losses[-20:])

    print_report("Shortcut models, Chapter 6", report, samples)
    save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples)


if __name__ == "__main__":
    main()