1 / 1

Posterior Mean and Variance

Diffusion ModelsProbabilityGenerative Models

Medium

Compute the mean and variance of the forward process posterior $q(x_{t-1} \mid x_t, x_0)$ for every sample in a batch.

\mu = \frac{\beta_t \sqrt{\bar\alpha_{t-1}}}{1 - \bar\alpha_t}\, x_0 + \frac{(1 - \bar\alpha_{t-1}) \sqrt{\alpha_t}}{1 - \bar\alpha_t}\, x_t, \qquad \tilde\beta_t = \frac{1 - \bar\alpha_{t-1}}{1 - \bar\alpha_t}\, \beta_t

Here betas holds the per-step variances $\beta_t$ of length $T$ , $\alpha_t = 1 - \beta_t$ , and $\bar\alpha_t$ is the product of $\alpha_0$ up to $\alpha_t$ . The value $\bar\alpha_{t-1}$ is the same product one step earlier, and at $t = 0$ it is defined as 1, which makes $\tilde\beta_0 = 0$ . The tensors x0 and x_t have shape (B, ...), and timesteps gives each sample its own $t$ , so the two mean coefficients of a sample scale all of its elements.

The log variance must stay finite at $t = 0$ . For $t \ge 1$ it is $\log \tilde\beta_t$ , and for $t = 0$ it is $\log \tilde\beta_1$ , the log of the variance at the next timestep.

Return a dict with three PyTorch tensors in the dtype and device of x0. The key "mean" holds $\mu$ with the shape of x0, "variance" holds $\tilde\beta_t$ of shape (B,), and "log_variance" holds the clipped log variance of shape (B,).

Examples

Input: x0 = [[[[1.0,-1.0],[0.5,0.0]]],[[[0.2,0.4],[-0.6,0.8]]]], x_t = [[[[0.8,-0.5],[0.9,0.3]]],[[[1.0,-0.2],[0.1,0.6]]]], timesteps = [1,2], betas = [0.1,0.2,0.3,0.4]

Output: {"mean":[[[[0.933182,-0.83735],[0.62631,0.095831]]],[[[0.574953,0.110827],[-0.260703,0.693963]]]],"variance":[0.071429,0.169355],"log_variance":[-2.639057,-1.775759]}

Explanation: The first image uses t = 1, where alpha_bars is 0.9 at t = 0 and 0.72 at t = 1, so its variance is 0.2 * 0.1 / 0.28, about 0.0714. The second image uses t = 2 and gets its own coefficients, which are broadcast over all of its pixels.

Input: x0 = [[2.0,0.0,-2.0]], x_t = [[1.5,0.4,-1.0]], timesteps = [0], betas = [0.1,0.2,0.3]

Output: {"mean":[[2.0,0.0,-2.0]],"variance":[0.0],"log_variance":[-2.639057]}

Input: x0 = [[1.0,1.0],[-1.0,0.5]], x_t = [[0.0,2.0],[0.3,-0.7]], timesteps = [0,1], betas = [0.05,0.1]

Output: {"mean":[[1.0,1.0],[-0.574053,0.107104]],"variance":[0.0,0.034483],"log_variance":[-3.367296,-3.367296]}

Hint 1

torch.cumprod(1 - betas, dim=0) gives alpha_bar.

Hint 2

torch.cat([torch.ones_like(ab[:1]), ab[:-1]]) gives the previous values.

Hint 3

Indexing with timesteps then .view(-1, *[1] * (x0.dim() - 1)) broadcasts per sample.

Requirements

Use PyTorch tensor operations.
Use the per-step alpha inside the square root of the x_t coefficient and the cumulative alpha_bar values everywhere else.
Broadcast each sample's coefficients over every non-batch dimension.
Leave all input tensors unchanged.

Constraints

x0 and x_t are floating-point PyTorch tensors of the same shape (B, ...) and dtype, with B between 1 and 4 and at least one non-batch dimension.
timesteps is an int64 tensor of shape (B,) with every value between 0 and T - 1.
betas has shape (T,) with T between 2 and 30, the same dtype as x0, and every value strictly between 0 and 1.
When the dtype is float32, every beta is at least 0.005. With float64 the betas can be as small as 0.0001.
Supported dtypes are float32 and float64. All values are finite.

