Skip to content

Discrete noise schedules

DiscreteCosineNoiseSchedule

Bases: DiscreteNoiseSchedule

A cosine discrete noise schedule.

Source code in bionemo/moco/schedules/noise/discrete_noise_schedules.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
145
class DiscreteCosineNoiseSchedule(DiscreteNoiseSchedule):
    """A cosine discrete noise schedule."""

    def __init__(self, nsteps: int, nu: Float = 1.0, s: Float = 0.008):
        """Initialize the CosineNoiseSchedule.

        Args:
            nsteps (int): Number of discrete steps.
            nu (Optional[Float]): Hyperparameter for the cosine schedule exponent (default is 1.0).
            s (Optional[Float]): Hyperparameter for the cosine schedule shift (default is 0.008).
        """
        super().__init__(nsteps=nsteps, direction=TimeDirection.DIFFUSION)
        self.nu = nu
        self.s = s

    def _generate_schedule(self, nsteps: Optional[int] = None, device: Union[str, torch.device] = "cpu") -> Tensor:
        """Generate the cosine noise schedule.

        Args:
            nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        if nsteps is None:
            nsteps = self.nsteps
        steps = (
            nsteps + 1
        )  #! matches OpenAI code https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py#L62
        x = torch.linspace(0, nsteps, steps, device=device)
        alphas_cumprod = torch.cos(((x / nsteps) ** self.nu + self.s) / (1 + self.s) * torch.pi * 0.5) ** 2
        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
        betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
        betas = torch.clip(betas, 0.001, 0.999)
        return 1 - betas

    def _clip_noise_schedule(self, alphas2: Tensor, clip_value: Float = 0.001) -> Tensor:
        """For a noise schedule given by alpha^2, this clips alpha_t / alpha_t-1. This may help improve stability during sampling.

        Args:
            alphas2 (Tensor): The noise schedule given by alpha^2.
            clip_value (Optional[Float]): The minimum value for alpha_t / alpha_t-1 (default is 0.001).

        Returns:
            Tensor: The clipped noise schedule.
        """
        alphas2 = torch.cat([torch.ones(1, device=alphas2.device), alphas2], dim=0)

        alphas_step = alphas2[1:] / alphas2[:-1]

        alphas_step = torch.clamp(alphas_step, min=clip_value, max=1.0)
        alphas2 = torch.cumprod(alphas_step, dim=0)

        return alphas2

__init__(nsteps, nu=1.0, s=0.008)

Initialize the CosineNoiseSchedule.

Parameters:

Name Type Description Default
nsteps int

Number of discrete steps.

 required
nu Optional[Float]

Hyperparameter for the cosine schedule exponent (default is 1.0).

 1.0
s Optional[Float]

Hyperparameter for the cosine schedule shift (default is 0.008).

 0.008
Source code in bionemo/moco/schedules/noise/discrete_noise_schedules.py 
 97
 98
 99
100
101
102
103
104
105
106
107
def __init__(self, nsteps: int, nu: Float = 1.0, s: Float = 0.008):
    """Initialize the CosineNoiseSchedule.

    Args:
        nsteps (int): Number of discrete steps.
        nu (Optional[Float]): Hyperparameter for the cosine schedule exponent (default is 1.0).
        s (Optional[Float]): Hyperparameter for the cosine schedule shift (default is 0.008).
    """
    super().__init__(nsteps=nsteps, direction=TimeDirection.DIFFUSION)
    self.nu = nu
    self.s = s

DiscreteLinearNoiseSchedule

Bases: DiscreteNoiseSchedule

A linear discrete noise schedule.

Source code in bionemo/moco/schedules/noise/discrete_noise_schedules.py 
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
class DiscreteLinearNoiseSchedule(DiscreteNoiseSchedule):
    """A linear discrete noise schedule."""

    def __init__(self, nsteps: int, beta_start: Float = 1e-4, beta_end: Float = 0.02):
        """Initialize the CosineNoiseSchedule.

        Args:
            nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
            beta_start (Optional[int]): starting beta value. Defaults to 1e-4.
            beta_end (Optional[int]): end beta value. Defaults to 0.02.
        """
        super().__init__(nsteps=nsteps, direction=TimeDirection.DIFFUSION)
        self.beta_start = beta_start
        self.beta_end = beta_end

    def _generate_schedule(self, nsteps: Optional[int] = None, device: Union[str, torch.device] = "cpu") -> Tensor:
        """Generate the cosine noise schedule.

        Args:
            nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        if nsteps is None:
            nsteps = self.nsteps
        betas = torch.linspace(self.beta_start, self.beta_end, nsteps, dtype=torch.float32, device=device)
        return 1 - betas

__init__(nsteps, beta_start=0.0001, beta_end=0.02)

Initialize the CosineNoiseSchedule.

Parameters:

Name Type Description Default
nsteps Optional[int]

Number of time steps. If None, uses the value from initialization.

 required
beta_start Optional[int]

starting beta value. Defaults to 1e-4.

 0.0001
beta_end Optional[int]

end beta value. Defaults to 0.02.

 0.02
