Blame for lecture_3/guidance.py

"""Steering a sampler by adding one gradient to the score.

Bayes' rule splits the conditional score into the unconditional score plus the gradient of the log
likelihood of the condition, so every form of control in this chapter is one extra term in the
same place. A classifier supplies that term directly, a conditional and an unconditional
prediction of one network supply it as their difference, and a differentiable reward supplies it
as its own gradient. When several objectives apply at once, we combine them through weights on
the simplex. We target the four-Gaussian mixture with one label per mode, where every score is a
closed form, so the guidance weight appears directly in the mode fractions.

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.37 through 3.45 of the notes. Prints the mode fractions at several guidance
weights, the chapter's worked guided score and guided backward step, the two candidates under two
weightings, and the schedule value at t = 0.6.

    python lecture_3/guidance.py              # 200 reverse steps per setting
    python lecture_3/guidance.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 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.
TARGET_LABEL = 0  # The mode at (2, 2), and every guided run is steered toward it.
REWARD_SCALE = 2.0  # Distance from the origin to a mode, used to standardize the two rewards.


# %% 02. The target, its labels and the exact scores
#
# Each mode has its own label, so p_t(x | c) is a single Gaussian with mean a_t mu_c and
# variance a_t^2 s^2 + sigma_t^2, and p_t(x) is the equal mixture of the four. Both scores are
# closed forms.


def schedule_tensors(num_steps: int) -> dict:
    """The DDPM tensors on a grid of K steps, read off abar_k = exp(-B(t_k))."""
    grid = (torch.arange(0, num_steps + 1).to(torch.get_default_dtype())
            / num_steps)  # [K+1] the times t_k.
    alphas_cumprod = torch.exp(-(BETA_MIN * grid + 0.5 * (BETA_MAX - BETA_MIN) * grid ** 2))
    previous = torch.cat([torch.ones(1), alphas_cumprod[:-1]])  # abar_{k-1}.
    alphas = alphas_cumprod / previous  # One-step retention alpha_k.
    betas = 1.0 - alphas  # Per-step variance beta_k.
    return {"betas": betas, "alphas": alphas, "alphas_cumprod": alphas_cumprod,
            "posterior_variance": betas * (1.0 - previous) / (1.0 - alphas_cumprod)}


def component_log_weights(x: torch.Tensor, a_t: float):
    """Log of each mode's unnormalized contribution at x, and the offsets to the four means."""
    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].
    return -(offset ** 2).sum(dim=2) / (2.0 * variance), offset, variance


def unconditional_score(x: torch.Tensor, a_t: float) -> torch.Tensor:
    """The score of the equal mixture at level a_t. [B, 2] -> [B, 2]."""
    log_weight, offset, variance = component_log_weights(x, a_t)
    responsibility = log_weight.softmax(dim=1)[:, :, None]  # [B, 4, 1] posterior over the modes.
    return (responsibility * (-offset / variance)).sum(dim=1)  # [B, 2].


def conditional_score(x: torch.Tensor, a_t: float, label: int) -> torch.Tensor:
    """The score of p_t(. | c), a single Gaussian centered on a_t mu_c. [B, 2] -> [B, 2]."""
    _, offset, variance = component_log_weights(x, a_t)
    return -offset[:, label, :] / variance  # [B, 2].


# %% 03. Conditioning the score and the guidance scale
#
# s_t(x | c) = s_t(x) + grad log p_t(c | x), and if we attach a scale gamma to the second term we
# sample the tilted density p_t(x) p_t(c | x)^gamma. At gamma = 0 generation is unconditional, at
# gamma = 1 it follows Bayes' rule, and at gamma > 1 the condition carries more weight.


def classifier_gradient(x: torch.Tensor, a_t: float, label: int) -> torch.Tensor:
    """grad log p_t(c | x), taken as the difference of the two scores. [B, 2] -> [B, 2]."""
    return conditional_score(x, a_t, label) - unconditional_score(x, a_t)


def guided_score(score: torch.Tensor, gradient: torch.Tensor, gamma: float) -> torch.Tensor:
    """s_theta + gamma grad log p_phi(c | x), the guided score of the notes."""
    return score + gamma * gradient


def guided_noise(eps: torch.Tensor, gradient: torch.Tensor, gamma: float, sigma_k: float):
    """eps_theta - gamma sigma_t grad log p_phi(c | x), the same guidance in the noise."""
    return eps - gamma * sigma_k * gradient


def classifier_free_field(unconditional, conditional, w: float):
    """F_uncond + w (F_cond - F_uncond), with no separate classifier."""
    return unconditional + w * (conditional - unconditional)


