Upload 87 files

0ba9d09 verified about 1 month ago

11.3 kB

"""Three interpolants between the prior and the data.
We can replace the weights 1 - t and t by any pair of schedules alpha_t and beta_t with
alpha_0 = 1, beta_0 = 0, alpha_1 = 0 and beta_1 = 1, and the conditional flow-matching argument
stays intact. The regression target then becomes the derivative alpha_dot X_0 + beta_dot X_1. The
straight line, the t-squared schedule and the trigonometric schedule differ in where the path goes,
how fast it gets there, and what the intermediate densities look like, so the same network trained
on each produces different samples.
Chapter 2, Section 2.18 of the notes. Trained on ``cis6270.data.gaussian_mixture_2d``. Prints the
chapter's table of positions and velocities at t = 0.25, the endpoint check for each schedule, the
identity covariance of the trigonometric schedule, and the samples each trained field produces at a
fine and at a coarse step budget.
    python lecture_2/interpolants.py                 # train three fields, about fifteen seconds
    python lecture_2/interpolants.py --steps 20      # a first look, with no time to learn
"""

# %% 01. Imports and configuration
import argparse
import math

import torch

from cis6270.data import gaussian_mixture_2d
from cis6270.metrics import mean_log_density, mode_report
from cis6270.nets import VelocityMLP
from cis6270.runner import print_report, save_run, set_seed, train

DIM = 2  # d, the width of the state.
HALF_PI = math.pi / 2.0  # The constant in the trigonometric schedule and in its derivative.
TRAIN_STEPS = 1000  # Optimizer steps for each of the three fields.
BATCH_SIZE = 256  # Endpoint pairs per step.
LEARNING_RATE = 3e-3  # Optimizer step size.
WIDTH = 128  # Hidden width of the velocity network.
SAMPLE_STEPS = 64  # Euler steps in the fine sampling budget.
COARSE_STEPS = 8  # Euler steps in the coarse budget, where a varying velocity costs accuracy.
NUM_SAMPLES = 1024  # Points generated per schedule and per budget.
DATA_POINTS = 4096  # Size of the training set drawn from the four-Gaussian target.
COVARIANCE_POINTS = 8192  # Endpoint pairs used for the empirical covariance of X_t.
GRID = torch.tensor([0.0, 0.25, 0.50, 0.75, 1.0])  # [5] the chapter's grid of times.


# %% 02. The three schedules
#
# Section 2.18. Each returns alpha_t, beta_t and their time derivatives, evaluated on a batch of
# times. The pair of schedules fixes the path and the pair of derivatives fixes the target.
def linear(t):
    """alpha_t = 1 - t and beta_t = t, the straight segment traversed at constant speed."""
    return 1.0 - t, t, -torch.ones_like(t), torch.ones_like(t)  # Four tensors of shape [B].


def quadratic(t):
    """alpha_t = 1 - t^2 and beta_t = t^2, the same segment entered slowly and left quickly."""
    return 1.0 - t ** 2, t ** 2, -2.0 * t, 2.0 * t  # U_t = 2t (X_1 - X_0), zero speed at t = 0.


def trigonometric(t):
    """alpha_t = cos(pi t / 2) and beta_t = sin(pi t / 2), a quarter circle in the coefficients."""
    angle = HALF_PI * t  # [B] the angle the coefficient pair has turned through.
    return (torch.cos(angle), torch.sin(angle),
            -HALF_PI * torch.sin(angle), HALF_PI * torch.cos(angle))


SCHEDULES = {"linear": linear, "quadratic": quadratic, "trigonometric": trigonometric}


def interpolate(schedule, x_0, x_1, t):
    """X_t = alpha_t X_0 + beta_t X_1 and U_t = alpha_dot X_0 + beta_dot X_1. Two [B, d] tensors."""
    alpha, beta, alpha_dot, beta_dot = schedule(t)  # [B] each, the two weights and their rates.
    x_t = alpha[:, None] * x_0 + beta[:, None] * x_1  # [B, d] the state on the conditional path.
    u_t = alpha_dot[:, None] * x_0 + beta_dot[:, None] * x_1  # [B, d] the regression target.
    return x_t, u_t


