mdlm.noise_mdlm

ContinousTimeNoiseSchedule

Bases: Module, NoiseSchedule

Base class for continuous time noise schedules for absorbing diffusion.

For absorbing diffusion in continuous time, we only need \(\sigma(t)\) and the integral \(\dot\sigma(t)\), which we call noise_rate and total_noise respectively.

Note

We assume that for continous time, \(t \in [0, 1]\).

__init__(antithetic_sampling=True, importance_sampling=False, grad=False, eps=0.001)

Parameters:

Name	Type	Description	Default
antithetic_sampling	bool	If true, the sampled time steps in a batch are sampled around points of a uniform grid over [0, 1], insted of sampling directly from a uniform distribution over [0, 1].	True
importance_sampling	bool	The goal is have a desired distribution over the noise level \(\sigma\), sampling of \(t\) is just a way of obtaining a value of \(\sigma\). Since \(\sigma(t)\) is non-linear function of \(t\), if we want to have a desired distribution over \(\sigma\) for training, which is indeed the case, we cannot simply sample \(t\) uniformly and then transform it to \(\sigma(t)\). Setting importance_sampling=True, will sample uniformly directly over \(\sigma\) in the range \([\sigma_{ ext{min}}, \sigma_{ ext{max}}]\).	False

noise_rate(t)

Return the noise level at time t.

total_noise(t)

Return the total noise at time t.

t_from_noise_rate(noise_rate)

Return the time step t from the noise level sigma.

t_from_total_noise(total_noise)

Return the time step t from the total noise.

forward(t)

Return the noise level at time t.

Parameters:

Name	Type	Description	Default
t	Float[Tensor, ' batch']	The time step tensor of shape (batch)	required

Returns: The noise level tensor of shape (batch) The total noise tensor of shape (batch)

sample_t(batch_size, device=torch.device('cpu'))

Sample a t uniformly from [0, 1].