Download lecture_3/samplers.py from ChatterjeeLab/CIS6270: direct link, hf CLI and curl.
- Browser
- Download file 11.8 kB
-
https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_3/samplers.py
- Command line
-
hf download hf://ChatterjeeLab/CIS6270/lecture_3/samplers.py
-
curl -L -o samplers.py https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_3/samplers.py
11.8 kB
| """Three samplers reading one score. | |
| Once we have a score, we can generate in three ways. Ancestral DDPM sampling comes from | |
| discretizing the reverse-time equation on a grid, and it injects fresh noise at every step. The | |
| probability-flow ODE comes from rewriting the Laplacian of the Fokker-Planck equation as a | |
| divergence, and it reproduces the same marginals with no noise at all. Langevin dynamics holds one | |
| density fixed and runs its score as a drift, so it needs a ladder of noise levels to cross between | |
| modes. We target the four-Gaussian mixture here, and its corrupted score is a closed form, so all | |
| three run on the same exact field. | |
| In Chapter 3 t = 0 holds a clean sample and corruption grows with t. To generate we run the same | |
| interval backward, from t = 1 to t = 0. | |
| Chapter 3, Sections 3.22, 3.26, 3.32, 3.34 and 3.35 of the notes. Prints the mode coverage of each | |
| sampler, the number of score evaluations each spends, and the chapter's one backward step from | |
| x = 1.2 under both dynamics. | |
| python lecture_3/samplers.py # 200 steps for each of the three | |
| python lecture_3/samplers.py --steps 20 # a coarse grid, for a first look | |
| """ | |
| # %% 01. Imports and configuration | |
| import math | |
| import torch | |
| from cis6270.data import GAUSSIAN_CENTERS, GAUSSIAN_STD | |
| from cis6270.metrics import mean_log_density, mode_report | |
| from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed | |
| BETA_MIN = 0.1 # beta(0) of the linear variance-preserving schedule. | |
| BETA_MAX = 20.0 # beta(1) of the same schedule. | |
| LADDER_START = 0.7 # Noisiest level of the Langevin ladder, in diffusion time. | |
| LADDER_END = 0.01 # Least noisy level of the ladder. | |
| # %% 02. The target and its exact score at any noise level | |
| # | |
| # If we corrupt a Gaussian mixture by X_t = a_t X_0 + sigma_t eps we get a Gaussian mixture with | |
| # means a_t mu_m and the common variance a_t^2 s^2 + sigma_t^2, so the score of p_t is a closed | |
| # form at every level. We count every call, since the cost of a sampler is the number of | |
| # evaluations a trained network would have served. | |
| def beta(t: float) -> float: | |
| """The noising rate beta(t) = 0.1 + 19.9 t of the linear schedule.""" | |
| return BETA_MIN + (BETA_MAX - BETA_MIN) * t | |
| def signal(t: float) -> float: | |
| """a_t = exp(-B(t)/2) with B(t) = 0.1 t + 9.95 t^2, the accumulated noising.""" | |
| return math.exp(-(BETA_MIN * t + 0.5 * (BETA_MAX - BETA_MIN) * t ** 2) / 2.0) | |
| class ExactScore: | |
| """The score of p_t at a given signal coefficient, with a counter on the calls.""" | |
| def __init__(self): | |
| self.calls = 0 # One call stands for one network evaluation. | |
| def __call__(self, x: torch.Tensor, a_t: float) -> torch.Tensor: | |
| """x: [B, 2], a_t the signal coefficient -> [B, 2].""" | |
| self.calls += 1 | |
| variance = a_t ** 2 * GAUSSIAN_STD ** 2 + (1.0 - a_t ** 2) # Width of p_t. | |
| offset = x[:, None, :] - a_t * GAUSSIAN_CENTERS.to(x)[None, :, :] # [B, 4, 2]. | |
| log_weight = -(offset ** 2).sum(dim=2) / (2.0 * variance) # [B, 4] equal mixture weights. | |
| responsibility = log_weight.softmax(dim=1)[:, :, None] # [B, 4, 1]. | |
| return (responsibility * (-offset / variance)).sum(dim=1) # [B, 2]. | |
| def make_schedule(num_steps: int) -> dict: | |
| """The DDPM tensors on a grid of K steps, read off the continuous schedule. | |
| We set abar_k = a_{t_k}^2 at t_k = k/K and recover alpha_k = abar_k / abar_{k-1}. These match | |
| the continuous marginals at every grid point. At K = 1000 the betas agree with the linear | |
| sequence from 1e-4 to 0.02 to the accuracy of log(1 - beta) = -beta. At a coarse K every beta_k | |
| stays below one, and beta(t_k)/K can rise above one there. | |
| """ | |
| grid = (torch.arange(0, num_steps + 1).to(torch.get_default_dtype()) | |
| / num_steps) # [K+1] the times t_k, with t_0 = 0. | |
| alphas_cumprod = torch.exp( | |
| -(BETA_MIN * grid + 0.5 * (BETA_MAX - BETA_MIN) * grid ** 2) | |
| ) # [K+1] abar_k, with abar_0 = 1 at the clean state. | |
| alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]]) # abar_{k-1}. | |
| alphas = alphas_cumprod / alphas_cumprod_prev # One-step retention alpha_k. | |
| betas = 1.0 - alphas # Per-step variance beta_k, zero at the placeholder index 0. | |
| return { | |
| "betas": betas, "alphas": alphas, "alphas_cumprod": alphas_cumprod, | |
| "posterior_variance": betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod), | |
| } | |
| # %% 03. Ancestral DDPM sampling | |
| # | |
| # x_{k-1} = mu_theta(x_k, k) + sqrt(beta-tilde_k) z with the reverse mean built from the noise | |
| # prediction, eps = -sigma_k s_k. Fresh noise enters at every step except the last. | |
| def ancestral_sample(score, schedule, num_points: int, generator=None) -> torch.Tensor: | |
| """Run the reverse chain from X_K ~ N(0, I) down to X_0. Returns [N, 2].""" | |
| num_steps = len(schedule["betas"]) - 1 | |
| x = torch.randn(num_points, 2, generator=generator) # [N, 2] the terminal Gaussian. | |
| for step in range(num_steps, 0, -1): # k = K down to k = 1. | |
| abar_k = float(schedule["alphas_cumprod"][step]) | |
| sigma_k = math.sqrt(1.0 - abar_k) # Corruption standard deviation at this step. | |
| eps_theta = -sigma_k * score(x, math.sqrt(abar_k)) # [N, 2] the noise the score implies. | |
| mean = (x - float(schedule["betas"][step]) / sigma_k * eps_theta) / math.sqrt( | |
| float(schedule["alphas"][step]) | |
| ) # Reverse mean mu_theta. | |
| if step > 1: # No noise at the final step. | |
| variance = float(schedule["posterior_variance"][step]) | |
| x = mean + math.sqrt(variance) * torch.randn(x.shape, generator=generator) | |
| else: | |
| x = mean | |
| return x # [N, 2]. | |
| # %% 04. The probability-flow ODE | |
| # | |
| # v_t(x) = -(1/2) beta(t) [x + s_t(x)] has the same marginals with no Brownian increment, and we | |
| # integrate it backward, x_{t - dt} = x_t - dt v_t(x_t). Once we draw X_T the whole trajectory is | |
| # fixed. | |
| def probability_flow_sample(score, num_points: int, steps: int, generator=None) -> torch.Tensor: | |
| """Euler integration of the probability-flow ODE from t = 1 to t = 0. Returns [N, 2].""" | |
| x = torch.randn(num_points, 2, generator=generator) # [N, 2] the terminal Gaussian. | |
| dt = 1.0 / steps # Positive step, taken in the direction of decreasing t. | |
| for step in range(steps, 0, -1): | |
| t = step * dt | |
| a_t = signal(t) | |
| velocity = -0.5 * beta(t) * (x + score(x, a_t)) # [N, 2] forward-time velocity. | |
| x = x - dt * velocity # One backward Euler step. | |
| return x # [N, 2]. | |
| # %% 05. Langevin dynamics and the noise ladder | |
| # | |
| # dX = s(X) d_tau + sqrt(2) dW holds p stationary at one fixed density, and the Euler-Maruyama | |
| # update is x' = x + eta s(x) + sqrt(2 eta) z. A chain run at the data density alone crosses | |
| # between modes slowly, so we visit the levels in decreasing order of corruption, with a few | |
| # iterations at each. | |
| def annealed_langevin_sample(score, num_points, levels, per_level, scale, generator=None): | |
| """Run ``per_level`` Langevin iterations at each level of the ladder. Returns [N, 2].""" | |
| x = torch.randn(num_points, 2, generator=generator) # [N, 2] a Gaussian start. | |
| for t in levels: | |
| a_t = signal(t) | |
| width = a_t ** 2 * GAUSSIAN_STD ** 2 + (1.0 - a_t ** 2) # Width of p_t at this level. | |
| eta = scale * width # Step size tied to the level, large when the density is broad. | |
| for _ in range(per_level): | |
| noise = torch.randn(x.shape, generator=generator) # [N, 2] fresh at every iteration. | |
| x = x + eta * score(x, a_t) + math.sqrt(2.0 * eta) * noise | |
| return x # [N, 2] at the least corrupted level of the ladder. | |
| def langevin_step(x: float, score_value: float, eta: float, z: float) -> float: | |
| """One Langevin update in one coordinate, the chapter's x + eta s(x) + sqrt(2 eta) z.""" | |
| return x + eta * score_value + math.sqrt(2.0 * eta) * z | |
| # %% 06. One backward step under both dynamics | |
| # | |
| # The score enters the reverse stochastic equation with the full g(t)^2 and the deterministic | |
| # velocity with half of it, and that factor of two is the only structural difference between the | |
| # two samplers. | |
| def one_step_comparison(x_t, beta_t, sigma_t, eps_theta, dt) -> dict: | |
| """The chapter's state advanced one step by the ODE and by the reverse-time equation.""" | |
| score_value = -eps_theta / sigma_t # The score the noise prediction implies. | |
| velocity = -0.5 * beta_t * (x_t + score_value) # Probability-flow velocity. | |
| reverse_drift = -0.5 * beta_t * x_t - beta_t * score_value # f - g^2 s. | |
| return { | |
| "score": score_value, "velocity": velocity, "ode_next": x_t - dt * velocity, | |
| "sde_drift": reverse_drift, "sde_mean": x_t - dt * reverse_drift, | |
| "sde_noise_std": math.sqrt(beta_t * dt), | |
| } | |
| # %% 07. The report | |
| def main() -> None: | |
| parser = common_parser(__doc__.split("\n")[0]) | |
| parser.add_argument("--levels", type=int, default=20, help="rungs of the Langevin ladder") | |
| parser.add_argument("--langevin-scale", type=float, default=0.15, help="step size over width") | |
| parser.set_defaults(steps=200, samples=2048) | |
| args = parser.parse_args() | |
| set_seed(args.seed) | |
| device = pick_device(args.device) # Two coordinates per point; a CPU is enough. | |
| results, counts = {}, {} | |
| generator = torch.Generator().manual_seed(args.seed) | |
| score = ExactScore() | |
| results["ancestral DDPM"] = ancestral_sample( | |
| score, make_schedule(args.steps), args.samples, generator, | |
| ) | |
| counts["ancestral DDPM"] = score.calls | |
| generator = torch.Generator().manual_seed(args.seed) # The same starting noise for each. | |
| score = ExactScore() | |
| results["probability flow"] = probability_flow_sample( | |
| score, args.samples, args.steps, generator, | |
| ) | |
| counts["probability flow"] = score.calls | |
| generator = torch.Generator().manual_seed(args.seed) | |
| score = ExactScore() | |
| per_level = max(args.steps // args.levels, 1) # Spend the same budget as the other two. | |
| ladder = [LADDER_START + (LADDER_END - LADDER_START) * i / (args.levels - 1) | |
| for i in range(args.levels)] # Decreasing corruption. | |
| results["annealed Langevin"] = annealed_langevin_sample( | |
| score, args.samples, ladder, per_level, args.langevin_scale, generator, | |
| ) | |
| counts["annealed Langevin"] = score.calls | |
| report = {} | |
| for name, points in results.items(): | |
| measured = mode_report(points) | |
| report[f"{name}: score evaluations"] = counts[name] | |
| report[f"{name}: modes found"] = measured["modes_found"] | |
| report[f"{name}: mode entropy"] = measured["mode_entropy"] # log 4 = 1.386 when even. | |
| report[f"{name}: mode fractions"] = measured["mode_fractions"] | |
| report[f"{name}: distance to mode"] = measured["nearest_distance"] # 0.44 on the target. | |
| report[f"{name}: mean log density"] = mean_log_density(points) | |
| step = one_step_comparison(1.2, 2.0, 0.5, 0.8, 0.1) | |
| report["score at the worked state"] = step["score"] # -1.6. | |
| report["probability-flow velocity"] = step["velocity"] # 0.4. | |
| report["ODE next state"] = step["ode_next"] # 1.16. | |
| report["reverse-time drift"] = step["sde_drift"] # 2.0. | |
| report["reverse-time mean"] = step["sde_mean"] # 1.0. | |
| report["reverse-time noise std"] = step["sde_noise_std"] # 0.4472. | |
| report["Langevin step from x=2"] = langevin_step(2.0, -2.0, 0.01, 0.5) # 2.0507. | |
| if not args.quiet: | |
| print(f"device={device} points={args.samples} steps={args.steps} seed={args.seed}") | |
| print_report("Three samplers on one exact score", report) | |
| save_run(args.out, config=vars(args), report=report) | |
| # %% 08. Run the script | |
| if __name__ == "__main__": | |
| main() | |