# %% 03. The endpoints and the velocity at one time
#
# Section 2.18. The two ends are correct when alpha_0 = 1, beta_0 = 0, alpha_1 = 0 and beta_1 = 1.
# Table 2.1 reads the three paths at X_0 = (2, 0), X_1 = (0, 2) and t = 0.25.
def endpoint_and_table(name):
    """The endpoint error of one schedule and its position and velocity at the chapter's input."""
    x_0 = torch.tensor([[2.0, 0.0]])  # [1, d] the prior endpoint of the table.
    x_1 = torch.tensor([[0.0, 2.0]])  # [1, d] the data endpoint of the table.
    schedule = SCHEDULES[name]
    ends, _ = interpolate(schedule, x_0.expand(2, DIM), x_1.expand(2, DIM),
                          torch.tensor([0.0, 1.0]))  # [2, d] the states at t = 0 and t = 1.
    start_error = (ends[0] - x_0[0]).abs().max()  # X_0 is returned at t = 0.
    finish_error = (ends[1] - x_1[0]).abs().max()  # X_1 is returned at t = 1.
    x_t, u_t = interpolate(schedule, x_0, x_1, torch.tensor([0.25]))  # One quarter of the interval.
    return {
        f"{name}_endpoint_error": float(torch.maximum(start_error, finish_error)),
        f"{name}_position_at_quarter": [round(float(v), 4) + 0.0 for v in x_t[0]],
        f"{name}_velocity_at_quarter": [round(float(v), 4) + 0.0 for v in u_t[0]],
    }


# %% 04. Identity covariance
#
# Section 2.18. For independent, centered endpoints with identity covariance the interpolation has
# covariance (alpha_t^2 + beta_t^2) I. That factor is cos^2 + sin^2 = 1 at every time for the
# trigonometric schedule, and below one between the ends for the other two.
def covariance_check(name, num_points, seed):
    """The coefficient squares on the grid, and the largest deviation of Cov(X_t) from the identity."""
    set_seed(seed)  # Same draw for every schedule, so the comparison is between schedules alone.
    x_0 = torch.randn(num_points, DIM)  # [N, d] a centered, white prior endpoint.
    x_1 = torch.randn(num_points, DIM)  # [N, d] a centered, white data endpoint, independent of it.
    identity = torch.eye(DIM)  # [d, d] the covariance the chapter's hypothesis assigns each endpoint.
    squares, deviations = [], []
    for t in GRID:  # One time at a time, since we measure the covariance over the batch.
        times = t.expand(num_points)  # [N] the same time for every pair.
        alpha, beta, _, _ = SCHEDULES[name](times)
        squares.append(round(float(alpha[0] ** 2 + beta[0] ** 2), 4))  # The exact variance factor.
        x_t, _ = interpolate(SCHEDULES[name], x_0, x_1, times)  # [N, d] the interpolated states.
        covariance = (x_t.T @ x_t) / num_points  # [d, d] the empirical second moment, centered at 0.
        deviations.append(float((covariance - identity).abs().max()))
    return {
        f"{name}_coefficient_squares": squares,
        f"{name}_covariance_deviation": max(deviations),
    }


# %% 05. Training one field per interpolant
#
# Section 2.18 and Section 2.19. We run the loop of the straight-line case with the target
# replaced, so any difference in the samples comes from the schedule alone.
def train_one(name, data, args):
    """Train v_theta against the conditional velocity of one schedule and return the loss history."""
    set_seed(args.seed)  # The same initialization and the same batches for every schedule.
    model = VelocityMLP(dim=DIM, width=args.width)  # v_theta(x, t).
    schedule = SCHEDULES[name]

    def loss_fn(step):
        index = torch.randint(0, len(data), (args.batch_size,))  # [B] which data point to pair.
        x_1 = data[index]  # [B, d] the data endpoint.
        x_0 = torch.randn_like(x_1)  # [B, d] the prior endpoint, independent of it.
        t = torch.rand(args.batch_size)  # [B] t ~ Uniform(0, 1).
        x_t, u_t = interpolate(schedule, x_0, x_1, t)  # [B, d] the state and the target.
        return (model(x_t, t) - u_t).square().sum(1).mean()  # ||v_theta(X_t, t) - U_t||^2.

    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
    return model, losses


