The AI Interview - Master AI/ML Interviews

Implement the forward diffusion process used in Denoising Diffusion Probabilistic Models (DDPMs). The forward process gradually adds Gaussian noise to data over a series of timesteps according to a variance schedule.

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
A noise sample (epsilon)

Write a function forward_diffusion(x_0, t, beta_start, beta_end, num_timesteps, noise) that computes the noisy sample x_t.

The function should:

Create a linear beta schedule from beta_start to beta_end with num_timesteps values
Compute alpha values where alpha = 1 - beta
Compute the cumulative product of alphas (alpha_bar)
Use the alpha_bar value at timestep t to compute the noisy sample

The key property of the forward diffusion is that we can sample x_t directly from x_0 without iterating through all previous timesteps.

Return the noisy sample x_t as a numpy array.

Examples

Example 1:

Input: x_0 = np.array([1.0]), t = 1, beta_start = 0.1, beta_end = 0.2, num_timesteps = 5, noise = np.array([1.0])

Output: [1.2649]

Explanation: 1. Create beta schedule: betas = [0.1, 0.125, 0.15, 0.175, 0.2] 2. Compute alphas = 1 - betas = [0.9, 0.875, 0.85, 0.825, 0.8] 3. Compute alpha_bar (cumulative product) = [0.9, 0.7875, 0.669375, ...] 4. At t=1: alpha_bar_t = 0.9 5. sqrt(0.9) = 0.9487, sqrt(1-0.9) = sqrt(0.1) = 0.3162 6. x_t = 0.9487 * 1.0 + 0.3162 * 1.0 = 1.2649

Starter Code

import numpy as np

def forward_diffusion(x_0: np.ndarray, t: int, beta_start: float, beta_end: float, num_timesteps: int, noise: np.ndarray) -> np.ndarray:
    """
    Apply forward diffusion process to add noise to input data.
    
    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: Noise array (same shape as x_0)
    
    Returns:
        Noisy sample x_t as numpy array
    """
    pass

Forward Diffusion Process

Examples

Starter Code