Try Similar Problems

Prediction ParameterizationMedium Reparameterized Gaussian SampleEasy Diagonal Gaussian KL DivergenceEasy Noise Prediction LossEasy Gaussian Mixture ScoreMedium DDPM Ancestral StepMedium Min-SNR Loss WeightingMedium Forward Diffusion SampleEasy DDIM StepHard Autoencoder Forward PassMedium Variational Autoencoder Forward PassMedium GAN Training StepHard

Posterior Mean and Variance

Diffusion ModelsProbabilityGenerative Models

Medium

Compute the mean and variance of the forward process posterior $q(x_{t-1} \mid x_t, x_0)$ for every sample in a batch.

\mu = \frac{\beta_t \sqrt{\bar\alpha_{t-1}}}{1 - \bar\alpha_t}\, x_0 + \frac{(1 - \bar\alpha_{t-1}) \sqrt{\alpha_t}}{1 - \bar\alpha_t}\, x_t, \qquad \tilde\beta_t = \frac{1 - \bar\alpha_{t-1}}{1 - \bar\alpha_t}\, \beta_t

The log variance must stay finite at $t = 0$ . For $t \ge 1$ it is $\log \tilde\beta_t$ , and for $t = 0$ it is $\log \tilde\beta_1$ , the log of the variance at the next timestep.

Examples

Input: x0 = [[[[1.0,-1.0],[0.5,0.0]]],[[[0.2,0.4],[-0.6,0.8]]]], x_t = [[[[0.8,-0.5],[0.9,0.3]]],[[[1.0,-0.2],[0.1,0.6]]]], timesteps = [1,2], betas = [0.1,0.2,0.3,0.4]

Output: {"mean":[[[[0.933182,-0.83735],[0.62631,0.095831]]],[[[0.574953,0.110827],[-0.260703,0.693963]]]],"variance":[0.071429,0.169355],"log_variance":[-2.639057,-1.775759]}

Input: x0 = [[2.0,0.0,-2.0]], x_t = [[1.5,0.4,-1.0]], timesteps = [0], betas = [0.1,0.2,0.3]

Output: {"mean":[[2.0,0.0,-2.0]],"variance":[0.0],"log_variance":[-2.639057]}

Input: x0 = [[1.0,1.0],[-1.0,0.5]], x_t = [[0.0,2.0],[0.3,-0.7]], timesteps = [0,1], betas = [0.05,0.1]

Output: {"mean":[[1.0,1.0],[-0.574053,0.107104]],"variance":[0.0,0.034483],"log_variance":[-3.367296,-3.367296]}

Hint 1

torch.cumprod(1 - betas, dim=0) gives alpha_bar.

Hint 2

torch.cat([torch.ones_like(ab[:1]), ab[:-1]]) gives the previous values.

Hint 3

Indexing with timesteps then .view(-1, *[1] * (x0.dim() - 1)) broadcasts per sample.

Requirements

Use PyTorch tensor operations.
Use the per-step alpha inside the square root of the x_t coefficient and the cumulative alpha_bar values everywhere else.
Broadcast each sample's coefficients over every non-batch dimension.
Leave all input tensors unchanged.

Constraints

x0 and x_t are floating-point PyTorch tensors of the same shape (B, ...) and dtype, with B between 1 and 4 and at least one non-batch dimension.
timesteps is an int64 tensor of shape (B,) with every value between 0 and T - 1.
betas has shape (T,) with T between 2 and 30, the same dtype as x0, and every value strictly between 0 and 1.
When the dtype is float32, every beta is at least 0.005. With float64 the betas can be as small as 0.0001.
Supported dtypes are float32 and float64. All values are finite.

Try Similar Problems

x0 =

Accepts: array

x_t =

Accepts: array

timesteps =

Accepts: array

betas =

Accepts: array