# %% 04. Reward guidance and several objectives at once
#
# A reward is a differentiable scalar that is larger for more desirable states, and exponential
# reweighting turns it into a density whose score is s_t(x) + eta grad R_t(x). With M objectives
# we clamp the weights to be nonnegative and divide by their sum, the simplex condition of the
# notes, and the gradient is linear in them.


def reward_model(z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """Two standardized properties of the state, its reach to the right and its reach upward.

    The notes train a small network on the latent and the time here. Two coordinates of a planar
    state need no network, so we read the properties off directly and keep the same signature.
    """
    return z / REWARD_SCALE  # [B, 2] one standardized property per coordinate.


def normalize_weights(lambdas, device):
    """Clamp the weights to be nonnegative and divide by their sum, the simplex condition."""
    w = torch.as_tensor(lambdas, dtype=torch.get_default_dtype(), device=device)
    w = w.clamp_min(0.0)
    return w / w.sum().clamp_min(1e-8)


def weighted_reward(reward_model, z, t, lambdas):
    """R_lambda(z) = sum_m lambda_m r_m(z), one scalar per batch item."""
    r_hat = reward_model(z, t)                       # [B, M]
    return (r_hat * normalize_weights(lambdas, z.device)).sum(dim=1)


def reward_gradient(reward_model, z, t, lambdas):
    """grad_z R_lambda(z), taken inside a no-grad sampler. [B, d] -> [B, d]."""
    with torch.enable_grad():
        state = z.detach().requires_grad_(True)
        R = weighted_reward(reward_model, state, t, lambdas).sum()
        grad, = torch.autograd.grad(R, state)
    return grad.detach()


def scalarize(values, lambdas) -> float:
    """The weighted sum of normalized objectives, R_lambda(x), on one candidate."""
    weights = normalize_weights(lambdas, torch.device("cpu"))
    return float((torch.as_tensor(values, dtype=torch.get_default_dtype()) * weights).sum())


# %% 05. Scheduled guidance
#
# A property predictor varies in reliability across the noise levels, so we let the strength
# depend on time, gamma(t) = gamma_max zeta(t). The schedule zeta(t) = 4 t (1 - t) vanishes
# at both ends of the interval and peaks at one in the middle.


def schedule_value(t: float) -> float:
    """zeta(t) = 4 t (1 - t), between zero and one."""
    return 4.0 * t * (1.0 - t)


# %% 06. The guided sampler
#
# The reverse step is the one of Section 3.26 with the guided noise prediction in place of the
# raw one. We apply classifier-free guidance first and the reward term second, the order the
# notes implements.


@torch.no_grad()
def sample_diffusion(schedule, num_points, label, w=0.0, eta=0.0, lambdas=(1.0, 0.0),
                     scheduled=False, generator=None) -> torch.Tensor:
    """Draw points under classifier-free guidance toward ``label`` and reward steering."""
    num_steps = len(schedule["betas"]) - 1
    z = 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.
        t = step / num_steps  # Normalized time, the input the network would receive.
        abar = float(schedule["alphas_cumprod"][step])
        a_k, sigma_k = math.sqrt(abar), math.sqrt(1.0 - abar)
        eps = -sigma_k * unconditional_score(z, a_k)  # [N, 2] unconditional noise.
        if w > 0.0:
            conditional = -sigma_k * conditional_score(z, a_k, label)  # [N, 2].
            strength = w * (schedule_value(t) if scheduled else 1.0)  # gamma(t) = gamma_max zeta.
            eps = classifier_free_field(eps, conditional, strength)
        if eta > 0.0:
            grad_R = reward_gradient(reward_model, z, torch.full((num_points,), t), lambdas)
            eps = eps - eta * sigma_k * grad_R  # Reward guidance in the noise parameterization.
        beta_k, alpha_k = float(schedule["betas"][step]), float(schedule["alphas"][step])
        mean = (z - beta_k / sigma_k * eps) / math.sqrt(alpha_k)  # Reverse mean mu_theta.
        if step > 1:
            variance = float(schedule["posterior_variance"][step])
            z = mean + math.sqrt(variance) * torch.randn(z.shape, generator=generator)
        else:
            z = mean
    return z  # [N, 2].


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    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.
    schedule = schedule_tensors(args.steps)

    def run(**settings) -> list:
        generator = torch.Generator().manual_seed(args.seed)  # The same starting noise each time.
        points = sample_diffusion(schedule, args.samples, TARGET_LABEL, generator=generator,
                                  **settings)
        return mode_report(points)["mode_fractions"]

    # The conditional score here is exact, so at w = 1 the condition already pins the mode and
    # the fractions saturate. The graded part of the sweep therefore sits below one.
    report = {}
    for w in (0.0, 0.1, 0.25, 0.5, 1.0):  # The target mode is the first entry of each list.
        report[f"mode fractions at w={w}"] = run(w=w)
    report["mode fractions, w=1 scheduled"] = run(w=1.0, scheduled=True)
    for eta in (0.5, 2.0):  # Reward steering toward the two right-hand modes.
        report[f"mode fractions at eta={eta}, lambda=(1,0)"] = run(eta=eta, lambdas=(1.0, 0.0))
    report["mode fractions at eta=2, lambda=(0.7,0.3)"] = run(eta=2.0, lambdas=(0.7, 0.3))

    # The two routes to one guided noise prediction, on a batch at a middle level.
    probe = torch.randn(256, 2, generator=torch.Generator().manual_seed(args.seed + 1))
    a_t, sigma_t = 0.8, 0.6
    by_classifier = guided_noise(-sigma_t * unconditional_score(probe, a_t),
                                 classifier_gradient(probe, a_t, TARGET_LABEL), 2.0, sigma_t)
    by_difference = classifier_free_field(-sigma_t * unconditional_score(probe, a_t),
                                          -sigma_t * conditional_score(probe, a_t, TARGET_LABEL),
                                          2.0)
    gap = float((by_classifier - by_difference).abs().max())
    report["classifier and classifier-free gap"] = gap

    # The chapter's worked numbers, all of them arithmetic.
    score, gradient = torch.tensor([[-1.0, -0.1]]), torch.tensor([[1.3, 1.1]])
    for gamma in (1.0, 2.0):
        direction = guided_score(score, gradient, gamma)[0]
        report[f"guided direction at gamma={gamma}"] = [round(float(v), 4) for v in direction]
        report[f"its length and angle at gamma={gamma}"] = [
            round(float(direction.norm()), 4),
            round(math.degrees(math.atan2(float(direction[1]), float(direction[0]))), 1),
        ]
    cfg_eps = float(classifier_free_field(torch.tensor(0.50), torch.tensor(0.10), 2.0))
    report["classifier-free eps at w=2"] = cfg_eps  # -0.30, past the conditional 0.10.
    report["its implied score"] = -cfg_eps / 0.5  # 0.6, at sigma_t = 0.5.

    # One guided backward step from the state of the deterministic sampler, through both
    # parameterizations. A disagreement between them is a missing sigma_t or a missing one half.
    x_t, beta_t, sigma_t, eps_theta = 1.2, 2.0, 0.5, 0.8  # The state of Section 3.35.
    gradient, gamma, dt = 0.6, 2.0, 0.1  # The classifier gradient, the scale and the step.
    unguided = -0.5 * beta_t * (x_t - eps_theta / sigma_t)  # 0.4, the unguided velocity.
    guided = unguided - 0.5 * gamma * beta_t * gradient  # -0.8, half of gamma g^2 on the gradient.
    guided_eps = float(guided_noise(torch.tensor(eps_theta), torch.tensor(gradient),
                                    gamma, sigma_t))  # 0.2, the same guidance in the noise.
    report["guided velocity and next state"] = [round(guided, 4),
                                                round(x_t - dt * guided, 4)]  # -0.8 and 1.28.
    report["guided eps and its velocity"] = [
        round(guided_eps, 4),  # 0.2, with implied score -0.4.
        round(-0.5 * beta_t * (x_t - guided_eps / sigma_t), 4),  # -0.8, the same velocity.
    ]

    # Two candidates under two weightings, and two objectives that disagree.
    candidates = ((1.2, 0.2), (0.5, 1.1))  # Normalized binding and solubility of two peptides.
    for lambdas in ((0.7, 0.3), (0.3, 0.7)):
        report[f"candidate rewards, lambda={lambdas}"] = [round(scalarize(v, lambdas), 4)
                                                          for v in candidates]
    competing = 0.6 * torch.tensor([1.0, 0.8]) + 0.4 * torch.tensor([0.5, -1.2])
    report["competing gradients combined"] = [round(float(v), 4) for v in competing]  # (0.8, 0).
    rewards = torch.tensor([1.2, 0.8, 0.3])
    report["selection weights at eta=2"] = [round(float(p), 4)
                                            for p in (2.0 * rewards).softmax(0)]
    report["schedule value at t=0.6"] = schedule_value(0.6)  # 0.96.

    if not args.quiet:
        print(f"device={device} points={args.samples} steps={args.steps} "
              f"target mode={GAUSSIAN_CENTERS[TARGET_LABEL].tolist()}")
    print_report("Guidance on the four-Gaussian target", report)
    save_run(args.out, config=vars(args), report=report)


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