File size: 10,467 Bytes
0ba9d09 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 | """The worked numbers of Chapter 2, asserted against the Lecture 2 scripts.
Every figure here is quoted in the chapter: the Euler landings at one, two and four steps, the flux
balance on [-1, 1], both sides of the continuity equation at x = 2, the single endpoint pair, the
average at the crossing, the nine-unit gap between the two objectives, the table of three
interpolants at t = 0.25, and the parameter count of the MNIST U-Net. Nothing in this file trains.
"""
import math
import sys
from pathlib import Path
import pytest
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # Run from a fresh clone.
from lecture_2 import flow_matching_2d as fm2d # noqa: E402
from lecture_2 import interpolants as interp # noqa: E402
from lecture_2 import velocity_fields as vf # noqa: E402
LOG2 = math.log(2.0)
# %% 01. Velocity fields, trajectories and the flow map (Sections 2.2 to 2.4)
def test_trajectories_match_the_chapter():
report = vf.trajectories()
assert report["constant_field_trajectory"] == [1.0, 1.5, 2.0, 2.5, 3.0]
assert report["expansion_x_half_from_two"] == pytest.approx(2.828, abs=1e-3)
assert report["expansion_endpoints"] == [-2.0, 0.0, 2.0]
assert report["expansion_outer_spacing"] == pytest.approx(4.0)
assert report["contraction_endpoints"] == [-0.5, 0.0, 0.5]
assert report["contraction_outer_spacing"] == pytest.approx(1.0)
def test_rotation_preserves_the_radius():
report = vf.trajectories()
assert report["rotation_quarter_turn"] == [0.0, 1.0]
assert report["rotation_half_turn"] == [-1.0, 0.0]
assert report["rotation_radius"] == pytest.approx(1.0, abs=1e-6)
assert report["rotation_squared_norm_rate"] == pytest.approx(0.0, abs=1e-6)
def test_flow_map_is_the_identity_at_zero_and_doubles_at_one():
report = vf.flow_map_report()
assert report["flow_map_identity_error"] == pytest.approx(0.0, abs=1e-6)
assert report["flow_map_jacobian_det"] == pytest.approx(2.0, abs=1e-6)
assert report["trajectory_gap_minimum"] == pytest.approx(2.0, abs=1e-6)
# %% 02. Euler's method (Section 2.5)
def test_four_euler_steps_land_at_the_chapter_value():
landing = vf.euler(lambda x, t: x, torch.tensor(1.0), 4) # dx/dt = x from x_0 = 1.
assert float(landing) == pytest.approx(2.44140625, abs=1e-6)
def test_euler_ladder_errors_and_ratios():
report = vf.euler_ladder(4)
assert report["euler_step_counts"] == [1, 2, 4]
assert report["euler_landings"] == [2.0, 2.25, pytest.approx(2.4414, abs=1e-4)]
assert report["euler_errors"] == [pytest.approx(0.7183, abs=1e-4),
pytest.approx(0.4683, abs=1e-4),
pytest.approx(0.2769, abs=1e-4)]
assert report["euler_error_halving_ratios"] == [pytest.approx(1.53, abs=5e-3),
pytest.approx(1.69, abs=5e-3)]
assert report["euler_relative_error_four_steps"] == pytest.approx(0.1019, abs=1e-4)
def test_the_euler_error_closes_like_one_over_the_step_count():
report = vf.euler_ladder(256) # The error times the step count settles at e/2 = 1.359.
scaled = report["euler_error_times_steps"]
assert scaled[-1] == pytest.approx(math.e / 2.0, abs=0.02)
assert report["euler_error_halving_ratios"][-1] == pytest.approx(2.0, abs=0.02)
# %% 03. The pushforward, the flux and the continuity equation (Sections 2.6 and 2.7)
def test_the_expansion_field_doubles_the_standard_deviation():
report = vf.pushforward_moments(4096, seed=0)
assert report["pushforward_mean"] == pytest.approx(0.0, abs=0.05)
assert report["pushforward_variance"] == pytest.approx(4.0, abs=0.2)
def test_flux_balance_on_the_unit_interval():
report = vf.flux_balance()
assert report["density_at_one"] == pytest.approx(0.2420, abs=1e-4)
assert report["flux_at_minus_one"] == pytest.approx(-0.1677, abs=1e-4)
assert report["flux_at_one"] == pytest.approx(0.1677, abs=1e-4)
assert report["flux_balance"] == pytest.approx(-0.3354, abs=1e-4)
assert report["enclosed_mass"] == pytest.approx(0.6827, abs=1e-4)
# The balance and the derivative of the enclosed mass agree to four figures.
assert report["enclosed_mass_rate"] == pytest.approx(report["flux_balance"], abs=1e-4)
def test_both_sides_of_the_continuity_equation_at_x_equals_two():
report = vf.continuity_check()
assert report["density_at_two"] == pytest.approx(0.05399, abs=1e-5)
assert report["density_gradient_at_two"] == pytest.approx(-0.10798, abs=1e-5)
assert report["velocity_divergence"] == pytest.approx(LOG2, abs=1e-6)
assert report["continuity_time_derivative"] == pytest.approx(0.11227, abs=1e-5)
assert report["continuity_minus_flux_divergence"] == pytest.approx(0.11227, abs=1e-5)
def test_log_density_along_a_trajectory_falls_by_log_two():
report = vf.log_density_along_trajectory()
assert report["log_density_rate"] == pytest.approx(-LOG2)
assert report["log_density_change"] == pytest.approx(-LOG2, abs=1e-6)
assert report["log_density_accumulated"] == pytest.approx(-LOG2)
assert report["prior_density_at_origin"] == pytest.approx(0.3989, abs=1e-4)
assert report["final_density_at_origin"] == pytest.approx(0.1995, abs=1e-4)
assert report["density_ratio_at_origin"] == pytest.approx(2.0, abs=1e-6)
# %% 04. The conditional path and the crossing (Sections 2.12 and 2.13)
def test_the_single_endpoint_pair():
report = fm2d.worked_pair()
assert report["worked_pair_path"] == [-2.0, -0.75, 0.5, 1.75, 3.0]
assert report["worked_pair_velocity"] == pytest.approx(5.0)
def test_the_average_at_the_crossing():
report = fm2d.crossing_average()
assert report["crossing_state"] == pytest.approx(2.0)
assert report["crossing_velocities"] == [4.0, -2.0]
assert report["crossing_average_velocity"] == pytest.approx(1.0)
def test_the_two_losses_differ_by_nine_at_the_crossing():
report = fm2d.crossing_gap()
assert report["crossing_conditional_loss"] == [10.0, 9.0, 10.0]
assert report["crossing_marginal_loss"] == [1.0, 0.0, 1.0]
assert report["crossing_constant_gap"] == [9.0, 9.0, 9.0]
assert report["crossing_target_variance"] == pytest.approx(9.0)
# %% 05. The marginal velocity and the equivalence (Sections 2.13 and 2.15)
def test_the_marginal_velocity_at_the_two_ends():
x = torch.randn(64, 2)
# v_0(x) = E[X_1] - x = -x, and v_1(x) = x - E[X_0] = x, since the coupling is independent.
assert torch.allclose(fm2d.marginal_velocity(x, torch.zeros(64)), -x, atol=1e-5)
assert torch.allclose(fm2d.marginal_velocity(x, torch.ones(64)), x, atol=1e-5)
def test_the_marginal_velocity_transports_the_prior_onto_the_data():
from cis6270.metrics import mean_log_density, mode_report
torch.manual_seed(0)
x_t = torch.randn(512, 2) # X_0 ~ p_0.
for step in range(50): # Euler on the exact field, with no network anywhere.
t = torch.full((512,), step / 50.0)
x_t = x_t + fm2d.marginal_velocity(x_t, t) / 50.0
assert mode_report(x_t)["modes_found"] == 4
assert mean_log_density(x_t) > -3.0 # The data themselves measure -2.14.
def test_the_losses_decompose_exactly_on_any_batch():
from cis6270.nets import VelocityMLP
torch.manual_seed(0)
model = VelocityMLP(dim=2, width=32)
x_0, x_1 = torch.randn(256, 2), torch.randn(256, 2) + 2.0
t = torch.rand(256)
loss_cfm, loss_fm, gap, cross = fm2d.objectives(model, x_0, x_1, t)
# We expand the square and get L_CFM = L_FM + cross + gap term by term, on every batch.
assert float(loss_cfm) == pytest.approx(float(loss_fm + cross + gap), rel=1e-5)
def test_the_two_objectives_share_a_gradient():
from cis6270.data import gaussian_mixture_2d
from cis6270.nets import VelocityMLP
torch.manual_seed(0)
model = VelocityMLP(dim=2, width=64)
data = gaussian_mixture_2d(4096, seed=0)
report = fm2d.equivalence_check(model, data, 8192)
assert report["loss_cfm"] == pytest.approx(report["loss_fm_plus_gap"], abs=0.3)
assert abs(report["loss_cross_term"]) < 0.3 # Zero in expectation, small on a large batch.
assert report["gradient_cosine"] > 0.99 # The residual is the batch's own sampling error.
# %% 06. The three interpolants (Section 2.18)
def test_every_schedule_reaches_the_right_endpoints():
for name in interp.SCHEDULES:
assert interp.endpoint_and_table(name)[f"{name}_endpoint_error"] == pytest.approx(0.0,
abs=1e-6)
def test_the_table_of_three_interpolants_at_one_quarter():
expected = { # Table 2.1, at X_0 = (2, 0), X_1 = (0, 2) and t = 0.25.
"linear": ([1.50, 0.50], [-2.0, 2.0]),
"quadratic": ([1.875, 0.125], [-1.0, 1.0]),
"trigonometric": ([1.848, 0.765], [-1.202, 2.902]),
}
for name, (position, velocity) in expected.items():
report = interp.endpoint_and_table(name)
assert report[f"{name}_position_at_quarter"] == pytest.approx(position, abs=1e-3)
assert report[f"{name}_velocity_at_quarter"] == pytest.approx(velocity, abs=1e-3)
def test_only_the_trigonometric_schedule_keeps_identity_covariance():
trig = interp.covariance_check("trigonometric", 8192, seed=0)
assert trig["trigonometric_coefficient_squares"] == [1.0, 1.0, 1.0, 1.0, 1.0]
assert trig["trigonometric_covariance_deviation"] < 0.05
linear = interp.covariance_check("linear", 8192, seed=0)
assert linear["linear_coefficient_squares"] == [1.0, 0.625, 0.5, 0.625, 1.0]
assert linear["linear_covariance_deviation"] > 0.4 # Cov(X_{1/2}) is half the identity.
def test_the_quadratic_schedule_has_covered_a_sixteenth_at_one_quarter():
# beta_{1/4} = (1/4)^2 = 0.0625, so the state has covered 6.25 percent of the segment.
_, beta, _, _ = interp.quadratic(torch.tensor([0.25]))
assert float(beta) == pytest.approx(0.0625)
_, beta_half, _, _ = interp.quadratic(torch.tensor([0.5]))
assert float(beta_half) == pytest.approx(0.25) # A quarter of the distance at the midpoint.
# %% 07. The MNIST U-Net (Section 2.20)
def test_the_unet_has_the_parameter_count_the_chapter_quotes():
from lecture_2 import flow_matching_mnist as mnist
model = mnist.FlowUNet() # The default base width of 32, so the widths are 32, 64 and 128.
assert sum(p.numel() for p in model.parameters()) == 471_265
|