Skip to content

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
class 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:
        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
    """

    def __init__(self, **dynamics_hparams: Any) -> None:
        """
        Initialize a diffusion process with specific dynamics parameters.

        Args:
            **dynamics_hparams: 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:
            AssertionError: If alpha or sigma is not provided in dynamics_hparams
        """
        super().__init__()
        assert "alpha" in dynamics_hparams
        assert "sigma" in dynamics_hparams
        alpha: Callable[[torch.Tensor], torch.Tensor] = dynamics_hparams["alpha"]
        sigma: Callable[[torch.Tensor], torch.Tensor] = dynamics_hparams["sigma"]
        self.alpha: Callable[[torch.Tensor], torch.Tensor] = alpha
        self.sigma: Callable[[torch.Tensor], torch.Tensor] = sigma
        self.alpha_prime: Callable[[torch.Tensor], torch.Tensor] = scalar_derivative(
            alpha
        )
        self.sigma_prime: Callable[[torch.Tensor], torch.Tensor] = scalar_derivative(
            sigma
        )

    def forward(
        self, x: torch.Tensor, t: torch.Tensor, eps: torch.Tensor
    ) -> torch.Tensor:
        """
        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).

        Args:
            x (torch.Tensor): The input data tensor of shape (N, *D), where N is the batch size
                             and D represents the data dimensions.
            t (torch.Tensor): The time parameter tensor of shape (N,) or broadcastable to x's shape,
                             with values typically in the range [0, 1].
            eps (torch.Tensor): The Gaussian noise tensor of shape (N, *D), where N is the batch size
                               and D represents the data dimensions.
        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.
        """
        alpha = pad_shape_back(self.alpha(t), x.shape)
        sigma = pad_shape_back(self.sigma(t), x.shape)
        return alpha * x + sigma * eps

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
def __init__(self, **dynamics_hparams: Any) -> None:
    """
    Initialize a diffusion process with specific dynamics parameters.

    Args:
        **dynamics_hparams: 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:
        AssertionError: If alpha or sigma is not provided in dynamics_hparams
    """
    super().__init__()
    assert "alpha" in dynamics_hparams
    assert "sigma" in dynamics_hparams
    alpha: Callable[[torch.Tensor], torch.Tensor] = dynamics_hparams["alpha"]
    sigma: Callable[[torch.Tensor], torch.Tensor] = dynamics_hparams["sigma"]
    self.alpha: Callable[[torch.Tensor], torch.Tensor] = alpha
    self.sigma: Callable[[torch.Tensor], torch.Tensor] = sigma
    self.alpha_prime: Callable[[torch.Tensor], torch.Tensor] = scalar_derivative(
        alpha
    )
    self.sigma_prime: Callable[[torch.Tensor], torch.Tensor] = scalar_derivative(
        sigma
    )

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
def forward(
    self, x: torch.Tensor, t: torch.Tensor, eps: torch.Tensor
) -> torch.Tensor:
    """
    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).

    Args:
        x (torch.Tensor): The input data tensor of shape (N, *D), where N is the batch size
                         and D represents the data dimensions.
        t (torch.Tensor): The time parameter tensor of shape (N,) or broadcastable to x's shape,
                         with values typically in the range [0, 1].
        eps (torch.Tensor): The Gaussian noise tensor of shape (N, *D), where N is the batch size
                           and D represents the data dimensions.
    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.
    """
    alpha = pad_shape_back(self.alpha(t), x.shape)
    sigma = pad_shape_back(self.sigma(t), x.shape)
    return alpha * x + sigma * eps

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
class FlowMatchingProcess(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.
    """

    def __init__(self) -> None:
        """
        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.
        """
        super().__init__(alpha=lambda t: 1 - t, sigma=lambda t: t)

__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
def __init__(self) -> None:
    """
    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.
    """
    super().__init__(alpha=lambda t: 1 - t, sigma=lambda t: t)

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
class OrnsteinUhlenbeckProcess(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.
    """

    def __init__(self) -> None:
        """
        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.
        """
        super().__init__(alpha=lambda t: torch.sqrt(1 - t**2), sigma=lambda t: t)

__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
def __init__(self) -> None:
    """
    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.
    """
    super().__init__(alpha=lambda t: torch.sqrt(1 - t**2), sigma=lambda t: t)

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
class VarianceExplodingProcess(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.
    """

    def __init__(self, sigma: Callable[[torch.Tensor], torch.Tensor]) -> None:
        """
        Initialize a Variance Exploding diffusion process.

        Args:
            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.
        """
        super().__init__(alpha=lambda t: torch.ones_like(t), sigma=sigma)

__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
def __init__(self, sigma: Callable[[torch.Tensor], torch.Tensor]) -> None:
    """
    Initialize a Variance Exploding diffusion process.

    Args:
        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.
    """
    super().__init__(alpha=lambda t: torch.ones_like(t), sigma=sigma)