Skip to content

Continuous noise transforms

ContinuousExpNoiseTransform

Bases: ABC

A base class for continuous schedules.

alpha = exp(- sigma) where 1 - alpha controls the masking fraction.

Source code in bionemo/moco/schedules/noise/continuous_noise_transforms.py 
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
class ContinuousExpNoiseTransform(ABC):
    """A base class for continuous schedules.

    alpha = exp(- sigma) where 1 - alpha controls the masking fraction.
    """

    def __init__(self, direction: TimeDirection):
        """Initialize the DiscreteNoiseSchedule.

        Args:
            direction : TimeDirection, required this defines in which direction the scheduler was built
        """
        self.direction = string_to_enum(direction, TimeDirection)

    def calculate_sigma(
        self,
        t: Tensor,
        device: Union[str, torch.device] = "cpu",
        synchronize: Optional[TimeDirection] = None,
    ) -> Tensor:
        """Calculate the sigma for the given time steps.

        Args:
            t (Tensor): The input tensor representing the time steps, with values ranging from 0 to 1.
            device (Optional[str]): The device to place the schedule on. Defaults to "cpu".
            synchronize (optional[TimeDirection]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction,
                this parameter allows to flip the direction to match the specified one. Defaults to None.

        Returns:
            Tensor: A tensor representing the sigma values for the given time steps.

        Raises:
            ValueError: If the input time steps exceed the maximum allowed value of 1.
        """
        if t.max() > 1:
            raise ValueError(f"Invalid value: max continuous time is 1, but got {t.max().item()}")

        if synchronize and self.direction != string_to_enum(synchronize, TimeDirection):
            t = 1 - t
        return self._calculate_sigma(t, device)

    @abstractmethod
    def _calculate_sigma(self, t: Tensor, device: Union[str, torch.device] = "cpu") -> Tensor:
        """Calculate the -log of the clean data value for the given time steps.

        Args:
            t (Tensor): The input tensor representing the time steps.
            device (Optional[str]): The device to place the schedule on. Defaults to "cpu".

        Returns:
            Tensor: A tensor representing the sigma values for the given time steps.
        """
        pass

    def sigma_to_alpha(self, sigma: Tensor) -> Tensor:
        """Converts sigma to alpha values by alpha = exp(- sigma).

        Args:
            sigma (Tensor): The input sigma tensor.

        Returns:
            Tensor: A tensor containing the alpha values.
        """
        return torch.exp(-1 * sigma)

__init__(direction)

Initialize the DiscreteNoiseSchedule.

Parameters:

Name Type Description Default
direction

TimeDirection, required this defines in which direction the scheduler was built

 required
Source code in bionemo/moco/schedules/noise/continuous_noise_transforms.py 
34
35
36
37
38
39
40
def __init__(self, direction: TimeDirection):
    """Initialize the DiscreteNoiseSchedule.

    Args:
        direction : TimeDirection, required this defines in which direction the scheduler was built
    """
    self.direction = string_to_enum(direction, TimeDirection)

calculate_sigma(t, device='cpu', synchronize=None)

Calculate the sigma for the given time steps.

Parameters:

Name Type Description Default
t Tensor

The input tensor representing the time steps, with values ranging from 0 to 1.

 required
device Optional[str]

The device to place the schedule on. Defaults to "cpu".

 'cpu'
synchronize optional[TimeDirection]

TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction, this parameter allows to flip the direction to match the specified one. Defaults to None.

 None

Returns:

Name Type Description
Tensor Tensor

A tensor representing the sigma values for the given time steps.

Raises:

Type Description
ValueError

If the input time steps exceed the maximum allowed value of 1.
Source code in bionemo/moco/schedules/noise/continuous_noise_transforms.py 
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
def calculate_sigma(
    self,
    t: Tensor,
    device: Union[str, torch.device] = "cpu",
    synchronize: Optional[TimeDirection] = None,
) -> Tensor:
    """Calculate the sigma for the given time steps.

    Args:
        t (Tensor): The input tensor representing the time steps, with values ranging from 0 to 1.
        device (Optional[str]): The device to place the schedule on. Defaults to "cpu".
        synchronize (optional[TimeDirection]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction,
            this parameter allows to flip the direction to match the specified one. Defaults to None.

    Returns:
        Tensor: A tensor representing the sigma values for the given time steps.

    Raises:
        ValueError: If the input time steps exceed the maximum allowed value of 1.
    """
    if t.max() > 1:
        raise ValueError(f"Invalid value: max continuous time is 1, but got {t.max().item()}")

    if synchronize and self.direction != string_to_enum(synchronize, TimeDirection):
        t = 1 - t
    return self._calculate_sigma(t, device)