Source code in bionemo/moco/schedules/noise/discrete_noise_schedules.py 
151
152
153
154
155
156
157
158
159
160
161
def __init__(self, nsteps: int, beta_start: Float = 1e-4, beta_end: Float = 0.02):
    """Initialize the CosineNoiseSchedule.

    Args:
        nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
        beta_start (Optional[int]): starting beta value. Defaults to 1e-4.
        beta_end (Optional[int]): end beta value. Defaults to 0.02.
    """
    super().__init__(nsteps=nsteps, direction=TimeDirection.DIFFUSION)
    self.beta_start = beta_start
    self.beta_end = beta_end

DiscreteNoiseSchedule

Bases: ABC

A base class for discrete noise schedules.

Source code in bionemo/moco/schedules/noise/discrete_noise_schedules.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 DiscreteNoiseSchedule(ABC):
    """A base class for discrete noise schedules."""

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

        Args:
           nsteps (int): number of discrete steps.
           direction (TimeDirection): required this defines in which direction the scheduler was built
        """
        self.nsteps = nsteps
        self.direction = string_to_enum(direction, TimeDirection)

    def generate_schedule(
        self,
        nsteps: Optional[int] = None,
        device: Union[str, torch.device] = "cpu",
        synchronize: Optional[TimeDirection] = None,
    ) -> Tensor:
        """Generate the noise schedule as a tensor.

        Args:
            nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
            synchronize (Optional[str]): 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 (default is None).
        """
        schedule = self._generate_schedule(nsteps, device)
        if synchronize and self.direction != string_to_enum(synchronize, TimeDirection):
            return torch.flip(schedule, dims=[0])
        else:
            return schedule

    @abstractmethod
    def _generate_schedule(self, nsteps: Optional[int] = None, device: Union[str, torch.device] = "cpu") -> Tensor:
        """Generate the noise schedule tensor.

        Args:
            nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        pass

    def calculate_derivative(
        self,
        nsteps: Optional[int] = None,
        device: Union[str, torch.device] = "cpu",
        synchronize: Optional[TimeDirection] = None,
    ) -> Tensor:
        """Calculate the time derivative of the schedule.

        Args:
            nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
            synchronize (Optional[str]): 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 (default is None).

        Returns:
            Tensor: A tensor representing the time derivative of the schedule.

        Raises:
            NotImplementedError: If the derivative calculation is not implemented for this schedule.
        """
        raise NotImplementedError("Derivative calculation is not implemented for this schedule.")

__init__(nsteps, direction)

Initialize the DiscreteNoiseSchedule.

Parameters:

Name Type Description Default
nsteps int

number of discrete steps.

 required
direction TimeDirection

required this defines in which direction the scheduler was built

 required
Source code in bionemo/moco/schedules/noise/discrete_noise_schedules.py 
31
32
33
34
35
36
37
38
39
def __init__(self, nsteps: int, direction: TimeDirection):
    """Initialize the DiscreteNoiseSchedule.

    Args:
       nsteps (int): number of discrete steps.
       direction (TimeDirection): required this defines in which direction the scheduler was built
    """
    self.nsteps = nsteps
    self.direction = string_to_enum(direction, TimeDirection)

calculate_derivative(nsteps=None, device='cpu', synchronize=None)

Calculate the time derivative of the schedule.

Parameters:

Name Type Description Default
nsteps Optional[int]

Number of time steps. If None, uses the value from initialization.

 None
device Optional[str]

Device to place the schedule on (default is "cpu").

 'cpu'
synchronize Optional[str]

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 (default is None).

 None

Returns:

Name Type Description
Tensor Tensor

A tensor representing the time derivative of the schedule.

Raises:

Type Description
NotImplementedError

If the derivative calculation is not implemented for this schedule.
Source code in bionemo/moco/schedules/noise/discrete_noise_schedules.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def calculate_derivative(
    self,
    nsteps: Optional[int] = None,
    device: Union[str, torch.device] = "cpu",
    synchronize: Optional[TimeDirection] = None,
) -> Tensor:
    """Calculate the time derivative of the schedule.

    Args:
        nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
        device (Optional[str]): Device to place the schedule on (default is "cpu").
        synchronize (Optional[str]): 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 (default is None).

    Returns:
        Tensor: A tensor representing the time derivative of the schedule.

    Raises:
        NotImplementedError: If the derivative calculation is not implemented for this schedule.
    """
    raise NotImplementedError("Derivative calculation is not implemented for this schedule.")

generate_schedule(nsteps=None, device='cpu', synchronize=None)

Generate the noise schedule as a tensor.

Parameters:

Name Type Description Default
nsteps Optional[int]

Number of time steps. If None, uses the value from initialization.

 None
device Optional[str]

Device to place the schedule on (default is "cpu").

 'cpu'
synchronize Optional[str]

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 (default is None).

 None
Source code in bionemo/moco/schedules/noise/discrete_noise_schedules.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def generate_schedule(
    self,
    nsteps: Optional[int] = None,
    device: Union[str, torch.device] = "cpu",
    synchronize: Optional[TimeDirection] = None,
) -> Tensor:
    """Generate the noise schedule as a tensor.

    Args:
        nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
        device (Optional[str]): Device to place the schedule on (default is "cpu").
        synchronize (Optional[str]): 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 (default is None).
    """
    schedule = self._generate_schedule(nsteps, device)
    if synchronize and self.direction != string_to_enum(synchronize, TimeDirection):
        return torch.flip(schedule, dims=[0])
    else:
        return schedule