The posterior a spatial encoder emits, as a distribution.
Parameters
Notes
Held as a small object rather than a pair of tensors because three call sites want three different things from it: training wants a sample and a KL, encoding for the U-Net wants a sample, and deterministic reconstruction wants the mean. Returning a tuple makes the third indistinguishable from the second at the call site, which is how a "deterministic" path quietly keeps sampling.
logvar is clamped. An encoder early in training can emit a
variance whose exponential overflows, and the resulting inf
reaches the loss as a nan several operations later.
Examples
>>> import lucid
>>> from lucid.models.generative.stable_diffusion import DiagonalGaussian
>>> post = DiagonalGaussian(lucid.zeros((1, 4, 8, 8)), lucid.zeros((1, 4, 8, 8)))
>>> post.mode().shape
(1, 4, 8, 8)
>>> float(post.kl().item())
0.0Used by 1
Constructors
1Instance methods
3KL against , averaged over the batch.
Returns
TensorA scalar.
Notes
Summed over channels and both spatial axes, then averaged over the batch — the latent is one event per image, not one per position. Averaging over positions instead would make the penalty depend on resolution, so the same weight would mean different things at 256 and 512 pixels.
Return the distribution's mode, which for a Gaussian is its mean.
Returns
Tensor(B, C, H, W).
Draw a reparameterised sample.
Returns
Tensor(B, C, H, W).