sigma_to_alpha(sigma)

Converts sigma to alpha values by alpha = exp(- sigma).

Parameters:

Name Type Description Default
sigma Tensor

The input sigma tensor.

 required

Returns:

Name Type Description
Tensor Tensor

A tensor containing the alpha values.
Source code in bionemo/moco/schedules/noise/continuous_noise_transforms.py 
82
83
84
85
86
87
88
89
90
91
def sigma_to_alpha(self, sigma: Tensor) -> Tensor:
    """Converts sigma to alpha values by alpha = exp(- sigma).

    Args:
        sigma (Tensor): The input sigma tensor.

    Returns:
        Tensor: A tensor containing the alpha values.
    """
    return torch.exp(-1 * sigma)

CosineExpNoiseTransform

Bases: ContinuousExpNoiseTransform

A cosine Exponential noise schedule.

Source code in bionemo/moco/schedules/noise/continuous_noise_transforms.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
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
137
138
139
140
141
142
143
144
class CosineExpNoiseTransform(ContinuousExpNoiseTransform):
    """A cosine Exponential noise schedule."""

    def __init__(self, eps: Float = 1.0e-3):
        """Initialize the CosineNoiseSchedule.

        Args:
            eps (Float): small number to prevent numerical issues.
        """
        self.direction = TimeDirection.DIFFUSION
        self.eps = eps

    def _calculate_sigma(self, t: Tensor, device: Union[str, torch.device] = "cpu") -> Tensor:
        """Calculate negative log of data interpolant fraction.

        Args:
            t (Tensor): The input tensor representing the time steps.
            device (Optional[str]): The device to place the schedule on. Defaults to "cpu".

        Returns:
            Tensor: A tensor representing the sigma values for the given time steps.
        """
        cos = torch.cos(t * torch.pi / 2).to(device)
        return -torch.log(self.eps + (1 - self.eps) * cos)

    def d_dt_sigma(self, t: Tensor, device: Union[str, torch.device] = "cpu") -> Tensor:
        """Compute the derivative of sigma with respect to time.

        Args:
            t (Tensor): The input tensor representing the time steps.
            device (Optional[str]): The device to place the schedule on. Defaults to "cpu".

        Returns:
            Tensor: A tensor representing the derivative of sigma with respect to time.

        Notes:
            The derivative of sigma as a function of time is given by:

            d/dt sigma(t) = d/dt (-log(cos(t * pi / 2) + eps))

            Using the chain rule, we get:

            d/dt sigma(t) = (-1 / (cos(t * pi / 2) + eps)) * (-sin(t * pi / 2) * pi / 2)

            This is the derivative that is computed and returned by this method.
        """
        cos = (1 - self.eps) * torch.cos(t * torch.pi / 2)
        sin = (1 - self.eps) * torch.sin(t * torch.pi / 2)
        scale = torch.pi / 2
        derivative = scale * sin / (cos + self.eps)
        return derivative.to(device)

__init__(eps=0.001)

Initialize the CosineNoiseSchedule.

Parameters:

Name Type Description Default
eps Float

small number to prevent numerical issues.

 0.001
Source code in bionemo/moco/schedules/noise/continuous_noise_transforms.py 
 97
 98
 99
100
101
102
103
104
def __init__(self, eps: Float = 1.0e-3):
    """Initialize the CosineNoiseSchedule.

    Args:
        eps (Float): small number to prevent numerical issues.
    """
    self.direction = TimeDirection.DIFFUSION
    self.eps = eps

d_dt_sigma(t, device='cpu')

Compute the derivative of sigma with respect to time.

Parameters:

Name Type Description Default
t Tensor

The input tensor representing the time steps.

 required
device Optional[str]

The device to place the schedule on. Defaults to "cpu".

 'cpu'

Returns:

Name Type Description
Tensor Tensor

A tensor representing the derivative of sigma with respect to time.
Notes

The derivative of sigma as a function of time is given by:

d/dt sigma(t) = d/dt (-log(cos(t * pi / 2) + eps))

Using the chain rule, we get:

