Diffusion Theory Guide

Fundamental principles of diffusion models for time series generation

Comprehensive guide to diffusion probabilistic models, uncertainty quantification, and their application to time series forecasting and generation.

Table of Contents

1. Introduction to Diffusion Models

Diffusion probabilistic models represent a class of generative models that learn to generate data by gradually denoising random noise. Originally developed for image generation, these models have shown remarkable success in time series forecasting and generation tasks.

Key Advantages of Diffusion Models for Time Series:

  • Uncertainty Quantification: Natural probabilistic framework provides uncertainty estimates
  • High-Quality Generation: Produce realistic and diverse time series samples
  • Flexible Conditioning: Can incorporate various types of conditioning information
  • Stable Training: More stable training compared to GANs
  • Theoretical Foundation: Well-grounded in stochastic differential equations

Core Concept

Diffusion models work by defining a forward process that gradually adds noise to data until it becomes pure noise, and a reverse process that learns to remove this noise to generate new data. For time series, this process preserves temporal dependencies while enabling controlled generation.

2. Mathematical Foundation

Stochastic Differential Equations

Diffusion models are based on stochastic differential equations (SDEs). The forward process is described by:

$$d\mathbf{x}_t = \mathbf{f}(\mathbf{x}_t, t)dt + g(t)d\mathbf{w}$$

where:

Discrete-Time Formulation

For practical implementation, we use a discrete-time formulation with $T$ timesteps:

$$q(\mathbf{x}_t | \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1-\beta_t}\mathbf{x}_{t-1}, \beta_t\mathbf{I})$$

The complete forward process is:

$$q(\mathbf{x}_{1:T} | \mathbf{x}_0) = \prod_{t=1}^T q(\mathbf{x}_t | \mathbf{x}_{t-1})$$

3. Forward Diffusion Process

The forward process gradually adds Gaussian noise to the data according to a predefined noise schedule $\{\beta_t\}_{t=1}^T$.

Noise Scheduling

Common noise schedules include:

Linear Schedule:

β_t = β_1 + (β_T - β_1) * (t-1)/(T-1)

Cosine Schedule:

β_t = 1 - α_t^2 / α_{t-1}^2 α_t = cos(π/2 * (t-1)/T)^2

Reparameterization

Using the reparameterization trick, we can sample $\mathbf{x}_t$ directly from $\mathbf{x}_0$:

$$\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon}$$

where $\bar{\alpha}_t = \prod_{s=1}^t (1-\beta_s)$ and $\boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})$.

4. Reverse Denoising Process

The reverse process learns to denoise the data by predicting the noise added at each timestep:

$$p_\theta(\mathbf{x}_{t-1} | \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))$$

Noise Prediction

The neural network $\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)$ is trained to predict the noise $\boldsymbol{\epsilon}$:

$$\boldsymbol{\mu}_\theta(\mathbf{x}_t, t) = \frac{1}{\sqrt{\alpha_t}}\left(\mathbf{x}_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)\right)$$

Conditional Generation

For conditional generation (e.g., forecasting given historical data), we modify the prediction:

$$\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t, \mathbf{c}) = \text{UNet}(\mathbf{x}_t, t, \mathbf{c})$$

where $\mathbf{c}$ represents the conditioning information (historical time series).

5. Time Series Adaptation

Temporal Encoding

For time series, we need to preserve temporal dependencies. This is achieved through:

Positional Encoding:

PE(pos, 2i) = sin(pos / 10000^(2i/d_model)) PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

Architecture Modifications

Key modifications for time series diffusion models:

Time Series UNet Architecture:

class TimeSeriesUNet(nn.Module): def __init__(self, input_dim, hidden_dim, num_layers): super().__init__() self.embedding = nn.Embedding(input_dim, hidden_dim) self.temporal_conv = nn.Conv1d(hidden_dim, hidden_dim, 3, padding=1) self.attention = nn.MultiheadAttention(hidden_dim, num_heads=8) self.noise_pred = nn.Linear(hidden_dim, input_dim) def forward(self, x, t, condition=None): # Temporal encoding and noise prediction pass

