File size: 12,398 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 237 238 239 240 241 242 243 244 245 246 247 248 249 | """Simulation-Free Score and Flow Matching: two closed-form targets at one bridge sample.
With the entropic plan computed on a minibatch, the bridge is a mixture of Brownian bridges between
pairs drawn from that plan, and a Brownian bridge at a time t is Gaussian. One normal draw gives a
sample, and at that sample both the bridge drift and the score of the bridge marginal are available
without simulating a path. We subtract half the noise variance times the score. That turns the
drift into the velocity of the probability flow, and its dependence on epsilon cancels.
Chapter 7, Section 7.19.3 of the notes. Runs on the endpoint pair x_0 = 0, x_1 = 2 at epsilon = 1,
and on the two-point batch of Section 7.8.1. Prints the two regression targets at t = 0.25 and at
t = 0.75, the scaled score at every time on the grid, and the entropic plan of the batch.
python lecture_7/sf2m.py # 101 times on the grid
python lecture_7/sf2m.py --steps 20 # a coarse time grid, for a first look
"""
# %% 01. Imports and configuration
from __future__ import annotations
import torch
from cis6270.runner import common_parser, print_report, save_run, set_seed
DTYPE = torch.float64 # The chapter quotes four figures of both regression targets.
X0, X1 = 0.0, 2.0 # The endpoint pair the chapter works through.
EPSILON = 1.0 # The reference noise variance per unit time.
DRAW = 0.4 # The normal draw z the chapter uses at both times.
# %% 02. One sample from the bridge
#
# The Brownian bridge at time t is Gaussian with mean the straight interpolation and variance
# epsilon t (1 - t), so a sample costs one normal draw and no intermediate step.
# Section 7.19.3.
def bridge_mean(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
"""m_t = (1 - t) x_0 + t x_1, the mean of the bridge between the two ends."""
return (1.0 - t) * x0 + t * x1
def bridge_std(t: torch.Tensor, epsilon: float = EPSILON) -> torch.Tensor:
"""sigma_t = sqrt(epsilon t (1 - t)), zero at both ends and largest at t = 1/2."""
return torch.sqrt(epsilon * t * (1.0 - t))
def bridge_sample(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor, z: torch.Tensor,
epsilon: float = EPSILON) -> torch.Tensor:
"""x = m_t + sigma_t z, one draw from the bridge at time t."""
return bridge_mean(x0, x1, t) + bridge_std(t, epsilon) * z
# %% 03. The two conditional quantities
#
# Both are conditional on the drawn pair, and neither depends on a quantity a network has to
# estimate first. Section 7.19.3.
def conditional_drift(x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
"""u_t(x | x_0, x_1) = (x_1 - x) / (1 - t), the remaining displacement over the remaining time."""
return (x1 - x) / (1.0 - t)
def conditional_score(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor,
epsilon: float = EPSILON) -> torch.Tensor:
"""grad log p_t(x | x_0, x_1) = (m_t - x) / (epsilon t (1 - t)), displacement over variance."""
return (bridge_mean(x0, x1, t) - x) / (epsilon * t * (1.0 - t))
# %% 04. From the bridge drift to a probability flow velocity
#
# We subtract half the noise variance times the score. That converts a diffusion drift into the
# velocity of the probability flow, and that flow has the same marginals with no noise at all.
# The second form is the straight-line velocity plus a coefficient times the displacement of the
# sample from the mean, with the dependence on epsilon cancelled. Section 7.19.3.
def flow_coefficient(t: torch.Tensor) -> torch.Tensor:
"""(1 - 2t) / (2 t (1 - t)), positive before t = 1/2 and negative after."""
return (1.0 - 2.0 * t) / (2.0 * t * (1.0 - t))
def flow_velocity(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor,
epsilon: float = EPSILON) -> torch.Tensor:
"""u_t - (epsilon / 2) grad log p_t, the velocity the first network regresses onto."""
return (conditional_drift(x1, x, t)
- 0.5 * epsilon * conditional_score(x0, x1, x, t, epsilon))
def flow_velocity_collected(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor,
t: torch.Tensor) -> torch.Tensor:
"""(x_1 - x_0) + (1 - 2t) / (2 t (1 - t)) (x - m_t), the same velocity with delta collected."""
return (x1 - x0) + flow_coefficient(t) * (x - bridge_mean(x0, x1, t))
def scaled_score(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor,
epsilon: float = EPSILON) -> torch.Tensor:
"""sigma_t grad log p_t = (m_t - x) / sigma_t = -z, with unit scale at every time."""
return bridge_std(t, epsilon) * conditional_score(x0, x1, x, t, epsilon)
# %% 05. The loss and the sampler
#
# With lambda(t) = sigma_t the score target is -z, so no part of the schedule dominates the
# gradient. At sampling time we fix a noise level g(t) and integrate the two networks together,
# with g^2 = epsilon recovering the reference noise and g = 0 the probability flow.
# Section 7.19.3.
def sf2m_loss(velocity: torch.Tensor, score: torch.Tensor, velocity_target: torch.Tensor,
score_target: torch.Tensor, t: torch.Tensor, epsilon: float = EPSILON
) -> torch.Tensor:
"""||v_theta - u^o_t||^2 + lambda(t)^2 ||s_theta - grad log p_t||^2, with lambda(t) = sigma_t."""
weight = bridge_std(t, epsilon) ** 2 # lambda(t)^2, and it turns the score target into -z.
return ((velocity - velocity_target) ** 2 + weight * (score - score_target) ** 2).mean()
def sampler_drift(velocity: torch.Tensor, score: torch.Tensor, noise_level: float
) -> torch.Tensor:
"""v_theta + (1/2) g(t)^2 s_theta, the drift of the family of samplers between the two ends."""
return velocity + 0.5 * noise_level ** 2 * score
# %% 06. The role of the coupling
#
# If we draw the pairs from the entropic plan at strength 2 epsilon, the learned equation at
# g^2 = epsilon is the Schrodinger bridge, and from any other coupling it is the Markov projection
# of that mixture of bridges. We compute the plan on a minibatch. Section 7.19.3.
def entropic_plan(x: torch.Tensor, y: torch.Tensor, a: torch.Tensor, b: torch.Tensor,
epsilon: float = EPSILON, sweeps: int = 200) -> torch.Tensor:
"""The entropic plan for the squared cost at strength 2 epsilon, by Sinkhorn on exp(-C / 2 eps)."""
C = (x[:, None] - y[None, :]) ** 2 # [n, m] the squared cost of the batch.
K = torch.exp(-C / (2.0 * epsilon)) # The Gibbs kernel a Brownian reference produces.
u, v = torch.ones_like(a), torch.ones_like(b)
for _ in range(sweeps):
u = a / (K @ v)
v = b / (K.T @ u)
return u[:, None] * K * v[None, :] # [n, m].
# %% 07. The report
def main() -> None:
parser = common_parser(__doc__.split("\n")[0])
parser.set_defaults(steps=101) # --steps is the number of times on the grid.
args = parser.parse_args()
set_seed(args.seed)
loud = not args.quiet
x0 = torch.tensor(X0, dtype=DTYPE)
x1 = torch.tensor(X1, dtype=DTYPE)
z = torch.tensor(DRAW, dtype=DTYPE)
worked = {}
for value in (0.25, 0.75):
t = torch.tensor(value, dtype=DTYPE)
x = bridge_sample(x0, x1, t, z)
worked[value] = {
"mean": float(bridge_mean(x0, x1, t)),
"std": float(bridge_std(t)),
"sample": float(x),
"drift": float(conditional_drift(x1, x, t)),
"score": float(conditional_score(x0, x1, x, t)),
"flow": float(flow_velocity(x0, x1, x, t)),
"collected": float(flow_velocity_collected(x0, x1, x, t)),
"scaled": float(scaled_score(x0, x1, x, t)),
"coefficient": float(flow_coefficient(t)),
}
if loud:
w = worked[value]
print(f"t = {value}, z = {DRAW}")
print(f" m_t = {w['mean']:.4f} sigma_t = {w['std']:.4f} x = {w['sample']:.4f}")
print(f" u_t = {w['drift']:.4f} grad log p_t = {w['score']:.4f}")
print(f" u^o_t = {w['flow']:.4f} collected form {w['collected']:.4f} "
f"coefficient {w['coefficient']:+.4f}")
print(f" sigma_t grad log p_t = {w['scaled']:.4f}, which is -z")
grid = torch.linspace(0.0, 1.0, max(args.steps, 3), dtype=DTYPE)[1:-1] # Interior times.
samples = bridge_sample(x0, x1, grid, z) # [levels] one bridge sample per time.
scaled = scaled_score(x0, x1, samples, grid)
scaled_gap = float((scaled + z).abs().max()) # The target is -z at every time.
flow_targets = flow_velocity(x0, x1, samples, grid)
score_targets = conditional_score(x0, x1, samples, grid)
loss_exact = float(sf2m_loss(flow_targets, score_targets, flow_targets, score_targets, grid))
loss_zero = float(sf2m_loss(torch.zeros_like(flow_targets), torch.zeros_like(score_targets),
flow_targets, score_targets, grid)) # An untrained pair of networks.
stratified = torch.special.ndtri(
(torch.arange(max(args.steps, 3), dtype=DTYPE) + 0.5) / max(args.steps, 3))
half = torch.tensor(0.5, dtype=DTYPE)
averaged = float(flow_velocity(x0, x1, bridge_sample(x0, x1, half, stratified), half).mean())
symmetry = worked[0.25]["flow"] + worked[0.75]["flow"] - 2.0 * (X1 - X0)
if loud:
print(f"\nOn {len(grid)} interior times")
print(f" largest gap between sigma_t grad log p_t and -z: {scaled_gap:.2e}")
print(f" u^o averaged over the draw at t = 1/2: {averaged:.4f}, "
f"against the straight-line velocity {X1 - X0:.0f}")
print(f" the two flow targets about the straight-line velocity: {symmetry:.2e}")
print(f" loss at the exact targets {loss_exact:.2e}, at a zero network {loss_zero:.4f}")
x = torch.tensor([0.0, 2.0], dtype=DTYPE) # The batch of Section 7.8.1.
y = torch.tensor([1.0, 3.0], dtype=DTYPE)
uniform = torch.full((2,), 0.5, dtype=DTYPE)
C = (x[:, None] - y[None, :]) ** 2
plan = entropic_plan(x, y, uniform, uniform)
ordered_cost = float(C[0, 0] + C[1, 1]) # 1 + 1, the arrows of the exact pairing.
crossed_cost = float(C[0, 1] + C[1, 0]) # 9 + 1, the arrows of the crossed pairing.
if loud:
print(f"\nThe two-point batch at epsilon = {EPSILON:g}")
print(f" exact transport pairs 0 with 1 and 2 with 3 at cost {ordered_cost:.0f}, "
f"against {crossed_cost:.0f} for the crossed pairing")
print(f" entropic plan at strength {2 * EPSILON:g}: "
f"{[[round(float(z_), 6) for z_ in row] for row in plan]}")
print(f" mass on the crossed pairing {float(plan[0, 1] + plan[1, 0]):.6f}")
report = {
"m_t at t = 0.25": worked[0.25]["mean"],
"sigma_t at t = 0.25": worked[0.25]["std"],
"bridge sample at t = 0.25": worked[0.25]["sample"],
"u_t at t = 0.25": worked[0.25]["drift"],
"grad log p_t at t = 0.25": worked[0.25]["score"],
"u^o_t at t = 0.25": worked[0.25]["flow"],
"coefficient at t = 0.25": worked[0.25]["coefficient"],
"sigma_t grad log p_t at t = 0.25": worked[0.25]["scaled"],
"bridge sample at t = 0.75": worked[0.75]["sample"],
"u_t at t = 0.75": worked[0.75]["drift"],
"u^o_t at t = 0.75": worked[0.75]["flow"],
"coefficient at t = 0.75": worked[0.75]["coefficient"],
"the two targets, sum minus 2 u": symmetry,
"largest gap of the scaled score": scaled_gap,
"u^o averaged over the draw": averaged,
"loss at the exact targets": loss_exact,
"loss at a zero network": loss_zero,
"sampler drift at g^2 = epsilon": float(sampler_drift(
torch.tensor(worked[0.25]["flow"], dtype=DTYPE),
torch.tensor(worked[0.25]["score"], dtype=DTYPE), EPSILON ** 0.5)),
"sampler drift at g = 0": worked[0.25]["flow"],
"ordered pairing of the batch": ordered_cost,
"crossed pairing of the batch": crossed_cost,
"entropic plan, ordered mass": float(plan[0, 0] + plan[1, 1]),
"entropic plan, crossed mass": float(plan[0, 1] + plan[1, 0]),
}
print_report("Simulation-Free Score and Flow Matching, Chapter 7 Section 7.19.3", report)
save_run(args.out, config=vars(args), report=report) # Written only with --out.
if __name__ == "__main__":
main()
|