coarse-grained-forecasting-benchmark's picture

Add anonymous reproduction code

41dd7e0 verified 4 days ago

7.14 kB

"""
core.py — EDM core math for conditional diffusion.
Reimplements the design of NVlabs/edm
(https://github.com/NVlabs/edm; specifically ``training/training_loop.py``)
following Karras et al. 2022, "Elucidating the Design Space of
Diffusion-Based Generative Models", NeurIPS 2022.
Method-only. No architecture, no data loading, no training loop.
Six components:
  - σ sampling at training time
  - preconditioning (c_skip, c_out, c_in, c_noise)
  - denoiser wrapper D_θ
  - training loss with λ(σ) weighting
  - σ schedule for sampling (ρ-parameterized)
  - Heun's 2nd-order sampler with Euler fallback at σ=0
"""

import torch
import math


# --- 1. Training-time σ sampling ---

def sample_training_sigma(batch_size, P_mean, P_std, device):
    """ see Table 1 in EDM paper. Formula for noise distribution under Training section.
    Returns: tensor of shape (batch_size,).
    """
    ln_sigma = torch.empty(batch_size, device=device).normal_(mean=P_mean, std=P_std)
    sigma = torch.exp(ln_sigma)
    return sigma


# --- 2. Preconditioning (EDM Table 1) ---

def c_skip(sigma, sigma_data):
    """Skip-connection scaling."""
    c_skip = sigma_data**2 / (sigma**2 +sigma_data**2)
    return c_skip

def c_out(sigma, sigma_data):
    """Output scaling."""
    c_out = sigma * sigma_data / (torch.sqrt(sigma_data**2+sigma**2))
    return c_out

def c_in(sigma, sigma_data):
    """Input scaling."""
    c_in=1 / (torch.sqrt(sigma**2+sigma_data**2))
    return c_in

def c_noise(sigma):
    """Time-like scalar fed into the network."""
    c_noise=0.25 * torch.log(sigma)
    return c_noise




# --- 3. Denoiser wrapper ---
def denoise(F_theta, x_noisy, sigma, x_0, sigma_data):
    """
    Wraps the raw network F_theta into the preconditioned denoiser D_theta. This corresponds to equation 7 in the EDM paper.
    D_θ(x, σ, x_0) = c_skip(σ)·x + c_out(σ) · F_θ(c_in(σ)·x, c_noise(σ), x_0)
    Used identically in both training (target reconstruction) and sampling (Heuns method)
    Returns predicted clean x_T (same shape as x_noisy).
    """
    # Step 1: normalize the noisy input
    x_normalized = c_in(sigma, sigma_data) * x_noisy

    # Step 2: embed the noise level
    sigma_embedded = c_noise(sigma)

    # Step 3: feed normalized inputs + conditioning through the network
    F_out = F_theta(x_normalized, sigma_embedded, x_0)

    # Step 4: combine via skip connection + scaled output
    D_theta = c_skip(sigma, sigma_data) * x_noisy + c_out(sigma, sigma_data) * F_out

    return D_theta


# --- 4. Loss ---

def lambda_weight(sigma, sigma_data):
    """Per-σ loss weighting (EDM Table 1, section training)."""
    lambda_sigma = (sigma**2 + sigma_data**2)/(sigma*sigma_data)**2
    return lambda_sigma

def edm_loss(F_theta, x_T_batch, x_0_batch, sigma_data, P_mean, P_std):
    """
    Compute the EDM training loss for one batch.
    Args:
        F_theta: the raw network
        x_T_batch: clean data, shape [B, D]
        x_0_batch: conditioning, shape [B, D]
        sigma_data: fixed scalar
        P_mean, P_std: log-normal sampling hyperparameters
    Returns:
        scalar loss
    """
    batch_size = x_T_batch.shape[0]
    device = x_T_batch.device

    # Step 1: sample sigma per example
    sigma = sample_training_sigma(batch_size, P_mean, P_std, device)
    # reshape for broadcasting: [B] -> [B, 1] for vector data
    sigma = sigma.view(-1, 1)

    # Step 2: sample noise
    epsilon = torch.randn_like(x_T_batch)

    # Step 3: add noise (EDM's simple forward process)
    x_noisy = x_T_batch + sigma * epsilon

    # Step 4: run through denoiser
    D_theta = denoise(F_theta, x_noisy, sigma, x_0_batch, sigma_data)

    # Step 5: compute weighted squared error
    squared_error = (D_theta - x_T_batch).pow(2).sum(dim=-1)  # [B]
    weighted_loss = lambda_weight(sigma.squeeze(-1), sigma_data) * squared_error  # [B]

    # Step 6: average over batch
    return weighted_loss.mean()


