Score SDE 三种随机微分方程代码解读

定义SDE类

定义了7个子函数

T: End time of the SDE.
sde:
marginal_prob: Parameters to determine the marginal distribution of the SDE, pt(x)p_t(x)pt(x).
prior_sampling: Generate one sample from the prior distribution, pT(x)p_T(x)pT(x).
prior_logp: Compute log-density of the prior distribution.
discretize: Discretize the SDE in the form: x_{i+1} = x_i + f_i(x_i) + G_i z_i.
Useful for reverse diffusion sampling and probabiliy flow sampling.
Defaults to Euler-Maruyama discretization.
reverse: Create the reverse-time SDE/ODE.

reverse

class SDE(abc.ABC):"""SDE abstract class. Functions are designed for a mini-batch of inputs."""def __init__(self, N):"""Construct an SDE.Args:N: number of discretization time steps."""super().__init__()self.N = N@property@abc.abstractmethoddef T(self):"""End time of the SDE."""pass@abc.abstractmethoddef sde(self, x, t):pass@abc.abstractmethoddef marginal_prob(self, x, t):"""Parameters to determine the marginal distribution of the SDE, $p_t(x)$."""pass@abc.abstractmethoddef prior_sampling(self, shape):"""Generate one sample from the prior distribution, $p_T(x)$."""pass@abc.abstractmethoddef prior_logp(self, z):"""Compute log-density of the prior distribution.Useful for computing the log-likelihood via probability flow ODE.Args:z: latent codeReturns:log probability density"""passdef discretize(self, x, t):"""Discretize the SDE in the form: x_{i+1} = x_i + f_i(x_i) + G_i z_i.Useful for reverse diffusion sampling and probabiliy flow sampling.Defaults to Euler-Maruyama discretization.Args:x: a torch tensort: a torch float representing the time step (from 0 to `self.T`)Returns:f, G"""dt = 1 / self.Ndrift, diffusion = self.sde(x, t)f = drift * dtG = diffusion * torch.sqrt(torch.tensor(dt, device=t.device))return f, Gdef reverse(self, score_fn, probability_flow=False):"""Create the reverse-time SDE/ODE.Args:score_fn: A time-dependent score-based model that takes x and t and returns the score.probability_flow: If `True`, create the reverse-time ODE used for probability flow sampling."""N = self.NT = self.Tsde_fn = self.sdediscretize_fn = self.discretize

定义reverse-time SDE类

    # Build the class for reverse-time SDE.class RSDE(self.__class__):def __init__(self):self.N = Nself.probability_flow = probability_flow@propertydef T(self):return Tdef sde(self, x, t):"""Create the drift and diffusion functions for the reverse SDE/ODE."""drift, diffusion = sde_fn(x, t)score = score_fn(x, t)drift = drift - diffusion[:, None, None, None] ** 2 * score * (0.5 if self.probability_flow else 1.)# Set the diffusion function to zero for ODEs.diffusion = 0. if self.probability_flow else diffusionreturn drift, diffusiondef discretize(self, x, t):"""Create discretized iteration rules for the reverse diffusion sampler."""f, G = discretize_fn(x, t)rev_f = f - G[:, None, None, None] ** 2 * score_fn(x, t) * (0.5 if self.probability_flow else 1.)rev_G = torch.zeros_like(G) if self.probability_flow else Greturn rev_f, rev_Greturn RSDE()

定义VPSDE类

