Study Plan

Noise Prediction Loss

Diffusion ModelsLoss FunctionsGenerative Models

Easy

Compute the simple DDPM training loss of a toy linear denoiser on one batch.

x_t = \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\, \epsilon, \qquad \hat\epsilon = x_t W + \frac{t}{T}\, w_t + b

\ell_i = \frac{1}{D} \sum_{d=1}^{D} \left(\hat\epsilon_{i,d} - \epsilon_{i,d}\right)^2, \qquad L = \frac{1}{B} \sum_{i=1}^{B} \ell_i

Here x0 holds the clean samples of shape (B, ...), timesteps holds one integer timestep $t$ per sample, noise holds the Gaussian noise $\epsilon$ with the shape of x0, and $\bar\alpha_t$ is alpha_bars at index $t$ . The length of alpha_bars is $T$ . Before the denoiser runs, every sample is flattened so that $x_t$ and $\epsilon$ have shape (B, D), where $D$ is the product of all non-batch sizes. The matrix $W$ is weight of shape (D, D), applied on the right as x_t @ weight, not transposed. The vector $w_t$ is time_weight of shape (D,) and $b$ is bias of shape (D,).

Return a dict with two keys. The key "loss" holds $L$ as a Python float, and the key "per_sample_loss" holds the values $\ell_i$ as a PyTorch tensor of shape (B,) with the dtype of x0.

Examples

Input: x0 = [[1.0,-1.0],[0.5,2.0]], timesteps = [0,3], noise = [[0.2,-0.4],[1.0,0.5]], alpha_bars = [0.9,0.7,0.4,0.1], weight = [[0.5,0.1],[-0.2,0.3]], time_weight = [1.0,-1.0], bias = [0.0,0.1]

Output: {"loss":0.214014,"per_sample_loss":[0.17454,0.253488]}

Explanation: The first sample sits at t = 0, so its time feature is 0 and only the linear map of x_t contributes. The second sample sits at t = 3 of T = 4, so the time feature 0.75 scales time_weight before the squared error is averaged over both features.

Input: x0 = [[[[1.0,0.0],[-1.0,0.5]]]], timesteps = [1], noise = [[[[0.3,-0.6],[0.9,0.0]]]], alpha_bars = [0.8,0.5], weight = [[1.0,0.0,0.0,0.0],[0.0,1.0,0.0,0.0],[0.0,0.0,1.0,0.0],[0.0,0.0,0.0,1.0]], time_weight = [0.2,0.2,0.2,0.2], bias = [0.0,0.0,0.0,0.0]

Output: {"loss":0.389296,"per_sample_loss":[0.389296]}

Input: x0 = [[3.0,-2.0,1.0]], timesteps = [2], noise = [[0.4,0.1,-0.7]], alpha_bars = [0.6,0.3,0.0], weight = [[1.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0]], time_weight = [0.0,0.0,0.0], bias = [0.0,0.0,0.0]

Output: {"loss":0.0,"per_sample_loss":[0.0]}

Hint 1

Reshaping x0 and noise to (B, -1) turns any sample shape into flat feature vectors that the matrix product can use.

Hint 2

Turning the integer timesteps into a floating-point column of shape (B, 1) lets the time feature broadcast against time_weight.

Hint 3

Taking the mean along dimension 1 gives one error per sample, and calling item on a 0-d tensor gives a Python float.

Requirements

Use PyTorch tensor operations.
Build x_t with each sample's own timestep, using 0-based indexing into alpha_bars.
Flatten every sample in row-major order (the order of reshape) before applying the denoiser.
Average the squared error over the D features for each sample, then average over the batch.
Leave all input tensors unchanged.

Constraints

x0 is a floating-point PyTorch tensor of shape (B, ...) with B between 1 and 4 and at least one non-batch dimension.
noise has the same shape and dtype as x0.
timesteps is an int64 tensor of shape (B,) with every value between 0 and T - 1.
alpha_bars has shape (T,) with T at least 1 and values in [0, 1].
weight has shape (D, D), and time_weight and bias have shape (D,), where D is the number of elements in one sample.
All floating-point inputs share one dtype, either float32 or float64, and all values are finite.

Try Similar Problems

Posterior Mean and VarianceMedium Prediction ParameterizationMedium Min-SNR Loss WeightingMedium Reparameterized Gaussian SampleEasy Diagonal Gaussian KL DivergenceEasy Forward Diffusion SampleEasy DDPM Ancestral StepMedium Denoising Score Matching LossMedium Dice LossEasy Log Loss (Binary Cross-Entropy)Easy REINFORCE GradientEasy Higher-Order Derivatives & Loss CurvatureMedium

Noise Prediction Loss

Diffusion ModelsLoss FunctionsGenerative Models

Easy

Compute the simple DDPM training loss of a toy linear denoiser on one batch.

x_t = \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\, \epsilon, \qquad \hat\epsilon = x_t W + \frac{t}{T}\, w_t + b

\ell_i = \frac{1}{D} \sum_{d=1}^{D} \left(\hat\epsilon_{i,d} - \epsilon_{i,d}\right)^2, \qquad L = \frac{1}{B} \sum_{i=1}^{B} \ell_i

Return a dict with two keys. The key "loss" holds $L$ as a Python float, and the key "per_sample_loss" holds the values $\ell_i$ as a PyTorch tensor of shape (B,) with the dtype of x0.

Examples

Input: x0 = [[1.0,-1.0],[0.5,2.0]], timesteps = [0,3], noise = [[0.2,-0.4],[1.0,0.5]], alpha_bars = [0.9,0.7,0.4,0.1], weight = [[0.5,0.1],[-0.2,0.3]], time_weight = [1.0,-1.0], bias = [0.0,0.1]

Output: {"loss":0.214014,"per_sample_loss":[0.17454,0.253488]}

Output: {"loss":0.389296,"per_sample_loss":[0.389296]}

Output: {"loss":0.0,"per_sample_loss":[0.0]}

Hint 1

Reshaping x0 and noise to (B, -1) turns any sample shape into flat feature vectors that the matrix product can use.

Hint 2

Turning the integer timesteps into a floating-point column of shape (B, 1) lets the time feature broadcast against time_weight.

Hint 3

Taking the mean along dimension 1 gives one error per sample, and calling item on a 0-d tensor gives a Python float.

Requirements

Use PyTorch tensor operations.
Build x_t with each sample's own timestep, using 0-based indexing into alpha_bars.
Flatten every sample in row-major order (the order of reshape) before applying the denoiser.
Average the squared error over the D features for each sample, then average over the batch.
Leave all input tensors unchanged.

Constraints

x0 is a floating-point PyTorch tensor of shape (B, ...) with B between 1 and 4 and at least one non-batch dimension.
noise has the same shape and dtype as x0.
timesteps is an int64 tensor of shape (B,) with every value between 0 and T - 1.
alpha_bars has shape (T,) with T at least 1 and values in [0, 1].
weight has shape (D, D), and time_weight and bias have shape (D,), where D is the number of elements in one sample.
All floating-point inputs share one dtype, either float32 or float64, and all values are finite.

Try Similar Problems

x0 =

Accepts: array

timesteps =

Accepts: array

noise =

Accepts: array

alpha_bars =

Accepts: array

weight =

Accepts: array

time_weight =

Accepts: array

bias =

Accepts: array