6. Uncertainty Quantification

Diffusion models provide natural uncertainty quantification through their probabilistic nature.

Epistemic Uncertainty

Model uncertainty is captured by sampling multiple times from the reverse process:

$$\text{Epistemic Uncertainty} = \text{Var}[\{\mathbf{x}_0^{(i)}\}_{i=1}^N]$$

Aleatoric Uncertainty

Data uncertainty is modeled through the noise variance in the reverse process:

$$\boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t) = \beta_t \mathbf{I} + \boldsymbol{\Sigma}_{\text{learned}}(\mathbf{x}_t, t)$$

Calibration

To ensure well-calibrated uncertainty estimates, we use:

Calibration Algorithm:

def calibrate_uncertainty(predictions, targets, confidence_level=0.95): # Calculate prediction intervals errors = np.abs(predictions - targets) quantile = np.quantile(errors, confidence_level) # Return calibrated intervals return predictions - quantile, predictions + quantile

7. Training Objectives

Denoising Loss

The primary training objective is the denoising loss:

$$\mathcal{L}_{\text{denoise}} = \mathbb{E}_{t,\boldsymbol{\epsilon}}[\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)\|^2]$$

Conditional Training

For conditional generation, we incorporate conditioning information:

$$\mathcal{L}_{\text{cond}} = \mathbb{E}_{t,\boldsymbol{\epsilon},\mathbf{c}}[\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t, \mathbf{c})\|^2]$$

Regularization

Additional regularization terms can improve performance:

$$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{denoise}} + \lambda_1 \mathcal{L}_{\text{temp}} + \lambda_2 \mathcal{L}_{\text{spec}} + \lambda_3 \mathcal{L}_{\text{div}}$$

8. Sampling Strategies

DDPM Sampling

Standard DDPM sampling uses the following algorithm:

DDPM Sampling Algorithm:

def ddpm_sample(model, T, condition=None): x = torch.randn(shape) # Start with noise for t in reversed(range(T)): # Predict noise epsilon_pred = model(x, t, condition) # Update x x = (x - beta_t/sqrt(1-alpha_bar_t) * epsilon_pred) / sqrt(alpha_t) if t > 0: x += sqrt(beta_t) * torch.randn_like(x) return x

DDIM Sampling

Deterministic sampling for faster generation:

$$\mathbf{x}_{t-1} = \sqrt{\bar{\alpha}_{t-1}}\left(\frac{\mathbf{x}_t - \sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)}{\sqrt{\bar{\alpha}_t}}\right) + \sqrt{1-\bar{\alpha}_{t-1}}\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)$$

Stochastic Sampling

For diverse generation, we can add controlled noise:

$$\mathbf{x}_{t-1} = \boldsymbol{\mu}_\theta(\mathbf{x}_t, t) + \sigma_t \mathbf{z}$$

where $\mathbf{z} \sim \mathcal{N}(0, \mathbf{I})$ and $\sigma_t$ controls the stochasticity.

9. Implementation Considerations

Computational Efficiency

Key considerations for efficient implementation:

Hyperparameter Tuning

Important hyperparameters to consider:

Key Hyperparameters:

  • T (Timesteps): 100-1000, affects quality vs. speed trade-off
  • Learning Rate: 1e-4 to 1e-3, use cosine annealing
  • Batch Size: 32-128, depends on model size and memory
  • Noise Schedule: Cosine or linear, affects generation quality

Evaluation Metrics

Common metrics for evaluating diffusion models on time series:

Evaluation Pipeline:

def evaluate_diffusion_model(model, test_data): # Generate samples samples = model.sample(num_samples=1000) # Calculate metrics accuracy = calculate_forecast_accuracy(samples, test_data) calibration = calculate_ece(samples, test_data) diversity = calculate_diversity(samples) return { 'accuracy': accuracy, 'calibration': calibration, 'diversity': diversity }