# %% 06. Euler sampling
#
# Section 2.19. We sample every schedule with the same integrator, so a reader comparing them only
# has to hold the step budget fixed.
@torch.no_grad()
def sample(model, num_samples, steps):
    """Integrate dX_t/dt = v_theta(X_t, t) from t = 0 to t = 1 on a uniform grid. Returns [N, d]."""
    x_t = torch.randn(num_samples, DIM)  # [N, d] the draw X_0 ~ p_0.
    dt = 1.0 / steps  # Delta t, the width of one interval.
    for step in range(steps):  # One network evaluation per step.
        x_t = x_t + dt * model(x_t, torch.full((num_samples,), step * dt))  # One Euler step.
    return x_t


def measure(model, args):
    """Mode coverage and mean exact log density at the fine and the coarse step budget."""
    set_seed(args.seed + 1)  # The same starting noise for every schedule.
    fine = sample(model, args.samples, args.sample_steps)  # [N, d] points from the fine grid.
    coarse = sample(model, args.samples, args.coarse_steps)  # [N, d] points from the coarse grid.
    report = mode_report(fine)
    return {
        "modes_found": report["modes_found"],
        "mode_entropy": report["mode_entropy"],
        "nearest_distance": report["nearest_distance"],
        "mean_log_density": mean_log_density(fine),
        "coarse_log_density": mean_log_density(coarse),
    }


# %% 07. The report
def main():
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--steps", type=int, default=TRAIN_STEPS, help="optimizer steps per field")
    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="pairs per step")
    parser.add_argument("--lr", type=float, default=LEARNING_RATE, help="learning rate")
    parser.add_argument("--width", type=int, default=WIDTH, help="hidden width of the network")
    parser.add_argument("--samples", type=int, default=NUM_SAMPLES, help="points to generate")
    parser.add_argument("--sample-steps", type=int, default=SAMPLE_STEPS, help="fine Euler steps")
    parser.add_argument("--coarse-steps", type=int, default=COARSE_STEPS, help="coarse Euler steps")
    parser.add_argument("--seed", type=int, default=0, help="random seed")
    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
    parser.add_argument("--quiet", action="store_true", help="print only the final report")
    args = parser.parse_args()

    data = gaussian_mixture_2d(DATA_POINTS, seed=args.seed)  # [N, d] the target p_1 = p_data.
    report = {}
    for name in SCHEDULES:  # Linear, quadratic, trigonometric, in the chapter's order.
        report.update(endpoint_and_table(name))  # Section 2.18, before any training.
        report.update(covariance_check(name, COVARIANCE_POINTS, args.seed))
        if not args.quiet:
            print(f"training the {name} schedule")
        model, losses = train_one(name, data, args)
        report[f"{name}_final_loss"] = sum(losses[-100:]) / len(losses[-100:])
        # Each loss regresses on its own target, so the three losses are not comparable with each
        # other. The three measurements below are comparable.
        for key, value in measure(model, args).items():
            report[f"{name}_{key}"] = value
        save_run(f"{args.out}/{name}" if args.out else "", model=model, losses=losses)

    report["data_mean_log_density"] = mean_log_density(data)  # The same measurement on real points.
    print_report("three interpolants on four Gaussians", report)
    save_run(args.out, config=vars(args), report=report)


