Lesson · 40 min · Free
DDPM Internals: Mean, Noise, Score
DDPM Internals: Mean, Noise, Score body { font-family: sans-serif; line-height: 1.6; margin: 20px; } h1, h2 { color: #2c3e50; } pre { background-color: #ecf0f1; padding: 15px; border-radius: 5px; overflow-x: auto; } code
DDPM Internals: Mean, Noise, Score
Welcome to this module on Denoising Diffusion Probabilistic Models (DDPMs), a powerful class of generative models that have shown remarkable success in diverse fields, including image generation and, increasingly, molecular design. In the context of AI in Drug Discovery, understanding DDPMs can unlock new avenues for generating novel molecular structures with desired properties. This lesson will delve into the core mathematical and conceptual underpinnings of DDPMs, focusing on the interplay of mean, noise, and score functions. At a high level, DDPMs operate through a two-phase process: a forward diffusion process and a reverse denoising process. The forward process gradually adds Gaussian noise to an initial data point (e.g., a molecule's representation) until it becomes pure noise. The reverse process then learns to reverse this noise addition, effectively transforming pure noise back into a coherent data point. The magic lies in how this reversal is learned, which heavily depends on estimating the noise at each step.
The Forward Diffusion Process: Adding Noise Systematically
The forward diffusion process is a fixed Markov chain that progressively adds Gaussian noise to an initial data point, x 0 . Each step t in this process transforms x t-1 into x t by adding a small amount of noise. This process is defined by a variance schedule, β 1 , ..., β T , which typically increases over time, meaning more noise is added in later steps. Mathematically, the transition from x t-1 to x t is given by: q(x_t | x_{t-1}) = N(x_t; sqrt(1 - β_t) * x_{t-1}, β_t * I) Here, N denotes a Gaussian distribution. This equation states that x t is sampled from a Gaussian distribution with a mean determined by a scaled version of x t-1 and a variance β t . A crucial property of this forward process is that we can directly sample x t at any arbitrary time step t, given x 0 , without needing to iterate through all intermediate steps. This is achieved by defining α t = 1 - β t and ᾱ t = Π s=1 t α s . Then, we can write: q(x_t | x_0) = N(x_t; sqrt(ᾱ_t) * x_0, (1 - ᾱ_t) * I) This formula is vital because it allows us to directly compute the noisy version of x 0 at any step t, which is essential for training the reverse process. From this, we can see that x t can be expressed as a combination of x 0 and a standard Gaussian noise vector ε: x_t = sqrt(ᾱ_t) * x_0 + sqrt(1 - ᾱ_t) * ε, where ε ~ N(0, I) This equation highlights the two components that constitute x t : a scaled version of the original data (the "mean" component) and a scaled noise component. The goal of the reverse process will be to predict this noise component (ε) at each step to subtract it and recover x t-1 from x t .
The Reverse Denoising Process: Estimating the Noise
The reverse process, parameterised by a neural network (often a U-Net architecture), aims to learn the conditional distribution p θ (x t-1 | x t ). This distribution is also Gaussian, and its mean and variance need to be predicted by the network. Theoretically, if β t is small enough, the reverse conditional distribution q(x t-1 | x t , x 0 ) is also Gaussian. The DDPM paper simplifies this by fixing the variance of the reverse process (or learning a simple variant) and focusing on learning the mean. The key insight is that the mean of the reverse distribution q(x t-1 | x t , x 0 ) can be expressed in terms of x t , x 0 , and the variance schedule. More importantly, it can be re-parameterised to depend on the noise ε that was added to x 0 to get x t . Specifically, the mean μ θ (x t , t) that the neural network learns to predict is typically formulated to predict the noise ε. The network, often denoted as ε θ (x t , t), takes the noisy data x t and the current time step t as input and outputs an estimate of the noise component ε. Once ε θ (x t , t) predicts the noise, we can derive an estimate for x 0 (the "predicted x 0 ") from the forward process equation: x_0_pred = (x_t - sqrt(1 - ᾱ_t) * ε_θ(x_t, t)) / sqrt(ᾱ_t) And then, the mean for sampling x t-1 from x t is: μ_θ(x_t, t) = (1 / sqrt(α_t)) * (x_t - (β_t / sqrt(1 - ᾱ_t)) * ε_θ(x_t, t)) The loss function used to train ε θ is simply the mean squared error between the predicted noise ε θ (x t , t) and the true noise ε that was sampled to generate x t from x 0 : L_t = || ε - ε_θ(x_t, t) ||^2 This elegant formulation allows the network to learn to "denoise" by predicting the noise component, which is directly related to the "score function" in score-based generative models. The score function, defined as ∇ x log p(x), essentially points in the direction of increasing probability density. While DDPMs don't explicitly compute the score function, their noise prediction objective is closely related to estimating it. By predicting the noise, the model implicitly learns the gradient of the log-density, which guides the reverse process towards high-probability regions of the data distribution. In the context of drug discovery, x 0 could represent a molecular graph, a fingerprint vector, or a 3D conformation. The forward process adds noise to this representation, and the reverse process learns to iteratively remove that noise to generate new, valid, and potentially novel molecular structures. The "mean" refers to the network's prediction for the denoised state, "noise" is what the network explicitly learns to predict, and the implicit "score" guides the generation towards chemically plausible and desired structures.
Code Example: Forward Diffusion Step (Conceptual)
This Python-like pseudocode illustrates a single forward diffusion step, showing how noise is added to an input data point. import torch def forward_diffusion_step(x_prev, beta_t, epsilon): """ Performs one step of the forward diffusion process. Args: x_prev (torch.Tensor): Data point at time t-1. beta_t (float): Variance parameter for the current step. epsilon (torch.Tensor): Random noise sampled from N(0, I). Returns: torch.Tensor: Noisy data point at time t. """ alpha_t = 1.0 - beta_t sqrt_alpha_t = torch.sqrt(torch.tensor(alpha_t)) sqrt_one_minus_alpha_t = torch.sqrt(torch.tensor(1.0 - alpha_t)) x_t = sqrt_alpha_t * x_prev + sqrt_one_minus_alpha_t * epsilon return x_t # Example usage (conceptual) # x0 = ... # Your initial molecular representation # beta_schedule = [0.0001, ..., 0.02] # A sequence of beta values # T = len(beta_schedule) # x_t = x0 # for t in range(T): # epsilon_t = torch.randn_like(x_t) # Sample noise # x_t = forward_diffusion_step(x_t, beta_schedule[t], epsilon_t) # # x_t is now progressively noisier
Code Example: Reverse Denoising Step (Conceptual)
This pseudocode demonstrates how a trained noise prediction model ( epsilon_model ) would be used in a reverse denoising step to generate a new sample. import torch # Assume epsilon_model is a trained neural network that predicts noise # It takes x_t and time_step as input, and outputs predicted_noise # epsilon_model = YourNoisePredictionNetwork() def reverse_denoising_step(x_t, time_step, beta_schedule, epsilon_model): """ Performs one step of the reverse denoising process to estimate x_{t-1}. Args: x_t (torch.Tensor): Noisy data point at time t. time_step (int): Current time step (t). beta_schedule (list): List of variance parameters [beta_1, ..., beta_T]. epsilon_model (nn.Module): Trained model to predict noise. Returns: torch.Tensor: Denoised data point at time t-1. """ beta_t = beta_schedule[time_step - 1] # assuming beta_schedule is 0-indexed for t=1..T alpha_t = 1.0 - beta_t alpha_bar_t = 1.0 for i in range(time_step): alpha_bar_t *= (1.0 - beta_schedule[i]) # Calculate alpha_bar_t up to current step # Predict the noise using the trained model predicted_noise = epsilon_model(x_t, torch.tensor(time_step)) # Calculate the mean for the reverse step (mu_theta) # This formula is derived from the re-parameterization trick mean_coeff_1 = 1.0 / torch.sqrt(torch.tensor(alpha_t)) mean_coeff_2 = beta_t / torch.sqrt(torch.tensor(1.0 - alpha_bar_t)) mu_theta = mean_coeff_1 * (x_t - mean_coeff_2 * predicted_noise) # For simplicity, we'll fix the variance to be beta_t for sampling # In practice, other variance schedules like ~beta_t or ~beta_tilde_t are used sigma_t = torch.sqrt(torch.tensor(beta_t)) # Sample x_{t-1} # If t > 1, we add noise; if t == 1, we don't add noise to get the final x0 if time_step > 1: z = torch.randn_like(x_t) x_prev = mu_theta + sigma_t * z else: x_prev = mu_theta # Final step, no additional noise return x_prev # Example usage (conceptual) # initial_noise = torch.randn(
Watch the full lesson — free
This topic is part of AI in Drug Discovery, a complete AI-narrated video course. Press play once and watch the entire lecture like a movie.
Start the course free →