Skip to content

Sub-Variance Preserving

Sub-Variance Preserving diffusion model implementation.

This module extends the basic variance preserving diffusion model with a modified diffusion process that uses a different noise formulation.

SubVariancePreserving

Bases: BaseDiffusion

Sub-Variance Preserving diffusion model implementation.

This class implements a variant of the variance preserving diffusion model with modified noise characteristics. It maintains a controlled level of variance throughout the diffusion process with a different formulation for the diffusion coefficient.

Source code in image_gen\diffusion\sub_vp.py 
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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
class SubVariancePreserving(BaseDiffusion):
    """Sub-Variance Preserving diffusion model implementation.

    This class implements a variant of the variance preserving diffusion model
    with modified noise characteristics. It maintains a controlled level of
    variance throughout the diffusion process with a different formulation
    for the diffusion coefficient.
    """

    def forward_sde(self, x: Tensor, t: Tensor, *args: Any, **kwargs: Any) -> Tuple[
            Tensor, Tensor]:
        """Calculate drift and diffusion coefficients for forward SDE.

        Args:
            x: The input tensor representing current state.
            t: Time steps tensor.
            *args: Additional positional arguments.
            **kwargs: Additional keyword arguments.

        Returns:
            A tuple of (drift, diffusion) tensors.
        """
        beta_t = self.schedule(t, *args, **kwargs).view(-1, 1, 1, 1)
        integral_beta_t = self.schedule.integral_beta(t, *args, **kwargs).view(
            -1, 1, 1, 1)
        exponential_term = torch.exp(-2 * integral_beta_t)
        g_squared = beta_t * (1 - exponential_term)
        diffusion = torch.sqrt(g_squared)
        drift = -0.5 * beta_t * x
        return drift, diffusion

    def forward_process(self, x0: Tensor, t: Tensor, *args: Any, **kwargs: Any) -> Tuple[
            Tensor, Tensor]:
        """Apply the forward diffusion process.

        Adds noise to the input according to the sub-variance preserving 
        schedule.

        Args:
            x0: The input tensor representing initial state.
            t: Time steps tensor.
            *args: Additional positional arguments.
            **kwargs: Additional keyword arguments.

        Returns:
            A tuple of (noisy_sample, noise) tensors.
        """
        integral_beta = self.schedule.integral_beta(t, *args, **kwargs)
        alpha_bar_t = torch.exp(-integral_beta).view(-1, 1, 1, 1)
        mu_x0 = torch.sqrt(alpha_bar_t) * x0
        sigma_t = torch.sqrt(1 - alpha_bar_t)
        noise = torch.randn_like(x0)
        xt = mu_x0 + sigma_t * noise
        return xt, noise

    def compute_loss(self, score: Tensor, noise: Tensor, t: Tensor,
                     *args: Any, **kwargs: Any) -> Tensor:
        """Compute loss between predicted score and actual noise.

        Args:
            score: The predicted noise tensor.
            noise: The actual noise tensor.
            t: Time steps tensor.
            *args: Additional positional arguments.
            **kwargs: Additional keyword arguments.

        Returns:
            A tensor representing the computed loss.
        """
        integral_beta = self.schedule.integral_beta(t, *args, **kwargs)
        alpha_bar_t = torch.exp(-integral_beta)
        sigma_t = torch.sqrt(1 - alpha_bar_t)
        sigma_t = sigma_t.view(score.shape[0], *([1] * (score.dim() - 1)))
        loss = (sigma_t * score + noise) ** 2
        return loss.sum(dim=tuple(range(1, loss.dim())))

    def config(self) -> dict:
        """Get configuration parameters for the diffusion model.

        Returns:
            A dictionary containing configuration parameters.
        """
        return self.schedule.config() if hasattr(self.schedule, "config") else {}

compute_loss(score, noise, t, *args, **kwargs)

Compute loss between predicted score and actual noise.

Parameters:

Name Type Description Default
score Tensor

The predicted noise tensor.

 required
noise Tensor

The actual noise tensor.

 required
t Tensor

Time steps tensor.

 required
*args Any

Additional positional arguments.

 ()
**kwargs Any

