File size: 8,736 Bytes
f15d29e | 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 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 | """
Copyright 2020 The Google Research Authors.
Copyright (c) Microsoft Corporation.
Licensed under the MIT License.
Based on code from https://github.com/yang-song/score_sde_pytorch
which is released under Apache licence.
Abstract SDE classes, Reverse SDE, and VE/VP SDEs.
Key changes:
- Adapted to work on batched pytorch_geometric style data
- Added '...given_score' methods so that score for a composite
state can be calculated in single forward pass of a shared score model,
and the scores for different fields then forwarded to the different reverse SDEs.
"""
import abc
from typing import Callable, Optional, Protocol, Tuple, Union
import numpy as np
import torch
from torch_scatter import scatter_add
from ...diffusion.corruption.corruption import B, Corruption, maybe_expand
from ...diffusion.data.batched_data import BatchedData
class ScoreFunction(Protocol):
def __call__(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
) -> torch.Tensor:
"""Calculate score.
Args:
x: Samples at which the score should be calculated. Shape [num_nodes, ...]
t: Timestep for each sample. Shape [num_samples,]
batch_idx: Indicates which sample each row of x belongs to. Shape [num_nodes,]
"""
pass
class SDE(Corruption):
"""Corruption using a stochastic differential equation."""
@abc.abstractmethod
def sde(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: Optional[BatchedData] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Returns drift f and diffusion coefficient g such that dx = f * dt + g * sqrt(dt) * standard Gaussian"""
pass # drift: (nodes_per_sample * batch_size, num_features), diffusion (batch_size,)
@abc.abstractmethod
def marginal_prob(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: Optional[BatchedData] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Returns mean and standard deviation of the marginal distribution of the SDE, $p_t(x)$."""
pass # mean: (nodes_per_sample * batch_size, num_features), std: (nodes_per_sample * batch_size, 1)
def mean_coeff_and_std(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: Optional[BatchedData] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Returns mean coefficient and standard deviation of marginal distribution at time t."""
return self.marginal_prob(
torch.ones_like(x), t, batch_idx, batch
) # mean_coeff: same shape as x, std: same shape as x
def sample_marginal(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: Optional[BatchedData] = None,
) -> torch.Tensor:
"""Sample marginal for x(t) given x(0).
Returns:
sampled x(t)
"""
mean, std = self.marginal_prob(x=x, t=t, batch_idx=batch_idx, batch=batch)
z = torch.randn_like(x)
return mean + std * z
class BaseVPSDE(SDE):
"""Base class for variance-preserving SDEs of the form
dx = - 0.5 * beta_t * x * dt + sqrt(beta_t) * z * sqrt(dt)
where z is unit Gaussian noise, or equivalently
dx = - 0.5 * beta_t *x * dt + sqrt(beta_t) * dW
"""
@abc.abstractmethod
def beta(self, t: torch.Tensor) -> torch.Tensor: ...
@abc.abstractmethod
def _marginal_mean_coeff(self, t: torch.Tensor) -> torch.Tensor:
"""This should be implemented to compute exp(-0.5 * int_0^t beta(s) ds). See equation (29) of Song et al."""
...
@property
def T(self) -> float:
return 1.0
def marginal_prob(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: Optional[BatchedData] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
mean_coeff = self._marginal_mean_coeff(t)
mean = maybe_expand(mean_coeff, batch_idx, x) * x
std = maybe_expand(torch.sqrt(1.0 - mean_coeff**2), batch_idx, x)
return mean, std
def prior_sampling(
self,
shape: Union[torch.Size, Tuple],
conditioning_data: Optional[BatchedData] = None,
batch_idx: B = None,
) -> torch.Tensor:
return torch.randn(*shape)
def prior_logp(
self,
z: torch.Tensor,
batch_idx: B = None,
batch: Optional[BatchedData] = None,
) -> torch.Tensor:
return unit_gaussian_logp(z, batch_idx)
def sde(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: Optional[BatchedData] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
beta_t = self.beta(t)
drift = -0.5 * maybe_expand(beta_t, batch_idx, x) * x
diffusion = maybe_expand(torch.sqrt(beta_t), batch_idx, x)
return drift, diffusion
class VPSDE(BaseVPSDE):
def __init__(self, beta_min: float = 0.1, beta_max: float = 20):
"""Variance-preserving SDE with drift coefficient changing linearly over time."""
super().__init__()
self.beta_0 = beta_min
self.beta_1 = beta_max
def beta(self, t: torch.Tensor) -> torch.Tensor:
return self.beta_0 + t * (self.beta_1 - self.beta_0)
def _marginal_mean_coeff(self, t: torch.Tensor) -> torch.Tensor:
log_mean_coeff = -0.25 * t**2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
return torch.exp(log_mean_coeff)
def unit_gaussian_logp(z: torch.Tensor, batch_idx: B = None) -> torch.Tensor:
shape = z.shape
N = np.prod(shape[1:])
if batch_idx is None:
logps = -N / 2.0 * np.log(2 * np.pi) - torch.sum(z**2, dim=tuple(range(1, z.ndim))) / 2.0
else:
if z.ndim > 2:
raise NotImplementedError
logps = -N / 2.0 * np.log(2 * np.pi) - scatter_add(torch.sum(z**2, dim=1), batch_idx) / 2.0
return logps
class VESDE(SDE):
def __init__(self, sigma_min: float = 0.01, sigma_max: float = 50.0):
"""Construct a Variance Exploding SDE.
The marginal standard deviation grows exponentially from sigma_min to sigma_max.
Args:
sigma_min: smallest sigma.
sigma_max: largest sigma.
"""
super().__init__()
self.sigma_min = sigma_min
self.sigma_max = sigma_max
@property
def T(self) -> float:
return 1.0
def sde(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: Optional[BatchedData] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
sigma = self.sigma_min * (self.sigma_max / self.sigma_min) ** t
drift = torch.zeros_like(x)
diffusion = maybe_expand(
sigma
* torch.sqrt(
torch.tensor(2 * (np.log(self.sigma_max) - np.log(self.sigma_min)), device=t.device)
),
batch_idx,
x,
)
return drift, diffusion
def marginal_prob(
self,
x: torch.Tensor,
t: torch.Tensor,
batch_idx: B = None,
batch: Optional[BatchedData] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
std = maybe_expand(self.sigma_min * (self.sigma_max / self.sigma_min) ** t, batch_idx, x)
mean = x
return mean, std
def prior_sampling(
self,
shape: Union[torch.Size, Tuple],
conditioning_data: Optional[BatchedData] = None,
batch_idx: B = None,
) -> torch.Tensor:
return torch.randn(*shape) * self.sigma_max
def prior_logp(
self,
z: torch.Tensor,
batch_idx: B = None,
batch: Optional[BatchedData] = None,
) -> torch.Tensor:
shape = z.shape
N = np.prod(shape[1:])
if batch_idx is not None:
return -N / 2.0 * np.log(2 * np.pi * self.sigma_max**2) - scatter_add(
torch.sum(z**2, dim=1), batch_idx
) / (2 * self.sigma_max**2)
else:
return -N / 2.0 * np.log(2 * np.pi * self.sigma_max**2) - torch.sum(
z**2, dim=tuple(range(1, z.ndim))
) / (2 * self.sigma_max**2)
def check_score_fn_defined(score_fn: Optional[Callable], fn_name_given_score: str):
"""Check that a reverse SDE has a score_fn. Give a useful error message if not."""
if score_fn is None:
raise ValueError(
f"This reverse SDE does not know its score_fn. You must either a) pass a score_fn when you construct this reverse SDE or b) call {fn_name_given_score} instead."
)
|