Upload 87 files

0ba9d09 verified about 1 month ago

8.95 kB

"""MeanFlow, the Eulerian identity written in the average-velocity parameterization.
If we substitute X_{s,t}(x) = x + (t-s) v_{s,t}(x) into the Eulerian identity, we are left with one
equation, v_{s,t}(x) = b(x,s) + (t-s) d/ds v_{s,t}(x), with the correction term a total derivative
along the trajectory. One forward-mode pass along the tangent (b, 1, 0) in the arguments (x, s, t)
computes that derivative without forming a Jacobian, and the target uses the realized velocity
x_1 - x_0 of the interpolant in place of b, so a step costs one forward pass and one derivative.
Chapter 6, Sections 6.13, 6.15 and 6.18 of the notes. Trained on the four Gaussian blobs of
``cis6270.data.gaussian_mixture_2d``. Prints the identity worked by hand on dx/dt = x, where the
average velocity is 1.718x, the instantaneous velocity is x and the correction is 0.718x, and then
one-step, two-step and four-step samples from the trained map.
    python lecture_6/meanflow.py                 # 2000 steps, about twenty seconds
    python lecture_6/meanflow.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.

DIAGONAL_FRACTION = 0.25  # Share of each batch drawn with t = s, where the identity is flow matching.


# %% 02. The average velocity of the scalar flow, in closed form
#
# The exact map of dx/dt = x multiplies the state by e^{t-s}, so the displacement over the elapsed
# time is x (e^{t-s} - 1)/(t-s). On the diagonal the limit is x, or b(x,t).
def scalar_average_velocity(x, s, t):
    """v_{s,t}(x) = x (e^{t-s} - 1)/(t - s) on the flow dx/dt = x."""
    elapsed = t - s
    return x * torch.expm1(elapsed) / elapsed  # expm1 keeps the limit accurate near the diagonal.


# %% 03. The identity, checked on the scalar flow by hand
#
# v_{s,t}(x) = b(x,s) + (t-s) d/ds v_{s,t}(x), with d/ds = partial_s + b(x,s) . grad_x. At s = 0,
# t = 1 the partial derivative is -x, the spatial gradient is e - 1, and the transport term is
# therefore (e-1)x, so the correction is x(e-2) = 0.718x.
def scalar_check(x_value: float = 1.0):
    """Return every term of the identity at s = 0, t = 1, each one computed here."""
    x = torch.tensor(x_value)
    s, t = torch.tensor(0.0), torch.tensor(1.0)
    average = scalar_average_velocity(x, s, t)  # v_{0,1}(x) = x(e-1).
    instant = x  # b(x, 0) = x on this flow.
    partial_s = torch.func.jvp(  # Derivative in the start time alone, with x and t held fixed.
        lambda ss: scalar_average_velocity(x, ss, t), (s,), (torch.ones_like(s),))[1]
    gradient_x = torch.func.jvp(  # Derivative in the state alone, the other half of d/ds.
        lambda xx: scalar_average_velocity(xx, s, t), (x,), (torch.ones_like(x),))[1]
    transport = instant * gradient_x  # b(x,s) . grad_x v, the second piece of the total derivative.
    total = partial_s + transport  # d/ds v_{s,t}(x) = x(e - 2).
    return {
        "v_0,1(1)": float(average),
        "b(1, 0)": float(instant),
        "partial_s v": float(partial_s),
        "grad_x v_0,1": float(gradient_x),
        "d/ds v_0,1": float(total),
        "identity gap": float((average - (instant + (t - s) * total)).abs()),
        "one map evaluation": float(x + (t - s) * average),
        "one Euler step": float(x + (t - s) * instant),
        "e": math.e,
    }


# %% 04. The Gaussian flow, where the correction vanishes
#
# A linear interpolant between paired Gaussians produces straight trajectories, and on a straight
# trajectory the average velocity equals the initial velocity, so the total derivative is zero.
def gaussian_check(x_value: float = 0.7):
    """v_{0,t}(x) = 2 + 2x for every end time, and b(x,0) is the same number."""
    x = torch.tensor(x_value)
    average = torch.stack([  # Three end times, all giving the same average velocity.
        (2.0 * t + (1.0 + 2.0 * t) * x - x) / t for t in (0.25, 0.5, 1.0)])
    instant = 2.0 + 2.0 * (x - 0.0) / 1.0  # b(x, 0) on the Gaussian flow.
    return {"v_0,t(0.7) spread": float(average.max() - average.min()),
            "v_0,1(0.7)": float(average[-1]), "b(0.7, 0)": float(instant)}


# %% 05. The MeanFlow target, by one forward-mode pass
#
# The tangent (x_1 - x_0, 1, 0) in the arguments (x, s, t) contracts the full Jacobian of v into
# the total derivative along the trajectory: the realized velocity in the state slot, one in the
# start-time slot, and zero in the end-time slot because the end time is held fixed.
def mean_flow_target(model, xs, s, t, realized):
    """b(x,s) + (t-s) d/ds v_{s,t}(x), detached, the right-hand side of the identity."""
    f = lambda xx, ss: model(xx, ss, t)
    _, total = torch.func.jvp(f, (xs, s), (realized, torch.ones_like(s)))  # [B, d] d/ds v.
    return (realized + (t - s)[:, None] * total).detach()  # [B, d] the target, held fixed.


def make_loss(model, data, batch_size, device):
    """Draw a coupling, draw times from the triangle, and regress v onto its own identity."""

    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.
        s = torch.rand(batch_size, device=device)  # [B] start time.
        t = s + (1.0 - s) * torch.rand(batch_size, device=device)  # [B] end time, s <= t <= 1.
        diagonal = torch.rand(batch_size, device=device) < DIAGONAL_FRACTION  # [B] which rows anchor.
        t = torch.where(diagonal, s, t)  # On those rows the correction term has a factor of zero.
        xs = (1 - s[:, None]) * x0 + s[:, None] * x1  # [B, 2] the interpolant I_s.
        realized = x1 - x0  # [B, 2] the derivative of the interpolant, available in closed form.
        target = mean_flow_target(model, xs, s, t, realized)
        return ((model(xs, s, t) - target) ** 2).mean()

    return loss_fn


# %% 06. Sampling, at any step count the reader chooses
@torch.no_grad()
def sample(model, x0, steps: int):
    """Cross [0, 1] in equal pieces, one evaluation of the map 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 = model.flow_map(x, s.expand(len(x)), t.expand(len(x)))  # X_{s,t}(x).
    return x  # [N, 2] generated points.


