File size: 3,890 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 | # Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import torch
from ...common.diffusion import corruption as sde_lib
from ...common.utils.data_utils import compute_lattice_polar_decomposition
from ...diffusion.corruption.corruption import Corruption, maybe_expand
from ...diffusion.data.batched_data import BatchedData
from ...diffusion.sampling import predictors_correctors as pc
from ...diffusion.sampling.predictors import AncestralSamplingPredictor
SampleAndMean = tuple[torch.Tensor, torch.Tensor]
class LatticeAncestralSamplingPredictor(AncestralSamplingPredictor):
@classmethod
def is_compatible(cls, corruption: Corruption) -> bool:
_super = super()
assert hasattr(_super, "is_compatible")
return _super.is_compatible(corruption) or isinstance(corruption, sde_lib.LatticeVPSDE)
def update_given_score(
self,
*,
x: torch.Tensor,
t: torch.Tensor,
dt: torch.Tensor,
batch_idx: torch.LongTensor,
score: torch.Tensor,
batch: BatchedData | None,
) -> SampleAndMean:
x_coeff, score_coeff, std = self._get_coeffs(
x=x,
t=t,
dt=dt,
batch_idx=batch_idx,
batch=batch,
)
# mean = (x + score * beta**2 - limit_mean)/(1-beta) + limit_mean
# <=> mean = x / (1-beta) + score * beta**2 / (1-beta) + limit_mean * (1 - 1/(1-beta))
# => mean_coeff = 1 - x_coeff = 1 - 1/(1-beta)
mean_coeff = 1 - x_coeff
# Sample random noise.
z = sde_lib.make_noise_symmetric_preserve_variance(torch.randn_like(x_coeff))
assert hasattr(self.corruption, "get_limit_mean") # mypy
mean = (
x_coeff * x
+ score_coeff * score
+ mean_coeff * self.corruption.get_limit_mean(x=x, batch=batch)
)
sample = mean + std * z
return sample, mean
# create a langevin corrector that accepts LatticeVPSDE
class LatticeLangevinDiffCorrector(pc.LangevinCorrector):
@classmethod
def is_compatible(cls, corruption: Corruption) -> bool:
_super = super()
assert hasattr(_super, "is_compatible")
return _super.is_compatible(corruption) or isinstance(corruption, sde_lib.LatticeVPSDE)
def step_given_score(
self,
*,
x: torch.Tensor,
batch_idx: torch.LongTensor | None,
score: torch.Tensor,
t: torch.Tensor,
dt: torch.Tensor,
) -> SampleAndMean:
assert isinstance(self.corruption, sde_lib.LatticeVPSDE)
alpha = self.get_alpha(t, dt=dt)
snr = self.snr
noise = torch.randn_like(x)
noise = sde_lib.make_noise_symmetric_preserve_variance(noise)
# [batch_size, ] or [num_atoms, ] if batch_idx is not None
grad_norm_square = torch.square(score).reshape(score.shape[0], -1).sum(dim=1)
noise_norm_square = torch.square(noise).reshape(noise.shape[0], -1).sum(dim=1)
# Average over items, leading to scalars.
grad_norm = grad_norm_square.sqrt().mean()
noise_norm = noise_norm_square.sqrt().mean()
# If gradient is zero (i.e., we are sampling from an improper distribution that's flat over the whole of R^n)
# the step_size blows up. Clip step_size to avoid this.
# The EGNN reports zero scores when there are no edges between nodes.
step_size = (snr * noise_norm / grad_norm) ** 2 * 2 * alpha
step_size = torch.minimum(step_size, self.max_step_size)
step_size[grad_norm == 0, :] = self.max_step_size
step_size = maybe_expand(step_size, batch_idx, score)
mean = x + step_size * score
x = mean + torch.sqrt(step_size * 2) * noise
x = compute_lattice_polar_decomposition(x)
mean = compute_lattice_polar_decomposition(mean)
return x, mean
|