class VPSDE(SDE):def __init__(self, beta_min=0.1, beta_max=20, N=1000):"""Construct a Variance Preserving SDE.Args:beta_min: value of beta(0)beta_max: value of beta(1)N: number of discretization steps"""super().__init__(N)self.beta_0 = beta_minself.beta_1 = beta_maxself.N = Nself.discrete_betas = torch.linspace(beta_min / N, beta_max / N, N)self.alphas = 1. - self.discrete_betasself.alphas_cumprod = torch.cumprod(self.alphas, dim=0)self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)self.sqrt_1m_alphas_cumprod = torch.sqrt(1. - self.alphas_cumprod)@propertydef T(self):return 1def sde(self, x, t):beta_t = self.beta_0 + t * (self.beta_1 - self.beta_0)drift = -0.5 * beta_t[:, None, None, None] * xdiffusion = torch.sqrt(beta_t)return drift, diffusiondef marginal_prob(self, x, t):log_mean_coeff = -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0mean = torch.exp(log_mean_coeff[:, None, None, None]) * xstd = torch.sqrt(1. - torch.exp(2. * log_mean_coeff))return mean, stddef prior_sampling(self, shape):return torch.randn(*shape)def prior_logp(self, z):shape = z.shapeN = np.prod(shape[1:])logps = -N / 2. * np.log(2 * np.pi) - torch.sum(z ** 2, dim=(1, 2, 3)) / 2.return logpsdef discretize(self, x, t):"""DDPM discretization."""timestep = (t * (self.N - 1) / self.T).long()beta = self.discrete_betas.to(x.device)[timestep]alpha = self.alphas.to(x.device)[timestep]sqrt_beta = torch.sqrt(beta)f = torch.sqrt(alpha)[:, None, None, None] * x - xG = sqrt_betareturn f, G

定义subVPSDE类

class subVPSDE(SDE):def __init__(self, beta_min=0.1, beta_max=20, N=1000):"""Construct the sub-VP SDE that excels at likelihoods.Args:beta_min: value of beta(0)beta_max: value of beta(1)N: number of discretization steps"""super().__init__(N)self.beta_0 = beta_minself.beta_1 = beta_maxself.N = N@propertydef T(self):return 1def sde(self, x, t):beta_t = self.beta_0 + t * (self.beta_1 - self.beta_0)drift = -0.5 * beta_t[:, None, None, None] * xdiscount = 1. - torch.exp(-2 * self.beta_0 * t - (self.beta_1 - self.beta_0) * t ** 2)diffusion = torch.sqrt(beta_t * discount)return drift, diffusion#边际概率函数，返回值是边际概率的mean和stddef marginal_prob(self, x, t):log_mean_coeff = -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0mean = torch.exp(log_mean_coeff)[:, None, None, None] * xstd = 1 - torch.exp(2. * log_mean_coeff)return mean, stddef prior_sampling(self, shape):return torch.randn(*shape)def prior_logp(self, z):shape = z.shapeN = np.prod(shape[1:])return -N / 2. * np.log(2 * np.pi) - torch.sum(z ** 2, dim=(1, 2, 3)) / 2.

定义VESDE类

class VESDE(SDE):def __init__(self, sigma_min=0.01, sigma_max=50, N=1000):"""Construct a Variance Exploding SDE.Args:sigma_min: smallest sigma.sigma_max: largest sigma.N: number of discretization steps"""super().__init__(N)self.sigma_min = sigma_minself.sigma_max = sigma_maxself.discrete_sigmas = torch.exp(torch.linspace(np.log(self.sigma_min), np.log(self.sigma_max), N))self.N = N@propertydef T(self):return 1def sde(self, x, t):sigma = self.sigma_min * (self.sigma_max / self.sigma_min) ** tdrift = torch.zeros_like(x)diffusion = sigma * torch.sqrt(torch.tensor(2 * (np.log(self.sigma_max) - np.log(self.sigma_min)),device=t.device))return drift, diffusiondef marginal_prob(self, x, t):std = self.sigma_min * (self.sigma_max / self.sigma_min) ** tmean = xreturn mean, stddef prior_sampling(self, shape):return torch.randn(*shape) * self.sigma_maxdef prior_logp(self, z):shape = z.shapeN = np.prod(shape[1:])return -N / 2. * np.log(2 * np.pi * self.sigma_max ** 2) - torch.sum(z ** 2, dim=(1, 2, 3)) / (2 * self.sigma_max ** 2)def discretize(self, x, t):"""SMLD(NCSN) discretization."""timestep = (t * (self.N - 1) / self.T).long()sigma = self.discrete_sigmas.to(t.device)[timestep]adjacent_sigma = torch.where(timestep == 0, torch.zeros_like(t),self.discrete_sigmas[timestep - 1].to(t.device))f = torch.zeros_like(x)G = torch.sqrt(sigma ** 2 - adjacent_sigma ** 2)return f, G