What the samplers and the loss need an SDE to provide.
Notes
Four things, and no more: the drift and diffusion of the forward process, the closed-form perturbation kernel that makes training a single network evaluation, and the prior the reverse process starts from. Everything else — reverse SDE, probability-flow ODE, Predictor-Corrector — is derived from these and lives with the model.
Examples
>>> import lucid
>>> from lucid.models.generative.score_sde._sde import make_sde
>>> sde = make_sde("vp")
>>> x = lucid.zeros(2, 3, 8, 8)
>>> t = lucid.zeros(2) + 0.5
>>> sde.drift(x, t).shape
(2, 3, 8, 8)
The drift follows the state and the diffusion does not: it is a
scalar per sample, which is what makes the noise isotropic.
>>> sde.diffusion(t).shape
(2,)
t_min is not zero. The reverse process is integrated backwards
and the score blows up at the origin, so sampling stops just short
of it.
>>> sde.t_min
0.001Used by 2
Properties
2Variance of the distribution prior_sampling draws from.
Needed by the exact-likelihood computation, which has to evaluate
rather than merely sample it — and VE's prior is
as wide as sigma_max while VP's is unit, so assuming either
would be wrong for the other.
The earliest time the SDE is defined at.
Instance methods
4, per sample.
, the deterministic part of the forward SDE.
The perturbation kernel .
Parameters
Returns
mean, std : TensorMean broadcast to x's shape, and the per-sample standard
deviation. Having this in closed form is what lets the loss
perturb a sample in one step rather than integrating.
Examples
>>> import lucid
>>> from lucid.models.generative.score_sde._sde import make_sde
>>> x = lucid.ones((2, 3, 4, 4))
>>> t = lucid.tensor([0.1, 0.9])
>>> vp = make_sde("vp")
>>> mean, std = vp.marginal_prob(x, t)
>>> mean.shape, std.shape
((2, 3, 4, 4), (2,))
>>> scale = mean[:, 0, 0, 0] # x is all ones, so this is the mean's factor
>>> bool(((scale**2 + std**2 - 1.0).abs() < 1e-5).all().item()) # VP
True
>>> ve = make_sde("ve")
>>> mean, std = ve.marginal_prob(x, t)
>>> mean is x # VE never moves the mean
True
>>> bool(lucid.allclose(std, ve.sigma(t)))
TrueDraw from the distribution the reverse process starts at.