Diffusions
This module contains functionality related to diffusions.
DiffusionProcess
Base class for implementing various diffusion processes.
A diffusion process defines how data evolves over time when noise is added according to specific dynamics. This class provides a framework for implementing different types of diffusion processes used in generative modeling.
The diffusion is parameterized by two functions: - alpha(t): Controls how much of the original signal is preserved at time t - sigma(t): Controls how much noise is added at time t
The forward process is defined as: x_t = alpha(t) * x_0 + sigma(t) * eps, where: - x_0 is the original data - x_t is the noised data at time t - eps is random noise sampled from a standard Gaussian distribution - t is the diffusion time parameter, typically in range [0, 1]
Attributes:
Name | Type | Description |
---|---|---|
alpha |
Callable
|
Function that determines signal preservation at time t, differentiable, maps any tensor to tensor of same shape |
sigma |
Callable
|
Function that determines noise level at time t, differentiable, maps any tensor to tensor of same shape |
alpha_prime |
Callable
|
Derivative of alpha, maps any tensor to tensor of same shape |
sigma_prime |
Callable
|
Derivative of sigma, maps any tensor to tensor of same shape |
Source code in src/diffusionlab/diffusions.py
7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
|
alpha = alpha
instance-attribute
alpha_prime = scalar_derivative(alpha)
instance-attribute
sigma = sigma
instance-attribute
sigma_prime = scalar_derivative(sigma)
instance-attribute
__init__(**dynamics_hparams)
Initialize a diffusion process with specific dynamics parameters.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
**dynamics_hparams
|
Any
|
Keyword arguments containing the dynamics parameters. Must include: - alpha: Callable that maps time t to signal coefficient - sigma: Callable that maps time t to noise coefficient |
{}
|
Raises:
Type | Description |
---|---|
AssertionError
|
If alpha or sigma is not provided in dynamics_hparams |
Source code in src/diffusionlab/diffusions.py
34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
|
forward(x, t, eps)
Forward pass of the dynamics model.
This method implements the forward diffusion process, which gradually adds noise to the input data according to the specified dynamics (alpha and sigma functions).
Parameters:
Name | Type | Description | Default |
---|---|---|---|
x
|
Tensor
|
The input data tensor of shape (N, *D), where N is the batch size and D represents the data dimensions. |
required |
t
|
Tensor
|
The time parameter tensor of shape (N,) or broadcastable to x's shape, with values typically in the range [0, 1]. |
required |
eps
|
Tensor
|
The Gaussian noise tensor of shape (N, *D), where N is the batch size and D represents the data dimensions. |
required |
Returns: torch.Tensor: The noised data at time t, computed as alpha(t) * x + sigma(t) * eps, of shape (N, *D) matching the input shape.
Source code in src/diffusionlab/diffusions.py
61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
|
FlowMatchingProcess
Bases: DiffusionProcess
Implements a Flow Matching diffusion process.
Flow Matching is a technique used in generative modeling where the goal is to learn a continuous transformation (flow) between a simple distribution and a complex data distribution.
In this implementation: - alpha(t) = 1 - t - sigma(t) = t
This creates a linear interpolation between the original data (at t=0) and the noise (at t=1), which is useful for training flow-based generative models.
Source code in src/diffusionlab/diffusions.py
139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
|
__init__()
Initialize a Flow Matching diffusion process with predefined dynamics.
The process uses: - alpha(t) = 1 - t - sigma(t) = t
Both functions map tensors of shape (N,) to tensors of the same shape. This creates a linear interpolation between the original data and noise.
Source code in src/diffusionlab/diffusions.py
155 156 157 158 159 160 161 162 163 164 165 166 |
|
OrnsteinUhlenbeckProcess
Bases: DiffusionProcess
Implements an Ornstein-Uhlenbeck diffusion process.
The Ornstein-Uhlenbeck process is a mean-reverting stochastic process that describes the velocity of a particle undergoing Brownian motion while being subject to friction.
In this implementation: - alpha(t) = sqrt(1 - t²) - sigma(t) = t
This process has properties that make it useful for certain generative modeling tasks, particularly when a smooth transition between clean and noisy states is desired.
Source code in src/diffusionlab/diffusions.py
111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
|
__init__()
Initialize an Ornstein-Uhlenbeck diffusion process with predefined dynamics.
The process uses: - alpha(t) = sqrt(1 - t²) - sigma(t) = t
Both functions map tensors of shape (N,) to tensors of the same shape.
Source code in src/diffusionlab/diffusions.py
126 127 128 129 130 131 132 133 134 135 136 |
|
VarianceExplodingProcess
Bases: DiffusionProcess
Implements a Variance Exploding (VE) diffusion process.
In a VE process, the signal component remains constant (alpha(t) = 1) while the noise component increases according to the provided sigma function. This leads to the variance of the process "exploding" as t increases.
The forward process is defined as: x_t = x_0 + sigma(t) * eps
This is used in models like NCSN (Noise Conditional Score Network) and Score SDE.
Source code in src/diffusionlab/diffusions.py
86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
|
__init__(sigma)
Initialize a Variance Exploding diffusion process.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
sigma
|
Callable
|
Function that determines how noise scales with time t. Should map a tensor of time values of shape (N,) to noise coefficients of the same shape. |
required |
Source code in src/diffusionlab/diffusions.py
99 100 101 102 103 104 105 106 107 108 |
|