# --- 5. Sampling-time σ schedule ---

def sampling_sigma_schedule(N, sigma_min, sigma_max, rho, device):
    """ρ-parameterized power-law schedule.
    Returns tensor of length N+1 with σ_0 = σ_max, σ_N = 0,
    and σ_i for 0 < i < N as in EDM Eq. 5.
    """
    # i goes from 0 to N-1
    i = torch.arange(N, device=device, dtype=torch.float32)

    # Apply the power-law formula
    inv_rho = 1.0 / rho
    sigma_max_inv = sigma_max ** inv_rho
    sigma_min_inv = sigma_min ** inv_rho
    sigmas = (sigma_max_inv + (i / (N - 1)) * (sigma_min_inv - sigma_max_inv)) ** rho

    # Append the final σ = 0 (we want to reach truly clean data)
    sigmas = torch.cat([sigmas, torch.zeros(1, device=device)])
    return sigmas


# --- 6. Heun sampler (EDM Algorithm 1) ---

def heun_sample(network, x_0, n_samples, N, sigma_min, sigma_max, rho, sigma_data, device):
    """Sample from p(x_T | x_0) via Heun's 2nd-order method.
    Initializes at x ~ N(0, σ_max² · I), iterates the σ schedule with
    Heun steps, falls back to Euler at the final step where σ_next = 0.
    Returns: tensor of shape (n_samples, x_0.shape[-1]).
    """
    # 1. Build the σ schedule (length N+1, ending at 0)
    sigmas = sampling_sigma_schedule(N, sigma_min, sigma_max, rho, device)

    # 2. Broadcast x_0 to (n_samples, D) so every generated sample uses
    #    the same conditioning but different starting noise.
    if x_0.ndim == 1:
        x_0_batch = x_0.unsqueeze(0).expand(n_samples, -1)
    else:
        x_0_batch = x_0
    D_dim = x_0_batch.shape[-1]

    # 3. Initialize at pure noise scaled to σ_max.
    #    y₀ ~ N(0, σ_max²·I) means y₀ = σ_max · ε for ε ~ N(0, I).
    y = sigma_max * torch.randn(n_samples, D_dim, device=device)

    # 4. Iterate Heun steps along the schedule.
    for i in range(N):
        sigma_cur = sigmas[i]       # current noise level
        sigma_next = sigmas[i + 1]  # target noise level for this step

        # Reshape σ to match batch shape for broadcasting
        sigma_cur_b = sigma_cur.expand(n_samples, 1)

        # --- Predictor: Euler step ---
        # Evaluate the denoiser at the current point
        D_cur = denoise(network, y, sigma_cur_b, x_0_batch, sigma_data)
        # Compute the slope dy/dσ at the current point
        d_cur = (y - D_cur) / sigma_cur
        # Euler step toward σ_next (note: σ_next − σ_cur is negative since σ decreases)
        y_pred = y + (sigma_next - sigma_cur) * d_cur

        # --- Corrector: Heun's average slope ---
        if sigma_next > 0:
            # Evaluate the denoiser at the predicted next point
            sigma_next_b = sigma_next.expand(n_samples, 1)
            D_next = denoise(network, y_pred, sigma_next_b, x_0_batch, sigma_data)
            # Slope at the predicted point
            d_next = (y_pred - D_next) / sigma_next
            # Take the average slope and redo the step
            y = y + (sigma_next - sigma_cur) * 0.5 * (d_cur + d_next)
        else:
            # Final step: σ_next = 0 would divide by zero in d_next.
            # Just use the Euler estimate. This is the only Euler-style step.
            y = y_pred
    return y

	"""
	core.py — EDM core math for conditional diffusion.

	Reimplements the design of NVlabs/edm
	(https://github.com/NVlabs/edm; specifically ``training/training_loop.py``)
	following Karras et al. 2022, "Elucidating the Design Space of
	Diffusion-Based Generative Models", NeurIPS 2022.

	Method-only. No architecture, no data loading, no training loop.

	Six components:
	- σ sampling at training time
	- preconditioning (c_skip, c_out, c_in, c_noise)
	- denoiser wrapper D_θ
	- training loss with λ(σ) weighting
	- σ schedule for sampling (ρ-parameterized)
	- Heun's 2nd-order sampler with Euler fallback at σ=0
	"""

	import torch
	import math


	# --- 1. Training-time σ sampling ---

	def sample_training_sigma(batch_size, P_mean, P_std, device):
	""" see Table 1 in EDM paper. Formula for noise distribution under Training section.
	Returns: tensor of shape (batch_size,).
	"""
	ln_sigma = torch.empty(batch_size, device=device).normal_(mean=P_mean, std=P_std)
	sigma = torch.exp(ln_sigma)
	return sigma


	# --- 2. Preconditioning (EDM Table 1) ---

	def c_skip(sigma, sigma_data):
	"""Skip-connection scaling."""
	c_skip = sigma_data2 / (sigma2 +sigma_data**2)
	return c_skip

	def c_out(sigma, sigma_data):
	"""Output scaling."""
	c_out = sigma * sigma_data / (torch.sqrt(sigma_data2+sigma2))
	return c_out

	def c_in(sigma, sigma_data):
	"""Input scaling."""
	c_in=1 / (torch.sqrt(sigma2+sigma_data2))
	return c_in

	def c_noise(sigma):
	"""Time-like scalar fed into the network."""
	c_noise=0.25 * torch.log(sigma)
	return c_noise




	# --- 3. Denoiser wrapper ---
	def denoise(F_theta, x_noisy, sigma, x_0, sigma_data):
	"""
	Wraps the raw network F_theta into the preconditioned denoiser D_theta. This corresponds to equation 7 in the EDM paper.
	D_θ(x, σ, x_0) = c_skip(σ)·x + c_out(σ) · F_θ(c_in(σ)·x, c_noise(σ), x_0)
	Used identically in both training (target reconstruction) and sampling (Heuns method)
	Returns predicted clean x_T (same shape as x_noisy).
	"""
	# Step 1: normalize the noisy input
	x_normalized = c_in(sigma, sigma_data) * x_noisy

	# Step 2: embed the noise level
	sigma_embedded = c_noise(sigma)

	# Step 3: feed normalized inputs + conditioning through the network
	F_out = F_theta(x_normalized, sigma_embedded, x_0)

	# Step 4: combine via skip connection + scaled output
	D_theta = c_skip(sigma, sigma_data) * x_noisy + c_out(sigma, sigma_data) * F_out

	return D_theta


	# --- 4. Loss ---

	def lambda_weight(sigma, sigma_data):
	"""Per-σ loss weighting (EDM Table 1, section training)."""
	lambda_sigma = (sigma2 + sigma_data2)/(sigmasigma_data)*2
	return lambda_sigma

	def edm_loss(F_theta, x_T_batch, x_0_batch, sigma_data, P_mean, P_std):
	"""
	Compute the EDM training loss for one batch.

	Args:
	F_theta: the raw network
	x_T_batch: clean data, shape [B, D]
	x_0_batch: conditioning, shape [B, D]
	sigma_data: fixed scalar
	P_mean, P_std: log-normal sampling hyperparameters

	Returns:
	scalar loss
	"""
	batch_size = x_T_batch.shape[0]
	device = x_T_batch.device

	# Step 1: sample sigma per example
	sigma = sample_training_sigma(batch_size, P_mean, P_std, device)
	# reshape for broadcasting: [B] -> [B, 1] for vector data
	sigma = sigma.view(-1, 1)

	# Step 2: sample noise
	epsilon = torch.randn_like(x_T_batch)

	# Step 3: add noise (EDM's simple forward process)
	x_noisy = x_T_batch + sigma * epsilon

	# Step 4: run through denoiser
	D_theta = denoise(F_theta, x_noisy, sigma, x_0_batch, sigma_data)

	# Step 5: compute weighted squared error
	squared_error = (D_theta - x_T_batch).pow(2).sum(dim=-1) # [B]
	weighted_loss = lambda_weight(sigma.squeeze(-1), sigma_data) * squared_error # [B]

	# Step 6: average over batch
	return weighted_loss.mean()


	# --- 5. Sampling-time σ schedule ---

	def sampling_sigma_schedule(N, sigma_min, sigma_max, rho, device):
	"""ρ-parameterized power-law schedule.

	Returns tensor of length N+1 with σ_0 = σ_max, σ_N = 0,
	and σ_i for 0 < i < N as in EDM Eq. 5.
	"""
	# i goes from 0 to N-1
	i = torch.arange(N, device=device, dtype=torch.float32)

	# Apply the power-law formula
	inv_rho = 1.0 / rho
	sigma_max_inv = sigma_max ** inv_rho
	sigma_min_inv = sigma_min ** inv_rho
	sigmas = (sigma_max_inv + (i / (N - 1)) * (sigma_min_inv - sigma_max_inv)) ** rho

	# Append the final σ = 0 (we want to reach truly clean data)
	sigmas = torch.cat([sigmas, torch.zeros(1, device=device)])
	return sigmas


	# --- 6. Heun sampler (EDM Algorithm 1) ---

	def heun_sample(network, x_0, n_samples, N, sigma_min, sigma_max, rho, sigma_data, device):
	"""Sample from p(x_T \| x_0) via Heun's 2nd-order method.

	Initializes at x ~ N(0, σ_max² · I), iterates the σ schedule with
	Heun steps, falls back to Euler at the final step where σ_next = 0.

	Returns: tensor of shape (n_samples, x_0.shape[-1]).
	"""
	# 1. Build the σ schedule (length N+1, ending at 0)
	sigmas = sampling_sigma_schedule(N, sigma_min, sigma_max, rho, device)

	# 2. Broadcast x_0 to (n_samples, D) so every generated sample uses
	# the same conditioning but different starting noise.
	if x_0.ndim == 1:
	x_0_batch = x_0.unsqueeze(0).expand(n_samples, -1)
	else:
	x_0_batch = x_0
	D_dim = x_0_batch.shape[-1]

	# 3. Initialize at pure noise scaled to σ_max.
	# y₀ ~ N(0, σ_max²·I) means y₀ = σ_max · ε for ε ~ N(0, I).
	y = sigma_max * torch.randn(n_samples, D_dim, device=device)

	# 4. Iterate Heun steps along the schedule.
	for i in range(N):
	sigma_cur = sigmas[i] # current noise level
	sigma_next = sigmas[i + 1] # target noise level for this step

	# Reshape σ to match batch shape for broadcasting
	sigma_cur_b = sigma_cur.expand(n_samples, 1)

	# --- Predictor: Euler step ---
	# Evaluate the denoiser at the current point
	D_cur = denoise(network, y, sigma_cur_b, x_0_batch, sigma_data)
	# Compute the slope dy/dσ at the current point
	d_cur = (y - D_cur) / sigma_cur
	# Euler step toward σ_next (note: σ_next − σ_cur is negative since σ decreases)
	y_pred = y + (sigma_next - sigma_cur) * d_cur

	# --- Corrector: Heun's average slope ---
	if sigma_next > 0:
	# Evaluate the denoiser at the predicted next point
	sigma_next_b = sigma_next.expand(n_samples, 1)
	D_next = denoise(network, y_pred, sigma_next_b, x_0_batch, sigma_data)
	# Slope at the predicted point
	d_next = (y_pred - D_next) / sigma_next
	# Take the average slope and redo the step
	y = y + (sigma_next - sigma_cur) * 0.5 * (d_cur + d_next)
	else:
	# Final step: σ_next = 0 would divide by zero in d_next.
	# Just use the Euler estimate. This is the only Euler-style step.
	y = y_pred
	return y