Diffusion Reconstruction Loss

Implement a function to compute the reconstruction loss used in training Denoising Diffusion Probabilistic Models (DDPMs).

In diffusion model training, we add noise to data using the forward process, then train a model to predict that noise. One way to evaluate the prediction is to reconstruct the original data from the noisy sample using the predicted noise, then compute the mean squared error against the true original.

Given:

• An original data sample x_0
• A timestep t (1-indexed)
• A linear beta schedule defined by beta_start, beta_end, and num_timesteps
• The noise that was added (epsilon)
• The model's predicted noise (epsilon_pred)

Write a function diffusion_loss(x_0, t, beta_start, beta_end, num_timesteps, noise, predicted_noise) that:

1. Creates a linear beta schedule from beta_start to beta_end with num_timesteps values
2. Computes alpha values where alpha = 1 - beta
3. Computes the cumulative product of alphas (alpha_bar)
4. Uses alpha_bar_t to compute the noisy sample x_t from x_0 and the true noise
5. Reconstructs x_0 from x_t using the predicted noise
6. Returns the mean squared error between the original x_0 and reconstructed x_0

import numpy as np

def diffusion_loss(x_0: np.ndarray, t: int, beta_start: float, beta_end: float, num_timesteps: int, noise: np.ndarray, predicted_noise: np.ndarray) -> float:
    """
    Compute the reconstruction loss for diffusion model training.
    
    Args:
        x_0: Original input data (numpy array)
        t: Timestep (1-indexed, from 1 to num_timesteps)
        beta_start: Starting value of linear beta schedule
        beta_end: Ending value of linear beta schedule
        num_timesteps: Total number of diffusion timesteps
        noise: True noise array (same shape as x_0)
        predicted_noise: Model's predicted noise (same shape as x_0)
    
    Returns:
        Mean squared error loss (float)
    """
    pass