Study Plan

Prediction Parameterization

Diffusion ModelsGenerative Models

Medium

Convert a diffusion model output into all three common prediction targets, the noise, the clean sample and the velocity.

x_t = a\, x_0 + b\, \epsilon, \qquad v = a\, \epsilon - b\, x_0, \qquad a = \sqrt{\bar\alpha_t}, \quad b = \sqrt{1 - \bar\alpha_t}

Here x_t holds the noisy samples of shape (B, ...), timesteps gives each sample its own $t$ , and $\bar\alpha_t$ is alpha_bars at index $t$ . The tensor model_output has the shape of x_t, and prediction_type says what it holds. With "epsilon" it is $\epsilon$ , with "x0" it is $x_0$ , and with "v" it is $v$ . Solve the two linear relations above for the two missing quantities, using each sample's own $a$ and $b$ for all of its elements.

Return a dict with the keys "x0", "epsilon" and "v". Each value is a PyTorch tensor with the shape, dtype and device of x_t, and the entry named by prediction_type equals model_output.

Examples

Input: model_output = [[[[0.5,-1.0],[0.0,1.0]]],[[[1.0,0.2],[-0.4,0.3]]]], x_t = [[[[1.0,-0.5],[0.3,0.8]]],[[[0.6,0.1],[-0.9,0.4]]]], timesteps = [0,2], alpha_bars = [0.9,0.6,0.25], prediction_type = "epsilon"

Output: {"x0":[[[[0.887426,-0.193713],[0.316228,0.509941]]],[[[-0.532051,-0.14641],[-1.10718,0.280385]]]],"epsilon":[[[[0.5,-1.0],[0.0,1.0]]],[[[1.0,0.2],[-0.4,0.3]]]],"v":[[[[0.193713,-0.887426],[-0.1,0.787426]]],[[[0.960769,0.226795],[0.758846,-0.09282]]]]}

Explanation: The model output is the noise, so it is returned unchanged under "epsilon". The first image uses alpha_bars[0] = 0.9 and the second uses alpha_bars[2] = 0.25, and each recovers x0 and v with its own pair of square root factors.

Input: model_output = [[1.0,-1.0,0.5],[0.0,2.0,-0.5]], x_t = [[0.9,-0.6,0.2],[0.4,1.0,0.0]], timesteps = [1,0], alpha_bars = [0.64,0.36], prediction_type = "x0"

Output: {"x0":[[1.0,-1.0,0.5],[0.0,2.0,-0.5]],"epsilon":[[0.375,0.0,-0.125],[0.666667,-1.0,0.666667]],"v":[[-0.575,0.8,-0.475],[0.533333,-2.0,0.833333]]}

Input: model_output = [[0.3,-0.8]], x_t = [[1.2,0.4]], timesteps = [3], alpha_bars = [0.95,0.7,0.4,0.05], prediction_type = "v"

Output: {"x0":[[-0.024076,0.869186]],"epsilon":[[1.236697,0.210986]],"v":[[0.3,-0.8]]}

Hint 1

Each pair of the three quantities determines the third, so recovering x0 first and then deriving the rest keeps the branches short.

Hint 2

alpha_bars[timesteps].view(-1, *[1] * (x_t.dim() - 1)) gives per-sample scalars that broadcast.

Requirements

Use PyTorch tensor operations.
Use 0-based indexing into alpha_bars with each sample's own timestep.
Broadcast the per-sample factors a and b over every non-batch dimension.
Raise a ValueError for any prediction_type other than the exact strings "epsilon", "x0" and "v".
Leave all input tensors unchanged, including model_output.

Constraints

x_t is a floating-point PyTorch tensor of shape (B, ...) with B between 1 and 4 and at least one non-batch dimension.
model_output has the same shape and dtype as x_t.
timesteps is an int64 tensor of shape (B,) with every value between 0 and T - 1.
alpha_bars has shape (T,) with T between 1 and 30, the dtype of x_t, and every value strictly between 0 and 1.
prediction_type is a Python string.
Supported dtypes are float32 and float64. All values are finite.

Try Similar Problems

Posterior Mean and VarianceMedium Noise Prediction LossEasy DDPM Ancestral StepMedium Min-SNR Loss WeightingMedium Reparameterized Gaussian SampleEasy Diagonal Gaussian KL DivergenceEasy Forward Diffusion SampleEasy DDIM StepHard Toy DDPM SamplerHard Autoencoder Forward PassMedium Variational Autoencoder Forward PassMedium GAN Training StepHard

Prediction Parameterization

Diffusion ModelsGenerative Models

Medium

Convert a diffusion model output into all three common prediction targets, the noise, the clean sample and the velocity.

x_t = a\, x_0 + b\, \epsilon, \qquad v = a\, \epsilon - b\, x_0, \qquad a = \sqrt{\bar\alpha_t}, \quad b = \sqrt{1 - \bar\alpha_t}

Return a dict with the keys "x0", "epsilon" and "v". Each value is a PyTorch tensor with the shape, dtype and device of x_t, and the entry named by prediction_type equals model_output.

Examples

Input: model_output = [[1.0,-1.0,0.5],[0.0,2.0,-0.5]], x_t = [[0.9,-0.6,0.2],[0.4,1.0,0.0]], timesteps = [1,0], alpha_bars = [0.64,0.36], prediction_type = "x0"

Output: {"x0":[[1.0,-1.0,0.5],[0.0,2.0,-0.5]],"epsilon":[[0.375,0.0,-0.125],[0.666667,-1.0,0.666667]],"v":[[-0.575,0.8,-0.475],[0.533333,-2.0,0.833333]]}

Input: model_output = [[0.3,-0.8]], x_t = [[1.2,0.4]], timesteps = [3], alpha_bars = [0.95,0.7,0.4,0.05], prediction_type = "v"

Output: {"x0":[[-0.024076,0.869186]],"epsilon":[[1.236697,0.210986]],"v":[[0.3,-0.8]]}

Hint 1

Each pair of the three quantities determines the third, so recovering x0 first and then deriving the rest keeps the branches short.

Hint 2

alpha_bars[timesteps].view(-1, *[1] * (x_t.dim() - 1)) gives per-sample scalars that broadcast.

Requirements

Use PyTorch tensor operations.
Use 0-based indexing into alpha_bars with each sample's own timestep.
Broadcast the per-sample factors a and b over every non-batch dimension.
Raise a ValueError for any prediction_type other than the exact strings "epsilon", "x0" and "v".
Leave all input tensors unchanged, including model_output.

Constraints

x_t is a floating-point PyTorch tensor of shape (B, ...) with B between 1 and 4 and at least one non-batch dimension.
model_output has the same shape and dtype as x_t.
timesteps is an int64 tensor of shape (B,) with every value between 0 and T - 1.
alpha_bars has shape (T,) with T between 1 and 30, the dtype of x_t, and every value strictly between 0 and 1.
prediction_type is a Python string.
Supported dtypes are float32 and float64. All values are finite.

Try Similar Problems

model_output =

Accepts: array

x_t =

Accepts: array

timesteps =

Accepts: array

alpha_bars =

Accepts: array

prediction_type =

Accepts: string