# %% 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_check()
    report.update(gaussian_check())
    if not args.quiet:
        print(f"on dx/dt = x at s=0, t=1: v = {report['v_0,1(1)']:.6f} = b {report['b(1, 0)']:.6f}"
              f" + correction {report['d/ds v_0,1']:.6f}")
        print(f"one map evaluation {report['one map evaluation']:.6f} against one Euler step "
              f"{report['one Euler step']:.6f}, exact e = {math.e:.6f}")

    model = TwoTimeMLP(dim=2, width=args.width).to(device)
    if not args.quiet:
        print(f"\ntraining MeanFlow  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 source points every run shares.
    samples: list[str] = []
    for steps in (1, 2, 4):
        points = sample(model, x0, steps)
        stats = mode_report(points)
        report[f"{steps}-step nearest distance"] = stats["nearest_distance"]
        report[f"{steps}-step mode entropy"] = stats["mode_entropy"]
        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}")
    one, four = sample(model, x0, 1), sample(model, x0, 4)
    report["one step against four"] = float((one - four).norm(dim=1).mean())
    report["final loss"] = sum(losses[-20:]) / len(losses[-20:])

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


if __name__ == "__main__":
    main()

	"""MeanFlow, the Eulerian identity written in the average-velocity parameterization.

	If we substitute X_{s,t}(x) = x + (t-s) v_{s,t}(x) into the Eulerian identity, we are left with one
	equation, v_{s,t}(x) = b(x,s) + (t-s) d/ds v_{s,t}(x), with the correction term a total derivative
	along the trajectory. One forward-mode pass along the tangent (b, 1, 0) in the arguments (x, s, t)
	computes that derivative without forming a Jacobian, and the target uses the realized velocity
	x_1 - x_0 of the interpolant in place of b, so a step costs one forward pass and one derivative.

	Chapter 6, Sections 6.13, 6.15 and 6.18 of the notes. Trained on the four Gaussian blobs of
	``cis6270.data.gaussian_mixture_2d``. Prints the identity worked by hand on dx/dt = x, where the
	average velocity is 1.718x, the instantaneous velocity is x and the correction is 0.718x, and then
	one-step, two-step and four-step samples from the trained map.

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

	DIAGONAL_FRACTION = 0.25 # Share of each batch drawn with t = s, where the identity is flow matching.


	# %% 02. The average velocity of the scalar flow, in closed form
	#
	# The exact map of dx/dt = x multiplies the state by e^{t-s}, so the displacement over the elapsed
	# time is x (e^{t-s} - 1)/(t-s). On the diagonal the limit is x, or b(x,t).
	def scalar_average_velocity(x, s, t):
	"""v_{s,t}(x) = x (e^{t-s} - 1)/(t - s) on the flow dx/dt = x."""
	elapsed = t - s
	return x * torch.expm1(elapsed) / elapsed # expm1 keeps the limit accurate near the diagonal.


	# %% 03. The identity, checked on the scalar flow by hand
	#
	# v_{s,t}(x) = b(x,s) + (t-s) d/ds v_{s,t}(x), with d/ds = partial_s + b(x,s) . grad_x. At s = 0,
	# t = 1 the partial derivative is -x, the spatial gradient is e - 1, and the transport term is
	# therefore (e-1)x, so the correction is x(e-2) = 0.718x.
	def scalar_check(x_value: float = 1.0):
	"""Return every term of the identity at s = 0, t = 1, each one computed here."""
	x = torch.tensor(x_value)
	s, t = torch.tensor(0.0), torch.tensor(1.0)
	average = scalar_average_velocity(x, s, t) # v_{0,1}(x) = x(e-1).
	instant = x # b(x, 0) = x on this flow.
	partial_s = torch.func.jvp( # Derivative in the start time alone, with x and t held fixed.
	lambda ss: scalar_average_velocity(x, ss, t), (s,), (torch.ones_like(s),))[1]
	gradient_x = torch.func.jvp( # Derivative in the state alone, the other half of d/ds.
	lambda xx: scalar_average_velocity(xx, s, t), (x,), (torch.ones_like(x),))[1]
	transport = instant * gradient_x # b(x,s) . grad_x v, the second piece of the total derivative.
	total = partial_s + transport # d/ds v_{s,t}(x) = x(e - 2).
	return {
	"v_0,1(1)": float(average),
	"b(1, 0)": float(instant),
	"partial_s v": float(partial_s),
	"grad_x v_0,1": float(gradient_x),
	"d/ds v_0,1": float(total),
	"identity gap": float((average - (instant + (t - s) * total)).abs()),
	"one map evaluation": float(x + (t - s) * average),
	"one Euler step": float(x + (t - s) * instant),
	"e": math.e,
	}


	# %% 04. The Gaussian flow, where the correction vanishes
	#
	# A linear interpolant between paired Gaussians produces straight trajectories, and on a straight
	# trajectory the average velocity equals the initial velocity, so the total derivative is zero.
	def gaussian_check(x_value: float = 0.7):
	"""v_{0,t}(x) = 2 + 2x for every end time, and b(x,0) is the same number."""
	x = torch.tensor(x_value)
	average = torch.stack([ # Three end times, all giving the same average velocity.
	(2.0 * t + (1.0 + 2.0 * t) * x - x) / t for t in (0.25, 0.5, 1.0)])
	instant = 2.0 + 2.0 * (x - 0.0) / 1.0 # b(x, 0) on the Gaussian flow.
	return {"v_0,t(0.7) spread": float(average.max() - average.min()),
	"v_0,1(0.7)": float(average[-1]), "b(0.7, 0)": float(instant)}


	# %% 05. The MeanFlow target, by one forward-mode pass
	#
	# The tangent (x_1 - x_0, 1, 0) in the arguments (x, s, t) contracts the full Jacobian of v into
	# the total derivative along the trajectory: the realized velocity in the state slot, one in the
	# start-time slot, and zero in the end-time slot because the end time is held fixed.
	def mean_flow_target(model, xs, s, t, realized):
	"""b(x,s) + (t-s) d/ds v_{s,t}(x), detached, the right-hand side of the identity."""
	f = lambda xx, ss: model(xx, ss, t)
	_, total = torch.func.jvp(f, (xs, s), (realized, torch.ones_like(s))) # [B, d] d/ds v.
	return (realized + (t - s)[:, None] * total).detach() # [B, d] the target, held fixed.


	def make_loss(model, data, batch_size, device):
	"""Draw a coupling, draw times from the triangle, and regress v onto its own identity."""

	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.
	s = torch.rand(batch_size, device=device) # [B] start time.
	t = s + (1.0 - s) * torch.rand(batch_size, device=device) # [B] end time, s <= t <= 1.
	diagonal = torch.rand(batch_size, device=device) < DIAGONAL_FRACTION # [B] which rows anchor.
	t = torch.where(diagonal, s, t) # On those rows the correction term has a factor of zero.
	xs = (1 - s[:, None]) * x0 + s[:, None] * x1 # [B, 2] the interpolant I_s.
	realized = x1 - x0 # [B, 2] the derivative of the interpolant, available in closed form.
	target = mean_flow_target(model, xs, s, t, realized)
	return ((model(xs, s, t) - target) ** 2).mean()

	return loss_fn


	# %% 06. Sampling, at any step count the reader chooses
	@torch.no_grad()
	def sample(model, x0, steps: int):
	"""Cross [0, 1] in equal pieces, one evaluation of the map 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 = model.flow_map(x, s.expand(len(x)), t.expand(len(x))) # X_{s,t}(x).
	return x # [N, 2] generated points.


	# %% 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_check()
	report.update(gaussian_check())
	if not args.quiet:
	print(f"on dx/dt = x at s=0, t=1: v = {report['v_0,1(1)']:.6f} = b {report['b(1, 0)']:.6f}"
	f" + correction {report['d/ds v_0,1']:.6f}")
	print(f"one map evaluation {report['one map evaluation']:.6f} against one Euler step "
	f"{report['one Euler step']:.6f}, exact e = {math.e:.6f}")

	model = TwoTimeMLP(dim=2, width=args.width).to(device)
	if not args.quiet:
	print(f"\ntraining MeanFlow 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 source points every run shares.
	samples: list[str] = []
	for steps in (1, 2, 4):
	points = sample(model, x0, steps)
	stats = mode_report(points)
	report[f"{steps}-step nearest distance"] = stats["nearest_distance"]
	report[f"{steps}-step mode entropy"] = stats["mode_entropy"]
	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}")
	one, four = sample(model, x0, 1), sample(model, x0, 4)
	report["one step against four"] = float((one - four).norm(dim=1).mean())
	report["final loss"] = sum(losses[-20:]) / len(losses[-20:])

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


	if __name__ == "__main__":
	main()