Noise Prediction Loss

The problem, from first principles

DDPM turns generation into a regression problem. Take a clean image, pick a random step t, draw noise eps, form the noisy image x_t, and train a network eps_theta(x_t, t) to predict the noise that was added. The loss is the mean squared error between the predicted and the true noise. Predicting the noise is mathematically equivalent to predicting the clean image: because x_t = sqrt(ab) * x0 + sqrt(1 - ab) * eps, rearranging gives x0 = (x_t - sqrt(1 - ab) * eps) / sqrt(ab). So any noise prediction implies an estimate of the clean image, which is what sampling algorithms use at every step.

From theory to code

Implement noise_prediction_loss and predict_x0.

Constraints

• noise_prediction_loss(eps_pred, eps) is the mean over all elements of (eps_pred - eps) ** 2, as a Python float.
• predict_x0(x_t, eps_pred, t, alpha_bars) returns (x_t - sqrt(1 - ab) * eps_pred) / sqrt(ab) with ab = alpha_bars[t] (0-based index).
• Do not modify the inputs.

Hints