class

ScoreSDEModel

extendsPretrainedModel

ScoreSDEModel(config: ScoreSDEConfig)

source edit

A time-dependent score network over one of the three SDEs.

Parameters

configScoreSDEConfig

Frozen configuration.

Notes

Reference: Song, Sohl-Dickstein, Kingma, Kumar, Ermon, and Poole, "Score-Based Generative Modeling through Stochastic Differential Equations", ICLR, 2021 (arXiv:2011.13456).

The network predicts the noise that was added; the score follows from the perturbation kernel, since for $x(t) = \mu + \sigma \epsilon$

\nabla_x \log p_{0t}(x(t) \mid x(0)) = -\frac{x(t) - \mu}{\sigma^2} = -\frac{\epsilon}{\sigma}.

Dividing by $\sigma$ rather than folding it into the network is what lets one trained model serve every sampler here — the samplers ask for a score and get one at whatever $t$ they are at.

Examples

>>> import lucid
>>> from lucid.models import score_sde_ve
>>> model = score_sde_ve(sample_size=8, base_channels=8,
...     channel_mult=(1,), num_res_blocks=1, resnet_groups=4,
...     attention_resolutions=()).eval()
>>> model.sde.__class__.__name__
'VESDE'

Used by 2

Constructors

dunder

init

→None

__init__(config: ScoreSDEConfig)

source edit

Instance methods

forward

→ScoreSDEOutput

forward(x: Tensor, t: Tensor | None = None)

source edit

Perturb, then score.

Parameters

xTensor

Clean samples, (N, C, H, W).

tTensor or None= None

Times to evaluate at. None draws them uniformly from [t_min, 1], which is what training does.

Returns

ScoreSDEOutput

Without loss — that belongs to the task wrapper.

predict_noise

→Tensor

predict_noise(x: Tensor, t: Tensor)

source edit

What the network thinks was added — same shape as x.

score

→Tensor

score(x: Tensor, t: Tensor)

source edit

$\nabla_x \log p_t(x)$ , the quantity every sampler wants.

Parameters

xTensor

Perturbed samples, (N, C, H, W).

tTensor

Times, (N,).

Returns

Tensor

The score, same shape as x.

Examples

>>> import lucid
>>> import lucid.nn as nn
>>> from lucid.models import score_sde_vp
>>> model = score_sde_vp(sample_size=8, base_channels=8,
...     channel_mult=(1,), num_res_blocks=1, resnet_groups=4,
...     attention_resolutions=()).eval()
>>> # The output layer starts at zero; give it weights so the check bites.
>>> _ = nn.init.normal_(model.unet.conv_out.weight, std=0.05)
>>> x, t = lucid.randn((2, 3, 8, 8)), lucid.tensor([0.05, 0.9])
>>> score = model.score(x, t)
>>> score.shape
(2, 3, 8, 8)
>>> _, std = model.sde.marginal_prob(x, t)
>>> noise = model.predict_noise(x, t)
>>> bool(lucid.allclose(score * std.reshape(2, 1, 1, 1), -noise, atol=1e-5))
True

>>> import lucid >>> from lucid.models import score_sde_ve >>> model = score_sde_ve(sample_size=8, base_channels=8, ... channel_mult=(1,), num_res_blocks=1, resnet_groups=4, ... attention_resolutions=()).eval() >>> model.sde.__class__.__name__ 'VESDE'

>>> import lucid >>> import lucid.nn as nn >>> from lucid.models import score_sde_vp >>> model = score_sde_vp(sample_size=8, base_channels=8, ... channel_mult=(1,), num_res_blocks=1, resnet_groups=4, ... attention_resolutions=()).eval() >>> # The output layer starts at zero; give it weights so the check bites. >>> _ = nn.init.normal_(model.unet.conv_out.weight, std=0.05) >>> x, t = lucid.randn((2, 3, 8, 8)), lucid.tensor([0.05, 0.9]) >>> score = model.score(x, t) >>> score.shape (2, 3, 8, 8) >>> _, std = model.sde.marginal_prob(x, t) >>> noise = model.predict_noise(x, t) >>> bool(lucid.allclose(score * std.reshape(2, 1, 1, 1), -noise, atol=1e-5)) True