Blame for lecture_3/samplers.py

"""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.


@torch.no_grad()
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.


@torch.no_grad()
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.


@torch.no_grad()
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()