d/dt sigma(t) = (-1 / (cos(t * pi / 2) + eps)) * (-sin(t * pi / 2) * pi / 2)

This is the derivative that is computed and returned by this method.

Source code in bionemo/moco/schedules/noise/continuous_noise_transforms.py 
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def d_dt_sigma(self, t: Tensor, device: Union[str, torch.device] = "cpu") -> Tensor:
    """Compute the derivative of sigma with respect to time.

    Args:
        t (Tensor): The input tensor representing the time steps.
        device (Optional[str]): The device to place the schedule on. Defaults to "cpu".

    Returns:
        Tensor: A tensor representing the derivative of sigma with respect to time.

    Notes:
        The derivative of sigma as a function of time is given by:

        d/dt sigma(t) = d/dt (-log(cos(t * pi / 2) + eps))

        Using the chain rule, we get:

        d/dt sigma(t) = (-1 / (cos(t * pi / 2) + eps)) * (-sin(t * pi / 2) * pi / 2)

        This is the derivative that is computed and returned by this method.
    """
    cos = (1 - self.eps) * torch.cos(t * torch.pi / 2)
    sin = (1 - self.eps) * torch.sin(t * torch.pi / 2)
    scale = torch.pi / 2
    derivative = scale * sin / (cos + self.eps)
    return derivative.to(device)

LogLinearExpNoiseTransform

Bases: ContinuousExpNoiseTransform

A log linear exponential schedule.

Source code in bionemo/moco/schedules/noise/continuous_noise_transforms.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
class LogLinearExpNoiseTransform(ContinuousExpNoiseTransform):
    """A log linear exponential schedule."""

    def __init__(self, eps: Float = 1.0e-3):
        """Initialize the CosineNoiseSchedule.

        Args:
            eps (Float): small value to prevent numerical issues.
        """
        self.direction = TimeDirection.DIFFUSION
        self.eps = eps

    def _calculate_sigma(self, t: Tensor, device: Union[str, torch.device] = "cpu") -> Tensor:
        """Calculate negative log of data interpolant fraction.

        Args:
            t (Tensor): The input tensor representing the time steps.
            device (Optional[str]): The device to place the schedule on. Defaults to "cpu".

        Returns:
            Tensor: A tensor representing the sigma values for the given time steps.
        """
        return -torch.log1p(-(1 - self.eps) * t).to(device)

    def d_dt_sigma(self, t: Tensor, device: Union[str, torch.device] = "cpu") -> Tensor:
        """Compute the derivative of sigma with respect to time.

        Args:
            t (Tensor): The input tensor representing the time steps.
            device (Optional[str]): The device to place the schedule on. Defaults to "cpu".

        Returns:
            Tensor: A tensor representing the derivative of sigma with respect to time.
        """
        derivative = (1 - self.eps) / (1 - (1 - self.eps) * t)
        return derivative.to(device)

__init__(eps=0.001)

Initialize the CosineNoiseSchedule.

Parameters:

Name Type Description Default
eps Float

small value to prevent numerical issues.

 0.001
Source code in bionemo/moco/schedules/noise/continuous_noise_transforms.py 
150
151
152
153
154
155
156
157
def __init__(self, eps: Float = 1.0e-3):
    """Initialize the CosineNoiseSchedule.

    Args:
        eps (Float): small value to prevent numerical issues.
    """
    self.direction = TimeDirection.DIFFUSION
    self.eps = eps

d_dt_sigma(t, device='cpu')

Compute the derivative of sigma with respect to time.

Parameters:

Name Type Description Default
t Tensor

The input tensor representing the time steps.

 required
device Optional[str]

The device to place the schedule on. Defaults to "cpu".

 'cpu'

Returns:

Name Type Description
Tensor Tensor

A tensor representing the derivative of sigma with respect to time.
Source code in bionemo/moco/schedules/noise/continuous_noise_transforms.py 
171
172
173
174
175
176
177
178
179
180
181
182
def d_dt_sigma(self, t: Tensor, device: Union[str, torch.device] = "cpu") -> Tensor:
    """Compute the derivative of sigma with respect to time.

    Args:
        t (Tensor): The input tensor representing the time steps.
        device (Optional[str]): The device to place the schedule on. Defaults to "cpu".

    Returns:
        Tensor: A tensor representing the derivative of sigma with respect to time.
    """
    derivative = (1 - self.eps) / (1 - (1 - self.eps) * t)
    return derivative.to(device)