Additional keyword arguments.

 {}

Returns:

Type Description
Tensor

A tensor representing the computed loss.
Source code in image_gen\diffusion\sub_vp.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def compute_loss(self, score: Tensor, noise: Tensor, t: Tensor,
                 *args: Any, **kwargs: Any) -> Tensor:
    """Compute loss between predicted score and actual noise.

    Args:
        score: The predicted noise tensor.
        noise: The actual noise tensor.
        t: Time steps tensor.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.

    Returns:
        A tensor representing the computed loss.
    """
    integral_beta = self.schedule.integral_beta(t, *args, **kwargs)
    alpha_bar_t = torch.exp(-integral_beta)
    sigma_t = torch.sqrt(1 - alpha_bar_t)
    sigma_t = sigma_t.view(score.shape[0], *([1] * (score.dim() - 1)))
    loss = (sigma_t * score + noise) ** 2
    return loss.sum(dim=tuple(range(1, loss.dim())))

config()

Get configuration parameters for the diffusion model.

Returns:

Type Description
dict

A dictionary containing configuration parameters.
Source code in image_gen\diffusion\sub_vp.py 
91
92
93
94
95
96
97
def config(self) -> dict:
    """Get configuration parameters for the diffusion model.

    Returns:
        A dictionary containing configuration parameters.
    """
    return self.schedule.config() if hasattr(self.schedule, "config") else {}

forward_process(x0, t, *args, **kwargs)

Apply the forward diffusion process.

Adds noise to the input according to the sub-variance preserving schedule.

Parameters:

Name Type Description Default
x0 Tensor

The input tensor representing initial state.

 required
t Tensor

Time steps tensor.

 required
*args Any

Additional positional arguments.

 ()
**kwargs Any

Additional keyword arguments.

 {}

Returns:

Type Description
Tuple[Tensor, Tensor]

A tuple of (noisy_sample, noise) tensors.
Source code in image_gen\diffusion\sub_vp.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def forward_process(self, x0: Tensor, t: Tensor, *args: Any, **kwargs: Any) -> Tuple[
        Tensor, Tensor]:
    """Apply the forward diffusion process.

    Adds noise to the input according to the sub-variance preserving 
    schedule.

    Args:
        x0: The input tensor representing initial state.
        t: Time steps tensor.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.

    Returns:
        A tuple of (noisy_sample, noise) tensors.
    """
    integral_beta = self.schedule.integral_beta(t, *args, **kwargs)
    alpha_bar_t = torch.exp(-integral_beta).view(-1, 1, 1, 1)
    mu_x0 = torch.sqrt(alpha_bar_t) * x0
    sigma_t = torch.sqrt(1 - alpha_bar_t)
    noise = torch.randn_like(x0)
    xt = mu_x0 + sigma_t * noise
    return xt, noise

forward_sde(x, t, *args, **kwargs)

Calculate drift and diffusion coefficients for forward SDE.

Parameters:

Name Type Description Default
x Tensor

The input tensor representing current state.

 required
t Tensor

Time steps tensor.

 required
*args Any

Additional positional arguments.

 ()
**kwargs Any

Additional keyword arguments.

 {}

Returns:

Type Description
Tuple[Tensor, Tensor]

A tuple of (drift, diffusion) tensors.
Source code in image_gen\diffusion\sub_vp.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def forward_sde(self, x: Tensor, t: Tensor, *args: Any, **kwargs: Any) -> Tuple[
        Tensor, Tensor]:
    """Calculate drift and diffusion coefficients for forward SDE.

    Args:
        x: The input tensor representing current state.
        t: Time steps tensor.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.

    Returns:
        A tuple of (drift, diffusion) tensors.
    """
    beta_t = self.schedule(t, *args, **kwargs).view(-1, 1, 1, 1)
    integral_beta_t = self.schedule.integral_beta(t, *args, **kwargs).view(
        -1, 1, 1, 1)
    exponential_term = torch.exp(-2 * integral_beta_t)
    g_squared = beta_t * (1 - exponential_term)
    diffusion = torch.sqrt(g_squared)
    drift = -0.5 * beta_t * x
    return drift, diffusion