# %% 08. Run the script
if __name__ == "__main__":
    main()

	"""Three interpolants between the prior and the data.

	We can replace the weights 1 - t and t by any pair of schedules alpha_t and beta_t with
	alpha_0 = 1, beta_0 = 0, alpha_1 = 0 and beta_1 = 1, and the conditional flow-matching argument
	stays intact. The regression target then becomes the derivative alpha_dot X_0 + beta_dot X_1. The
	straight line, the t-squared schedule and the trigonometric schedule differ in where the path goes,
	how fast it gets there, and what the intermediate densities look like, so the same network trained
	on each produces different samples.

	Chapter 2, Section 2.18 of the notes. Trained on ``cis6270.data.gaussian_mixture_2d``. Prints the
	chapter's table of positions and velocities at t = 0.25, the endpoint check for each schedule, the
	identity covariance of the trigonometric schedule, and the samples each trained field produces at a
	fine and at a coarse step budget.

	python lecture_2/interpolants.py # train three fields, about fifteen seconds
	python lecture_2/interpolants.py --steps 20 # a first look, with no time to learn
	"""

	# %% 01. Imports and configuration
	import argparse
	import math

	import torch

	from cis6270.data import gaussian_mixture_2d
	from cis6270.metrics import mean_log_density, mode_report
	from cis6270.nets import VelocityMLP
	from cis6270.runner import print_report, save_run, set_seed, train

	DIM = 2 # d, the width of the state.
	HALF_PI = math.pi / 2.0 # The constant in the trigonometric schedule and in its derivative.
	TRAIN_STEPS = 1000 # Optimizer steps for each of the three fields.
	BATCH_SIZE = 256 # Endpoint pairs per step.
	LEARNING_RATE = 3e-3 # Optimizer step size.
	WIDTH = 128 # Hidden width of the velocity network.
	SAMPLE_STEPS = 64 # Euler steps in the fine sampling budget.
	COARSE_STEPS = 8 # Euler steps in the coarse budget, where a varying velocity costs accuracy.
	NUM_SAMPLES = 1024 # Points generated per schedule and per budget.
	DATA_POINTS = 4096 # Size of the training set drawn from the four-Gaussian target.
	COVARIANCE_POINTS = 8192 # Endpoint pairs used for the empirical covariance of X_t.
	GRID = torch.tensor([0.0, 0.25, 0.50, 0.75, 1.0]) # [5] the chapter's grid of times.


	# %% 02. The three schedules
	#
	# Section 2.18. Each returns alpha_t, beta_t and their time derivatives, evaluated on a batch of
	# times. The pair of schedules fixes the path and the pair of derivatives fixes the target.
	def linear(t):
	"""alpha_t = 1 - t and beta_t = t, the straight segment traversed at constant speed."""
	return 1.0 - t, t, -torch.ones_like(t), torch.ones_like(t) # Four tensors of shape [B].


	def quadratic(t):
	"""alpha_t = 1 - t^2 and beta_t = t^2, the same segment entered slowly and left quickly."""
	return 1.0 - t 2, t 2, -2.0 * t, 2.0 * t # U_t = 2t (X_1 - X_0), zero speed at t = 0.


	def trigonometric(t):
	"""alpha_t = cos(pi t / 2) and beta_t = sin(pi t / 2), a quarter circle in the coefficients."""
	angle = HALF_PI * t # [B] the angle the coefficient pair has turned through.
	return (torch.cos(angle), torch.sin(angle),
	-HALF_PI * torch.sin(angle), HALF_PI * torch.cos(angle))


	SCHEDULES = {"linear": linear, "quadratic": quadratic, "trigonometric": trigonometric}


	def interpolate(schedule, x_0, x_1, t):
	"""X_t = alpha_t X_0 + beta_t X_1 and U_t = alpha_dot X_0 + beta_dot X_1. Two [B, d] tensors."""
	alpha, beta, alpha_dot, beta_dot = schedule(t) # [B] each, the two weights and their rates.
	x_t = alpha[:, None] * x_0 + beta[:, None] * x_1 # [B, d] the state on the conditional path.
	u_t = alpha_dot[:, None] * x_0 + beta_dot[:, None] * x_1 # [B, d] the regression target.
	return x_t, u_t


	# %% 03. The endpoints and the velocity at one time
	#
	# Section 2.18. The two ends are correct when alpha_0 = 1, beta_0 = 0, alpha_1 = 0 and beta_1 = 1.
	# Table 2.1 reads the three paths at X_0 = (2, 0), X_1 = (0, 2) and t = 0.25.
	def endpoint_and_table(name):
	"""The endpoint error of one schedule and its position and velocity at the chapter's input."""
	x_0 = torch.tensor([[2.0, 0.0]]) # [1, d] the prior endpoint of the table.
	x_1 = torch.tensor([[0.0, 2.0]]) # [1, d] the data endpoint of the table.
	schedule = SCHEDULES[name]
	ends, _ = interpolate(schedule, x_0.expand(2, DIM), x_1.expand(2, DIM),
	torch.tensor([0.0, 1.0])) # [2, d] the states at t = 0 and t = 1.
	start_error = (ends[0] - x_0[0]).abs().max() # X_0 is returned at t = 0.
	finish_error = (ends[1] - x_1[0]).abs().max() # X_1 is returned at t = 1.
	x_t, u_t = interpolate(schedule, x_0, x_1, torch.tensor([0.25])) # One quarter of the interval.
	return {
	f"{name}_endpoint_error": float(torch.maximum(start_error, finish_error)),
	f"{name}_position_at_quarter": [round(float(v), 4) + 0.0 for v in x_t[0]],
	f"{name}_velocity_at_quarter": [round(float(v), 4) + 0.0 for v in u_t[0]],
	}


	# %% 04. Identity covariance
	#
	# Section 2.18. For independent, centered endpoints with identity covariance the interpolation has
	# covariance (alpha_t^2 + beta_t^2) I. That factor is cos^2 + sin^2 = 1 at every time for the
	# trigonometric schedule, and below one between the ends for the other two.
	def covariance_check(name, num_points, seed):
	"""The coefficient squares on the grid, and the largest deviation of Cov(X_t) from the identity."""
	set_seed(seed) # Same draw for every schedule, so the comparison is between schedules alone.
	x_0 = torch.randn(num_points, DIM) # [N, d] a centered, white prior endpoint.
	x_1 = torch.randn(num_points, DIM) # [N, d] a centered, white data endpoint, independent of it.
	identity = torch.eye(DIM) # [d, d] the covariance the chapter's hypothesis assigns each endpoint.
	squares, deviations = [], []
	for t in GRID: # One time at a time, since we measure the covariance over the batch.
	times = t.expand(num_points) # [N] the same time for every pair.
	alpha, beta, _, _ = SCHEDULES[name](times)
	squares.append(round(float(alpha[0] 2 + beta[0] 2), 4)) # The exact variance factor.
	x_t, _ = interpolate(SCHEDULES[name], x_0, x_1, times) # [N, d] the interpolated states.
	covariance = (x_t.T @ x_t) / num_points # [d, d] the empirical second moment, centered at 0.
	deviations.append(float((covariance - identity).abs().max()))
	return {
	f"{name}_coefficient_squares": squares,
	f"{name}_covariance_deviation": max(deviations),
	}


	# %% 05. Training one field per interpolant
	#
	# Section 2.18 and Section 2.19. We run the loop of the straight-line case with the target
	# replaced, so any difference in the samples comes from the schedule alone.
	def train_one(name, data, args):
	"""Train v_theta against the conditional velocity of one schedule and return the loss history."""
	set_seed(args.seed) # The same initialization and the same batches for every schedule.
	model = VelocityMLP(dim=DIM, width=args.width) # v_theta(x, t).
	schedule = SCHEDULES[name]

	def loss_fn(step):
	index = torch.randint(0, len(data), (args.batch_size,)) # [B] which data point to pair.
	x_1 = data[index] # [B, d] the data endpoint.
	x_0 = torch.randn_like(x_1) # [B, d] the prior endpoint, independent of it.
	t = torch.rand(args.batch_size) # [B] t ~ Uniform(0, 1).
	x_t, u_t = interpolate(schedule, x_0, x_1, t) # [B, d] the state and the target.
	return (model(x_t, t) - u_t).square().sum(1).mean() # \|\|v_theta(X_t, t) - U_t\|\|^2.

	losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
	return model, losses


	# %% 06. Euler sampling
	#
	# Section 2.19. We sample every schedule with the same integrator, so a reader comparing them only
	# has to hold the step budget fixed.
	@torch.no_grad()
	def sample(model, num_samples, steps):
	"""Integrate dX_t/dt = v_theta(X_t, t) from t = 0 to t = 1 on a uniform grid. Returns [N, d]."""
	x_t = torch.randn(num_samples, DIM) # [N, d] the draw X_0 ~ p_0.
	dt = 1.0 / steps # Delta t, the width of one interval.
	for step in range(steps): # One network evaluation per step.
	x_t = x_t + dt * model(x_t, torch.full((num_samples,), step * dt)) # One Euler step.
	return x_t


	def measure(model, args):
	"""Mode coverage and mean exact log density at the fine and the coarse step budget."""
	set_seed(args.seed + 1) # The same starting noise for every schedule.
	fine = sample(model, args.samples, args.sample_steps) # [N, d] points from the fine grid.
	coarse = sample(model, args.samples, args.coarse_steps) # [N, d] points from the coarse grid.
	report = mode_report(fine)
	return {
	"modes_found": report["modes_found"],
	"mode_entropy": report["mode_entropy"],
	"nearest_distance": report["nearest_distance"],
	"mean_log_density": mean_log_density(fine),
	"coarse_log_density": mean_log_density(coarse),
	}


	# %% 07. The report
	def main():
	parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
	parser.add_argument("--steps", type=int, default=TRAIN_STEPS, help="optimizer steps per field")
	parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="pairs per step")
	parser.add_argument("--lr", type=float, default=LEARNING_RATE, help="learning rate")
	parser.add_argument("--width", type=int, default=WIDTH, help="hidden width of the network")
	parser.add_argument("--samples", type=int, default=NUM_SAMPLES, help="points to generate")
	parser.add_argument("--sample-steps", type=int, default=SAMPLE_STEPS, help="fine Euler steps")
	parser.add_argument("--coarse-steps", type=int, default=COARSE_STEPS, help="coarse Euler steps")
	parser.add_argument("--seed", type=int, default=0, help="random seed")
	parser.add_argument("--out", type=str, default="", help="directory for the run's output")
	parser.add_argument("--quiet", action="store_true", help="print only the final report")
	args = parser.parse_args()

	data = gaussian_mixture_2d(DATA_POINTS, seed=args.seed) # [N, d] the target p_1 = p_data.
	report = {}
	for name in SCHEDULES: # Linear, quadratic, trigonometric, in the chapter's order.
	report.update(endpoint_and_table(name)) # Section 2.18, before any training.
	report.update(covariance_check(name, COVARIANCE_POINTS, args.seed))
	if not args.quiet:
	print(f"training the {name} schedule")
	model, losses = train_one(name, data, args)
	report[f"{name}_final_loss"] = sum(losses[-100:]) / len(losses[-100:])
	# Each loss regresses on its own target, so the three losses are not comparable with each
	# other. The three measurements below are comparable.
	for key, value in measure(model, args).items():
	report[f"{name}_{key}"] = value
	save_run(f"{args.out}/{name}" if args.out else "", model=model, losses=losses)

	report["data_mean_log_density"] = mean_log_density(data) # The same measurement on real points.
	print_report("three interpolants on four Gaussians", report)
	save_run(args.out, config=vars(args), report=report)


	# %% 08. Run the script
	if __name__ == "__main__":
	main()